Type design for the AST of my language remembering token locations

Question

I wrote a parser and evaluator for a simple programming language. Here is a simplified version of the types for the AST:

data Value = IntV Int | FloatV Float | BoolV Bool
data Expr = IfE Value [Expr] | VarDefE String Value
type Program = [Expr]

I want error messages to tell the line and column of the source code in which the error occured. For example, if the value in an If expression is not a boolean, I want the evaluator to show an error saying "expected boolean at line x, column y", with x and y referring to the location of the value.

So, what I need to do is redefine the previous types so that they can store the relevant locations of different things. One option would be to add a location to each constructor for expressions, like so:

type Location = (Int, Int)

data Expr = IfE Value [Expr] Location | VarDef String Value Location

This clearly isn't optimal, because I have to add those Location fields to every possible expression, and if for example a value contained other values, I would need to add locations to that value too:

{-
this would turn into FunctionCall String [Value] [Location], 
with one location for each value in the function call
-}
data Value = ... | FunctionCall String [Value]

I came up with another solution, which allows me to add locations to everything:

data Located a = Located Location a
type LocatedExpr = Located Expr
type LocatedValue = Located Value

data Value = IntV Int | FloatV Float | BoolV Bool | FunctionCall String [LocatedValue]
data Expr = IfE LocatedValue [LocatedExpr] | VarDef String LocatedValue
data Program = [LocatedExpr]

However I don't like this that much. First of all, it clutters the definition of the evaluator and pattern matching has an extra layer every time. Also, I don't think saying that a function call takes located values as arguments is quite right. Function calls should take values as arguments, and locations should be metadata that doesn't interfere with the evaluator.

I need help redefining my types so that the solution is as clean as possible. Maybe there is a language extension or a design pattern I don't know about that could be helpful.

Answer 1

There are many ways to annotate an AST! This is half of what’s known as the AST typing problem, the other half being how you manage an AST that changes over the course of compilation. The problem isn’t exactly “solved”: all of the solutions have tradeoffs, and which one to pick depends on your expected use cases. I’ll go over a few that you might like to investigate at the end.

Whichever method you choose for organising the actual data types, if it makes pattern-matching ugly or unwieldy, the natural solution is PatternSynonyms.

Considering your first example:

{-# Language PatternSynonyms #-}

type Location = (Int, Int)

data Expr
  = LocatedIf     Value [Expr] Location
  | LocatedVarDef String Value Location

-- Unidirectional pattern synonyms which ignore the location:

pattern If :: Value -> [Expr] -> Expr
pattern If val exprs <- LocatedIf val exprs _loc

pattern VarDef :: String -> Value -> Expr
pattern VarDef name expr <- LocatedVarDef name expr _loc

-- Inform GHC that matching ‘If’ and ‘VarDef’ is just as good
-- as matching ‘LocatedIf’ and ‘LocatedVarDef’.

{-# Complete If, VarDef #-}

This may be sufficiently tidy for your purposes already. But here are a few more tips that I find helpful.

Put annotations first: when adding an annotation type to an AST directly, I often prefer to place it as the first parameter of each constructor, so that it can be conveniently partially applied.

data LocatedExpr
  = LocatedIf     Location Value [Expr]
  | LocatedVarDef Location String Value

If the annotation is a location, then this also makes it more convenient to obtain when writing certain kinds of parsers, along the lines of AnnotatedIf <$> (getSourceLocation <* ifKeyword) <*> value <*> many expr in a parser combinator library.

Parameterise your annotations: I often make the annotation type into a type parameter, so that GHC can derive some useful classes for me:

{-# Language
    DeriveFoldable,
    DeriveFunctor,
    DeriveTraversable #-}

data AnnotatedExpr a
  = AnnotatedIf     a Value [Expr]
  | AnnotatedVarDef a String Value
  deriving (Functor, Foldable, Traversable)

type LocatedExpr = AnnotatedExpr Location

-- Get the annotation of an expression.
-- (Total as long as every constructor is annotated.)
exprAnnotation :: AnnotatedExpr a -> a
exprAnnotation = head

-- Update annotations purely.
mapAnnotations
  :: (a -> b)
  -> AnnotatedExpr a -> AnnotatedExpr b
mapAnnotations = fmap

-- traverse, foldMap, &c.

If you want “doesn’t interfere”, use polymorphism: you can enforce that the evaluator can’t inspect the annotation type by being polymorphic over it. Pattern synonyms still let you match on these expressions conveniently:

pattern If :: Value -> [AnnotatedExpr a] -> AnnotatedExpr a
pattern If val exprs <- AnnotatedIf _anno val exprs

-- …

eval :: AnnotatedExpr a -> Value
eval expr = case expr of
  If     val exprs -> -- …
  VarDef name expr -> -- …

Unannotated terms aren’t your enemy: a term without source locations is no good for error reporting, but I think it’s still a good idea to make the pattern synonyms bidirectional for the convenience of constructing unannotated terms with a unit () annotation. (Or something equivalent, if you use e.g. Maybe Location as the annotation type.)

The reason is that this is quite convenient for writing unit tests, where you want to check the output, but want to use Eq instead of pattern matching, and don’t want to have to compare all the source locations in tests that aren’t concerned with them. Using the derived classes, void :: (Functor f) => f a -> f () strips out all the annotations on an AST.

import Control.Monad (void)

type BareExpr = AnnotatedExpr ()

-- One way to define bidirectional synonyms, so e.g.
-- ‘If’ can be used as either a pattern or a constructor.

pattern If :: Value -> [BareExpr] -> BareExpr
pattern If val exprs = AnnotatedIf () val exprs

-- …

stripAnnotations :: AnnotatedExpr a -> BareExpr
stripAnnotations = void

Equivalently, you could use GADTs / ExistentialQuantification to say data AnyExpr where { AnyExpr :: AnnotatedExpr a -> AnyExpr } / data AnyExpr = forall a. AnyExpr (AnnotatedExpr a); that way, the annotations have exactly as much information as (), but you don’t need to fmap over the entire tree with void in order to strip it, just apply the AnyExpr constructor to hide the type.

---

Finally, here are some brief introductions to a few AST typing solutions.

• Annotate each AST node with a tag (e.g. a unique ID), then store all metadata like source locations, types, and whatever else, separately from the AST:

  import Data.IntMap (IntMap)
  
  -- More sophisticated/stronglier-typed tags are possible.
  newtype Tag = Tag Int
  
  newtype TagMap a = TagMap (IntMap a)
  
  data Expr
    = If     !Tag Value [Expr]
    | VarDef !Tag String Expr
  
  type Span = (Location, Location)
  type SourceMap = TagMap Span
  type CommentMap = TagMap (Span, String)
  parse
    :: String             -- Input
    -> Either ParseError
      ( Expr              -- Parsed expression
      , SourceMap         -- Source locations of tags
      , CommentMap        -- Sideband for comments
      -- …
      )
  

  The advantage is that you can very easily mix in arbitrary new types of annotations anywhere, without affecting the AST itself, and avoid rewriting the AST just to change annotations. You can think of the tree and annotation tables as a kind of database, where the tags are the “foreign keys” relating them. A downside is that you must be careful to maintain these tags when you do rewrite the AST.

  I don’t know if this approach has an established name; I think of it as just “tagging” or a “tagged AST”.

• recursion-schemes and/or Data Types à la Carte^PDF: separate out the “recursive” part of an annotated expression tree from the “annotation” part, and use Fix to tie them back together, with Compose (or Cofree) to add annotations in the middle.

  data ExprF e
    = IfF     Value [e]
    | VarDefF String e
    -- …
    deriving (Foldable, Functor, Traversable, …)
  
  -- Unannotated: Expr ~ ExprF (ExprF (ExprF (…)))
  type Expr = Fix ExprF
  
  -- With a location at each recursive step:
  --
  -- LocatedExpr ~ Located (ExprF (Located (ExprF (…))))
  type LocatedExpr = Fix (Compose Located ExprF)
  
  data Located a = Located Location a
    deriving (Foldable, Functor, Traversable, …)
  -- or: type Located = (,) Location
  

  A distinct advantage is that you get a bunch of nice traversal stuff like cata for free-ish, so you can avoid having to write manual traversals over your AST over and over. A downside is that it adds some pattern clutter to clean up, as does the “à la carte” approach, but they do offer a lot of flexibility.

• Trees That Grow^PDF is overkill for just source locations, but in a serious compiler it’s quite helpful. If you expect to have more than one annotation type (such as inferred types or other analysis results) or an AST that changes over time, then you add a type parameter for the “compilation phase” (parsed, renamed, typechecked, desugared, &c.) and select field types or enable & disable constructors based on that index.

  A really unfortunate downside of this is that you often have to rewrite the tree even in places nothing has changed, because everything depends on the “phase”. An alternative that I use is to add one type parameter for each type of phase or annotation that can vary independently, e.g. data Expr annotation termVarName typeVarName, and abstract over that with type and pattern synonyms. This lets you update indices independently and still use classes like Functor and Bitraversable.