How do you express a computation that recursively descend a tree and keep a fresh local state upon each split?

Question

Context

The question in the title is actually a generic version of a specific problem I am having. Feel free to answer the general question or this specific question below.

I am implementing some function that traverse a pure untyped lambda calculus AST and replacing the variables with De-Bruijn indices. There are two representation of the AST, external (with variable names) and internal (with indices):

type Id = String

data TermEx                         
    = VarE Id                       
    | LamE Id TermEx                
    | AppE TermEx TermEx            

data TermIn                         
    = VarI Int                  
    | LamI TermIn               
    | AppI TermIn TermIn        

A quick refresher on how De-Bruijn's indexing works is found on page 6 of this pdf:

http://ttic.uchicago.edu/~pl/classes/CMSC336-Winter08/lectures/lec4.pdf

Current Solution

A regular recursive function that does what I want:

encode' :: TermEx -> TermIn
encode' = go [] where
go en te = case te of 
    VarE x   -> case elemIndex x en of 
        Just i -> VarI 
    LamE x t  -> LamI $ go (x:en) t
    AppE t t' -> AppI (go en t) (go en t')    -- * see comment below

Comment: note function application AppI constitutes to a split in the AST, and at each split, a fresh "local" state is spawned.

Question

Ideally, I want to describe this as some monadic computation that keeps track of a new local state each time the AST splits, my first attempt:

type DeBruijn = forall m. (Monad m, Functor m) => StateT [Id] m TermIn

does not work since all branches will share the same state, throwing off the index. So how do you describe this seemingly very common pattern of computation?

Answer 1

You need the Reader monad transformer:

import Control.Monad.Trans.Reader
import Control.Applicative
import Data.Maybe
import Data.List

type DeBruijnT = ReaderT [Id]

encode :: (Applicative m, Monad m) => TermEx -> DeBrujinT m TermIn
encode t = case t of
  VarE x    -> reader (maybe (error "!") VarI . elemIndex x)
  LamE x t  -> withReaderT (x:) $ LamI <$> encode t
  AppE t t' -> AppI <$> encode t <*> encode t'

Here, withReaderT executes its computation (2nd arg) in local environment, which is the result of applying 1st arg to current environment, so, when new bindings are introduced in branches after split, they would not mess up in the same environment.

Also, if you really want to create new environments on split only, and keep old one on new variable introductions, you can use State monad transformer in such way:

type DeBruijnT = StateT [Id]

encode :: (Applicative m, Monad m) => TermEx -> DeBruijnT m TermIn
encode t = case t of
  VarE x    -> gets (maybe (error "!") VarI . elemIndex x)
  LamE x t  -> withStateT (x:) $ LamI <$> encode t
  AppE t t' -> do
    s <- get
    AppI <$> lift (evalStateT (encode t ) s)
         <*> lift (evalStateT (encode t') s)

Here, instead of binding subcomputations on split to current monadic thread, new monadic threads are spawned and evaluated with current state.