Section. NbE with shared normal forms [nbe-nf7]

Now we have the needed infrastructure to implement our idea of improving NbE by creating a thunk of a normal form for a function argument a in the case of weakly normalising App f a. This thunk will be used like a cache: later if the function argument a is duplicated in different places in the function body, full normalisation of a will use the same thunk of a.

To do this, we change the type of Val6 to have an additional constructor Cached7 t v, where t is (a thunk of) the normal form of the value v. We shall have the invariant that Cached is never nested, i.e. Cached _ (Cached _ _) is not allowed.

data Val7 = Closure7 Tm Subst7 | Spine7 Int (SList Val7)
          | Cached7 ~TmS Val7
type Subst7 = SList Val7

Up to now, we have always used the type Tm of terms to store normal forms, but this time we are going to put normal forms in lazy substitutions so we need a new version of Tm that supports lazy shifting. To do this, we modify the constructors Abs and App of Tm to store an additional integer lazy tag (an integer i means that shift0 i has been logically applied to this term):

-- Terms with lazy tagging for `shift0 i`
data TmS = AbsS Int TmS | AppS Int TmS TmS | VarS Int deriving Show

Shifting is again O(1) because we just need to accumulate it in the lazy tags:

instance Shift0 TmS where
  shift0 i (AbsS j b) = AbsS (i + j) b
  shift0 i (AppS j f a) = AppS (i + j)  f a
  shift0 i (VarS j) = VarS (i + j)

Similarly, shifting a value is O(1) because shift0 for SList a and Tms is O(1) and there is no two consecutive layers of Cached:

instance Shift0 Val7 where
  shift0 :: Int -> Val7 -> Val7
  shift0 i (Closure7 b sub) = Closure7 b (shift0 i sub)
  shift0 i (Spine7 j as) = Spine7 (i + j) (shift0 i as)
  shift0 i (Cached7 n v) = Cached7 (shift0 i n) (shift0 i v)

Let's define two auxiliary functions for ignoring and preparing cached normal forms of values:

ignoreCache :: Val7 -> Val7
ignoreCache (Cached7 _ v) = v
ignoreCache v = v

makeCache :: Val7 -> Val7
makeCache v@(Cached7 _ _) = v
makeCache v = Cached7 (reify7 v) v

Our new weak normaliser now makes cached normal forms for function arguments (and ignores the cached normal forms for functions):

wnf7 :: Tm -> Subst7 -> Val7
wnf7 (Var n) sub = lookupSL sub n
wnf7 (Abs b) sub = Closure7 b sub
wnf7 (App f a) sub =
  let f' = wnf7 f sub
      ~a' = makeCache (wnf7 a sub)
  in case ignoreCache f' of
       Closure7 b sub' -> wnf7 b (Cons a' sub')
       Spine7 v as -> Spine7 v (Cons a' as)

When we reify a value, we can use a cached normal form when there is one, otherwise everything stays the same as reify6:

reify7 :: Val7 -> TmS
reify7 (Closure7 t sub) = AbsS 0 (reify7 (wnf7 t (Cons (Spine7 0 Nil) (shift0 1 sub))))
reify7 (Spine7 v as) = foldrSL (\a r -> AppS 0 r (reify7 a)) (VarS v) as
reify7 (Cached7 c _) = c

Reflection doesn't change at all:

reflect7 :: Int -> SList Val7
reflect7 n = foldr (\x rs -> Cons (Spine7 x Nil) rs) Nil [0 .. n]

With these ingredients, our new normaliser is again reify7 (wnf7 t (reflect7 n). Wait! Now reify7 produces an element of TmS, a normal form containing lazy shifts, and of course what we really want is a normal form of type Tm, without any lazy shifts. Therefore we need a function forceShifts :: TmS -> Tm triggers the lazy shifts stored in the resulting normal form (of type TmS) and produces a normal form (of type Tm) that doesn't have lazy shifts.

This function doesn't sound hard, does it? It turns out this function is in fact a bit tricky to get efficient... We don't want this forceShifts function to be more expensive than the actual normalisation process itself, so the naive implementation below isn't good enough:

forceShiftsSpec :: TmS -> Tm
forceShiftsSpec (VarS v) = Var v
forceShiftsSpec (AbsS i b) =
  shift 0 i (Abs (forceShiftsSpec b))
forceShiftsSpec (AppS i f a) =
  shift 0 i (App (forceShiftsSpec f) (forceShiftsSpec a))

The difficulty here is that we can't reduce the composition of two shifting operations into a single one. We do have shift 0 i . shift 0 j = shift 0 (i + j), but unfortunately for the case of AbsS i b above, we have a lambda abstraction, and shift 0 i (Abs b) becomes Abs (shift 1 i b). In general, the composition of two shifts isn't equal to a single shift:

shift x i . shift y j  =/=  shift _ _

To understand shifting functions better, we represent any function S : Int -> Int transforming de Bruijn indices as a sequence:

ss = [s0, s1, ..., sN]

which says that the variable of de Bruijn index 0 should be transformed to s0, the variable of de Bruijn index 1 should be transformed to s1, etc. Therefore the function S is just S v = ss ! v. Then the crucial observation is that the composition S . shift 0 i is

(S . shift 0 i) v = S (v + i) = ss ! (v + i) = (drop i ss) ! v

That is to say, the composite function S . shift 0 i can be represented by the sequence drop i ss. Similarly, when we descend from the context of Abs b to the context of b, the sequence that represents the function S changes from ss to 0 : map (1+) ss.

Therefore to maintain a compact representation of composites of shifting functions, we need a functional data structure that maintains a sequence of integers with the following operations:

  1. pushing a 0 to the head of the sequence,
  2. adding 1 to all the integers in sequence,
  3. dropping the first i elements of the sequence for any given integer i,
  4. accessing the i-th element for any given integer i.

We need the data structure to be functional because when for App f a the computation forks into two branches for f and a respectively, and modifying the sequence in one branch shall not affect the computation in other branches.

We may call this interface 'stack with bulk pop'. Our lazy tagging list SList Int supports the first two operations in O(1) and the other two operations in O(i). The version of forceShifts with SList Int is then as follows.

instance Shift0 Int where
  shift0 i j = i + j

forceShifts :: SList Int -> TmS -> Tm
forceShifts ss (VarS v) = Var (lookupSL ss v)
forceShifts ss (AbsS i b) =
  let ss' = Cons 0 (shift0 1 (dropSL i ss))
  in Abs (forceShifts ss' b)
forceShifts ss (AppS i f a) =
  let ss' = dropSL i ss
  in App (forceShifts ss' f)
         (forceShifts ss' a)

This gives us the new normaliser, normalisation by evaluation with shared normal forms:

nf7 :: Tm -> Tm
nf7 t = let n = fv t
            ss = foldr (\v rs -> Cons v rs) Nil [0 .. n]
        in forceShifts ss (reify7 (wnf7 t (reflect7 n)))

Indeed, nf7 is much faster on the nbeAdversarialExploit program that nf5 was slow:

-- >>> fv (nf7 nbeAdversarialExploit)
-- 0

Unfortunately, our SList Int is too slow for dropping the first i elements. A better implementation of stack with bulk pop would be using balanced trees that support log-time splitting, such as finger trees or red-black trees, which will give us an O(log i)-time implementation of dropping i elements. In our concrete application, i is bound by the depth of lambda abstractions, so i will not be too big in practice.

However, if we want a version of nf7 that has better or equal asymptotic complexity as NbE on all families of input programs, a logarithmic factor is not good enough either. We need pushing, bulk popping, and adding all to be O(1) (accessing in O(i) time is fine because for NbE we also did linear lookup for variables) I do know the existence of a fully persistent data structure in the literature that support all these operations in O(1), namely the one in the paper Improved Algorithms for Finding Level Ancestors in Dynamic Trees by Alstrup and Holm. However, this is a rather complex data structure so we won't implement it here.