Section. The lazy tagging trick [nbe-lazy-tagging]

The function wnf1 above uses the function beta to substitute the actual argument a' for the binder of the lambda abstraction. Recall that the function beta is defined by

beta b a = shift 0 (-1) (subst b 0 (shift 0 1 a))

So beta b a is not a cheap operation. We need to first use shift to bring the argument a to the next de Bruijn level, then traverses the function body b to do the actual substitution, and then use shift again to bring the result back to the previous de Bruijn level.

Can we do substitution more efficiently? Last year I mentioned on Mastodon that I had a data structure for implementing (capture-avoiding) substitution that avoids the cost of traversing the term. I plan to talk about this data structure more at this year's Logic Colloquium, but here because of some special property of how wnf1 works, we actually don't need any fancy data structure. A very common algorithmic trick suffices.

The trick that we are going to use to make substitution efficient is lazy tagging. The trick is best explained by a simple example. Suppose we have a binary tree storing integers

data Tree = Leaf Int | Node Tree Tree

and suppose that we would like to perform the operation of adding some integer to every integer of a tree:

add :: Int -> Tree -> Tree
add n (Leaf i) = Leaf (n + i)
add n (Node l r) = Node (add n l) (add n r)

The operation add n t works in time O(size t). The trick of lazy tagging optimises add to O(1) by changing the tree datatype to

data Tree' = Leaf' Int | Node'  Int Tree' Tree'

where the constructor Node now has an additional Int field (the lazy tag) storing a number that is logically added to all the leaves of the tree. With this additional field, add can be implemented as

add' :: Int -> Tree' -> Tree'
add' n (Leaf' i) = Leaf' (n + i)
add' n (Node' m l r) = Node' (n + m) l r

On the other hand, whenever we want to access the two sub-trees of a node, we should always push down the lazy tag first:

push :: Tree' -> Tree'
push (Node' m l r)
  | m /= 0 = Node' 0 (add' m l) (add' m r)
push t = t

This push operation is also O(1) because add' is O(1), so having an additional push whenever we want to pattern match a Node' doesn't change the time complexity. For example, if we want to sum the integers of a tree, we do

sumTree :: Tree' -> Int
sumTree (Leaf' i) = i
sumTree t = case push t of Node' _ l r -> sumTree l + sumTree r

In Haskell, we can even use view patterns and pattern synonyms to make the call of push completely transparent.

The reader might be wondering how lazy tagging is different from laziness of Haskell? Why don't we simply make our trees lazy by writing

data Tree'' = Leaf'' Int | Node'' ~Tree ~Tree

The key difference is because of this line:

add' n (Node' m l r) = Node' (n + m) l r

which combines two lazy tags n and m into one tag n + m. In comparison, lazy evaluation of Haskell isn't clever enough to do this; after applying add for Tree'' a few times, we would have nested thunks add n_1 (add n_2 (... (add n_m t))), for which we need O(m) rather than O(1) time to push these added integers one-level down.