The SPJ trick for typesafe ADT annotations

I came across this technique in SPJ's book, Implementing functional languages: a tutorial (see page 217). In retrospect I guess it's pretty obvious, but I figured I'd write it up since it's kind of nifty and I recently had occasion to use it at work.

Here's a simple ADT:

data Expr = 
  Primitive Val |
  Variable Name |
  App Expr Expr |
  Lambda Name Expr

Say we're writing a compiler and we want to annotate each node with its type, or information about free variables, etc. We could add an extra constructor to hold these sorts of annotations:

data Expr = 
  Primitive Val |
  Variable Name |
  App Expr Expr |
  Lambda Name Expr |
  Annotated Annotation Expr

data Annotation = TypeAnnotation Type | FreeVariables [Name] | ...

This is okay, except that functions expecting certain annotation info have no guarantee it is actually present, or might receive annotations of the wrong type. (Also, how do we know what order the annotations are nested in, if multiple are expected to be present? The types don't tell us anything.)

Okay, let's try again:

data Expr a = 
  Primitive Val |
  Variable Name |
  App Expr Expr |
  Lambda Name Expr |
  Annotation a Expr

Slightly better... now we can at least guarantee that the annotation is of a certain type. But we still don't have any assurance that the annotation is attached to the expected nodes.

Going one step further, we have:

data Expr a = 
  Primitive a Val |
  Variable a Name |
  App a Expr Expr |
  Lambda a Name Expr |

This guarantees the annotation is present, and has a definite type, but there's no uniform way to pull out the annotation info from each constructor. We could write a function, annotation that pattern matches on the Expr to retrieve the annotation. But, it turns out we can do one better:

data Expr a = Expr a (ExprImpl a)

data ExprImpl a = 
  Primitive Val |
  Variable Name |
  App (Expr a) (Expr a) |
  Lambda Name (Expr a)

This is the technique described by SPJ. All nodes in the ADT are guaranteed to have the annotation, and access of the annotation is uniform. Expr can even just be a type alias for (a, ExprImpl a).