Section. Stateful Shifts Forcing [nbe-nf8]

If we look at forceShifts again, we see that we use the sequence in a very predictable manner:

forceShifts ss (AppS i f a) =
  let ss' = dropSL i ss
  in App (forceShifts ss' f)
         (forceShifts ss' a)

We first drop the first i-elements of the sequence ss, and then we do recursion for f and a. Both the computations forceShifts ss' f and forceShifts ss' a can subsequently modify the sequence, which is the reason why I said we need a persistent implementation of the sequence. However, we can implement forceShifts with a mutable sequence, as long as we make the two recursion calls sequential and ensure that forceShifts restores the sequence after it finishes its computation. Doing this will imply that we give up the possibility of parallelising forceShifts, but this is an acceptable cost.

With mutation, stacks with bulk pops can be easily implemented as a mutable array, for which we remember an integer delta that is logically added to all integers stored in the array and an integer end that is logically the top of the stack:

data IndicesArr s = IndicesArr
  { delta :: Int
  , end :: Int
  , arr :: STUArray s Int Int
  }

type M s a = IndicesArr s -> ST s a

For mutable array we use Haskell's STUArray, whose operations are encapsulated in the ST monad for a pure interface. We define a type M s a = IndicesArr s -> ST s a (a reader monad transformer applied to the ST monad); most subsequent functions will return this type. (These are Haskell specific things. If we implement the algorithm in a language where mutable state is supported as an ambient effect, there is no need to define this M).

To access the v-th element from the top of the stack, we just read the array (which is O(1)). We also need to remember to add delta to the result, since delta is logically added to all integers stored in the stack:

lookupIA :: Int -> M s Tm
lookupIA v IndicesArr{..} =
  do i <- readArray arr (end - v - 1)
     return (Var (delta + i))

Adding an integer to the stack just adds to the lazy delta:

addIA :: Int -> M s a -> M s a
addIA i m ia@IndicesArr{..} = m ia{delta = delta + i}

Similarly, dropping i elements just moves the end cursor:

dropIA :: Int -> M s a -> M s a
dropIA i m ia@IndicesArr{..} = m ia{end = end - i}

The only non-trivial operation is pushing an integer to the top of the stack. There are two things that we need to be careful: 1. if we want to push i, we actually need to write i - delta to the array, since delta is logically added to all the integers in the stack; 2. we need to restore the original content of the array after we finish our computation, where 'our computation' is passed in as an argument m.

pushIA :: Int -> M s a -> M s a
pushIA i m ia@IndicesArr{..} = do
  oldV <- readArray arr end
  writeArray arr end (i - delta)
  a <- m ia{end = end + 1}
  writeArray arr end oldV
  return a

With these ingredients we can write our stateful implementation of forceShifts, which is essentially the same as forceShifts before except that the data structure is changed:

forceShiftsST :: TmS -> M s Tm
forceShiftsST (VarS v) = lookupIA v
forceShiftsST (AbsS i b) = dropIA i (addIA 1 (pushIA 0 (\ia -> Abs <$> forceShiftsST b ia)))
forceShiftsST (AppS i f a) = dropIA i (\ia -> App <$> forceShiftsST f ia <*> forceShiftsST a ia)

The function forceShifts8 invokes forceShiftsST with a sufficiently large initial array (we could also use the classic doubling vector trick if we don't want to pre-determine the array size).

forceShifts8 :: Int -> TmS -> Tm
forceShifts8 n t = runST $
  do let m = 1 + n + depth t
     indices <- newGenArray (0, m) (\i ->
       if i <= n
        then return (n - i)
        else return (-1))
     forceShiftsST t (IndicesArr 0 (n+1) indices) where

  -- Maximum number of nested abstractions
  depth :: TmS -> Int
  depth (VarS v) = 0
  depth (AbsS _ b) = 1 + depth b
  depth (AppS _ f a) = max (depth f) (depth a)

Our final normaliser is then

nf8 :: Tm -> Tm
nf8 t = let n = fv t
        in forceShifts8 n (reify7 (wnf7 t (reflect7 n)))

Our normaliser is asymptotically better than or equal to the standard NbE normaliser nf5 on all families of input programs. In particular, it is faster on the adversarial test that we saw earlier:

-- >>> fv (nf8 nbeAdversarialExploit)
-- 0

On this test nf7 is slightly faster than our nf8 because there is not much index shifting in this test, but we can construct the following example to differentiate nf8 from nf7:

(\x0. \x1. \x2. ... \xN. (\y . c y y ... y) x0) (\z. c)

where c is a free variable. The normal form of \z. c is shared for all the y-s but this normal form is shifted by N when it comes to the context of c y y ... y.

manyShifts :: Int -> Tm
manyShifts n = Abs (go n) `App` (Abs (Var 1)) where
  go 0 = Abs (foldr (\_ r -> r `App` Var 0) (Var (n + 2)) [1 .. n]) `App` Var n
  go i = Abs (go (i-1))

-- >>> manyShifts 3
-- (\x1. \x2. \x3. \x4. (\x5. x0 x5 x5 x5) x1) (\x1. x0)

-- This computes immediately on my machine:
-- >>> fv (nf8 (manyShifts 10000))
-- 0

-- and this is much slower:
-- >>> fv (nf7 (manyShifts 10000))
-- 0