Phase-index your expressions

Implementing a programming language is often just massaging trees through a sequence of transformations into close-but-not-identical forms: resolving names, inferring types, or desugaring syntax. It's easy to mess things up with all the boilerplate required for tree traversal and rewriting, so it would be nice to prove to ourselves that we've done each step across all the cases of our syntax tree.

There's a nifty paper called Trees that Grow that addresses this problem. The authors use type families to annotate each branch of syntax, parameterizing them by the phase of processing. We'll look at an example here that uses GADTs, a closely related feature that is better-behaved in my experience, particularly when it comes to typeclass deriving.

Consider an arithmetic expression language with types:

-- Two types: ints and functions
data Ty =
    TyInt
  | TyFun Ty Ty
  deriving stock (Show)

-- Our first attempt at an expression datatype:
-- vars, ints, addition, and lambdas
data Exp0 =
    ExpVar0 String
  | ExpInt0 Int
  | ExpPlus0 Exp0 Exp0
  | ExpLam0 String (Maybe Ty) Exp0
  -- Lambdas may have argument types
  deriving stock (Show)

We might parse surface syntax like \x -> x + 1 or \(x: Int) -> x + 1 in which lambdas elide or provide types for their arguments. At some point in the pipeline we'll want to always have those argument types on hand - there's no sense repeating type inference! We can write our inference pass to stick a Just value in every lambda, but we'll still have to litter the rest of our code with error "impossible" cases to unwrap the Maybe. God forbid an actual bug creep in, making one of the impossible cases possible...

Let's try another representation:

-- We choose as many phases as we want:
-- Each phase determines a concrete variant
-- of our expression type.
data Phase = Parse | Infer
  deriving stock (Show)

-- After parsing, lambdas MAY have argument types.
-- After type inference, they MUST.
data ArgTy (p :: Phase) where
  ArgTyParse :: Maybe Ty -> ArgTy Parse
  ArgTyInfer :: Ty -> ArgTy Infer

-- GADTs require standalone deriving for this
deriving stock instance Show (ArgTy p)

-- Our final expression type is indexed by Phase.
-- The phase controls which constructors of
-- ArgTy are allowed.
data Exp (p :: Phase) =
    ExpVar String
  | ExpInt Int
  | ExpPlus (Exp p) (Exp p)
  | ExpLam String (ArgTy p) (Exp p)
  deriving stock (Show)

-- Compilation phases are reflected in
-- expression types. (Impls omitted below)

parseExp :: String -> Maybe (Exp Parse)
parseExp = const Nothing

inferTypes :: Exp Parse -> Maybe (Exp Infer)
inferTypes = const Nothing

Phase never shows up at runtime, nor do we need it to. Using GADTs here makes downstream deriving work simply (like Show for Exp), while type families require explicit enumeration in contexts. We should recognize that this is a tradeoff: GADTs are a closed union of constructors for a single type; type families are an open union of types. GADTs work better for you, but type families work better for others. (However, in practice you'll rarely need that external extensibility.)

Another great feature of the GADT encoding is that you can work with Exp p without knowing the phase p at all, making it great for dumping debug information polymorphically. For example, we can render Exp p -> String for any phase if we simply handle all of the ArgTy constructors.