Indexed Traversal for Binders in DeBruijn Indexed Terms

Posted on March 8, 2015

This post is supposed to be a quick and dirty illustration of an idea I recently had to save some code lines when handling DeBruijn indices, so I won’t get into details. Imagine writing an AST with DeBruijn indexed terms for a dependently typed functional language, like so:

{-# LANGUAGE BangPatterns, RankNTypes, LambdaCase, FlexibleContexts #-}
import Control.Lens
import Control.Applicative
import Data.Monoid (All (..), getAll)

data Term
  = TmFree  String          -- ^ free variables
  | TmBound Int             -- ^ bound variables
  | TmLam   Term Term       -- ^ lambda abstraction, explicitly typed
  | TmPi    Term Term       -- ^ pi abstraction, explicitly typed
  | TmApp   Term Term       -- ^ application
  | TmCon   String          -- ^ constructor
  deriving (Eq, Show)

We have free variables identified by strings, applications, constructors, lambda and pi abstractions, and finally bound variables. Lambda and pi abstractions carry the type of their argument first, and then the term they abstract over, e.g. TmLam (TmCon "Int") (TmBound 0) represents λ(x : Int) -> x. Pi and lambda abstractions share the same DeBruijn binding levels.

Now we want to implement binding aware functions on these terms, like open or close (instantiate and abstract in bound). A first attempt might look like this:

-- | @open a b@ replaces the outermost bound variable in @b@ with @a@.
open :: Term -> Term -> Term
open otm = go 0 where
  go !n (TmBound m) | n == m = otm
  go !n (TmLam t x) = TmLam (go n t) (go (n + 1) x)
  go !n (TmPi t x)  = TmPi (go n t) (go (n + 1) x)
  go !n (TmApp f x) = TmApp (go n f) (go n x)
  go !n tm          = tm

-- | @close v a@ replaces the free variable @v@ with its binding level in @a@.
close :: String -> Term -> Term
close v = go 0 where
  go !n (TmFree w) | v == w = TmBound n
  go !n (TmLam t x) = TmLam (go n t) (go (n + 1) x)
  go !n (TmPi t x)  = TmPi (go n t) (go (n + 1) x)
  go !n (TmApp f x) = TmApp (go n f) (go n x)
  go !n tm          = tm

These functions look quite similar to each other; indeed, a lot of functions will roughly have this form. In addition to that, if we add new recursive terms (like let bindings) to Term, all of the functions will increase in size.

A first idea might be to use generic programming like uniplate to abstract over traversing the sub-terms of a term. But this fails, as it can not detect which sub-terms bind a variable (or multiple ones, consider multi argument lambdas). But having a traversal into the subterms is quite a nice idea to start with, if we could carry extra information (i.e. how many variables are bound) per “hole” in the traversal. Enter indexed traversals:

-- | 
-- @binds n@ is an indexed traversal into the subterms of the term: If a
-- subterm binds @m@ new variables, its index is @n + m@.
binds :: Int -> IndexedTraversal' Int Term Term
binds !n fp = \case
  TmApp f x -> TmApp <$> go 0 f <*> go 0 x
  TmLam t x -> TmLam <$> go 0 t <*> go 1 x
  TmPi  t x -> TmPi  <$> go 0 t <*> go 1 x
  tm        -> pure tm
  where go m = indexed fp (n + m)

The first argument of binds is just for convenience: n + binds 0 fp = binds n fp. Now we can rewrite open and close much more compactly:

open' :: Term -> Term -> Term
open' ot = go 0 where
  go !n (TmBound i) | i == n = ot
  go !n tm = iover (binds n) go tm
close' :: String -> Term -> Term
close' v = go 0 where
  go !n (TmFree w) | v == w = TmBound n
  go !n tm = iover (binds n) go tm

Say we want to extend this with a function that checks whether a term is locally closed, i.e. does not use indices larger than the number of binders it is contained in. binds naturally scales to do this as well:

locallyClosed :: Term -> Bool
locallyClosed = getAll . go 0 where
  go !n (TmBound m) = All (m < n)
  go !n tm = ifoldMapOf (binds n) go tm