Design of immutable, typeable, trees

Question

Here is a design problem I have faced repeatedly. Suppose you're building a compiler, how do you store the types in the trees?

Consider a simple Expr and Type hierarchy, and assume that Plus and Equals are polymorphic (plus on booleans in just ||, for instance).

trait Type
case object BoolType extends Type
case object IntType extends Type
case object Untyped extends Type

trait Expr { var tpe : Type = Untyped }

case class Var(id : String) extends Expr
case class Plus(l : Expr, r : Expr) extends Expr
case class Equals(l : Expr, r : Expr) extends Expr
// ...

Assume further that I do not know the type of identifiers when I construct the expression trees, and therefore cannot know the type by construction. Now a typical typechecking function could look like this:

def typeCheck(env : Map[String,Type])(expr : Expr) : Expr = expr match {
  case Var(id) =>
    expr.tpe = env(id)
    expr

  case Plus(l,r) =>
    val tl = typeCheck(env)(l)
    val tr = typeCheck(env)(r)
    assert(tl == tr)
    expr.tpe = tl
    expr

  // etc.
}

This is rather straightforward to write, but comes with two major problems:

• Exprs are mutable. No one likes mutation.
• Typed and untyped expressions cannot be distinguished. I cannot write a function whose signature specifies that its argument must be a typed expression.

So my question is the following. What is a good way (I dare not say design pattern) to define possibly untyped trees such that:

1. I need to define the Expr hierarchy only once.
2. Typed and untyped trees have distinct types and I can choose to make them incompatible.

Edit: One more requirement is that it should work for type systems with an unbounded and unpredictable number of types (think: case class ClassType(classID : String) extends Type, for instance).

Answer 1

This is a perfect use-case for type-level programming!

First, we need a type-level Option so that we can represent untyped trees in terms of type-level None and typed trees of type X in terms of type-level Some[X]:

// We are restricting our type-level option to
// only (potentially) hold subtypes of `Type`.
sealed trait IsTyped
sealed trait Untyped extends IsTyped
sealed trait Typed[T <: Type] extends IsTyped

Next, we lay out our type system hierarchy:

sealed trait Type

// We can create complicated subhierarchies if we want.
sealed trait SimpleType extends Type
sealed trait CompoundType extends Type

sealed trait PrimitiveType extends Type
sealed trait UserType extends Type

// Declaring our types.
case object IntType extends SimpleType with PrimitiveType

case object BoolType extends SimpleType with PrimitiveType

// A type with unbounded attributes.
case class ClassType(classId: String) extends CompoundType with UserType

// A type that depends statically on another type.
case class ArrayType(elemType: Type) extends CompoundType with PrimitiveType

Now, all that's left is to declare our expression tree:

sealed trait Expr[IT <: IsTyped] { val getType: Option[Type] }

// Our actual expression types.
case class Var[IT <: IsTyped](id: String, override val getType: Option[Type] = None) extends Expr[IT]

case class Plus[IT <: IsTyped](l: Expr[IT], r: Expr[IT], override val getType: Option[Type] = None) extends Expr[IT]

case class Equals[IT <: IsTyped](l: Expr[IT], r: Expr[IT], override val getType: Option[Type] = None) extends Expr[IT]

case class ArrayLiteral[IT](elems: List[Expr[_ :< IsTyped]], override val getType: Option[Type] = None) extends Expr[IT]

---

EDIT:

A simple but complete type-checking function:

def typeCheck(expr: Expr[Untyped], env: Map[String, Type]): Option[Expr[Typed[_ :< Type]]] = expr match {
  case Var(id, None) if env isDefinedAt id => Var[Typed[_ <: Type]](id, Some(env(id)))
  case Plus(r, l, None) => for {
      lt <- typeCheck(l, env)
      IntType <- lt.getType
      rt <- typeCheck(r, env)
      IntType <- rt.getType
    } yield Plus[Typed[IntType]](lt, rt, Some(IntType))
  case Equals(r, l, None) => for {
      lt <- typeCheck(l, env)
      lType <- lt.getType
      rt <- typeCheck(r, env)
      rType <- rt.getType
      if rType == lType
    } yield Equals[Typed[BoolType]](lt, rt, Some(BoolType))
  case ArrayLiteral(elems, None) => {
    val elemst: List[Option[Expr[Typed[_ <: Type]]]] =
      elems map { typeCheck(_, env) }
    val elemType: Option[Type] = if (elemst.isEmpty) None else elemst map { elem =>
      elem map { _.getType }
    } reduce { (elemType1, elemType2) =>
      for {
        et1 <- elemType1
        et2 <- elemType2
        if et1 == et2
      } yield et1
    }
    if (elemst forall { _.isDefined }) elemType map { et =>
      ArrayLiteral[Typed[ArrayType]](elemst map { _.get }, ArrayType(et))
    } else None
  }
  case _ => None
}