Detecting a macro annotated type within the body of a macro annotation

Question

I want to use macro annotations (macro-paradise, Scala 2.11) to generate synthetic traits within an annotated trait's companion object. For example, given some STM abstraction:

trait Var[Tx, A] {
  def apply()         (implicit tx: Tx): A
  def update(value: A)(implicit tx: Tx): Unit
}

I want to define a macro annotation txn such that:

@txn trait Cell[A] {
  val value: A
  var next: Option[Cell[A]]
}

Will be re-synthesised into:

object Cell {
  trait Txn[-Tx, A] {
    def value: A
    def next: Var[Option[Cell.Txn[Tx, A]]] // !
  }
}
trait Cell[A] {
  val value: A
  var next: Option[Cell[A]]
}

I got as far as producing the companion object, the inner trait, and the value member. But obviously, in order for the next member to have the augmented type (instead of Option[Cell[A]], I need Option[Cell.Txn[Tx, A]]), I need to pattern match the type tree and rewrite it.

For example, say I find the next value definition in the original Cell trait like this:

case v @ ValDef(vMods, vName, tpt, rhs) =>

How can I analyse tpt recursively to rewrite any type X[...] being annotated with @txn to X.Txn[Tx, ...]? Is this even possible, given like in the above example that X is yet to be processed? Should I modify Cell to mix-in a marker trait to be detected?

---

So the pattern matching function could start like this:

val tpt1 = tpt match {
  case tq"$ident" => $ident  // obviously don't change this?
  case ??? => ???
}

Answer 1

I would like to preface my answer with the disclaimer that this kind of stuff is not easy to do in current Scala. In an ideal macro system, we would like to typecheck tpt in its lexical context, then walk through the structure of the resulting type replacing X[A] with X.Txn[Tx, A] for those X that are subtypes of TxnMarker, and then use the resulting type in the macro expansion.

However, this kind of happy mixing of untyped trees (that come into macro annotations) and typed trees (that typechecker emits) is incompatible with how compiler internals work (some details about that can be found in Scala macros: What is the difference between typed (aka typechecked) an untyped Trees), so we'll have to approximate.

This will be an approximation, because both in 2.10 and 2.11 our macros are unhygienic, meaning that they are vulnerable to name clashes (e.g. in the final expansion tree Var in Var[...] could bind to something unrelated if e.g. that trait we're rewriting contains a type member called Var). Unfortunately at the moment there's just one way to address this problem robustly, and it is very-very hard to carry out without deep understanding of compiler internals, so I'm not going into those details here.

import scala.reflect.macros.whitebox._
import scala.language.experimental.macros
import scala.annotation.StaticAnnotation

trait Var[T]
trait TxnMarker

object txnMacro {
  def impl(c: Context)(annottees: c.Expr[Any]*): c.Expr[Any] = {
    import c.universe._
    // NOTE: this pattern is only going to work with simple traits
    // for a full pattern that captures all traits, refer to Denys's quasiquote guide:
    // http://den.sh/quasiquotes.html#defns-summary
    val q"$mods trait $name[..$targs] extends ..$parents { ..$stats }" = annottees.head.tree
    def rewire(tpt: Tree): Tree = {
      object RewireTransformer extends Transformer {
        override def transform(tree: Tree): Tree = tree match {
          case AppliedTypeTree(x @ RefTree(xqual, xname), a :: Nil) =>
            val dummyType = q"type SomeUniqueName[T] = $x[T]"
            val dummyTrait = q"$mods trait $name[..$targs] extends ..$parents { ..${stats :+ dummyType} }"
            val dummyTrait1 = c.typecheck(dummyTrait)
            val q"$_ trait $_[..$_] extends ..$_ { ..${_ :+ dummyType1} }" = dummyTrait1
            def refersToSelf = dummyTrait1.symbol == dummyType1.symbol.info.typeSymbol
            def refersToSubtypeOfTxnMarker = dummyType1.symbol.info.baseClasses.contains(symbolOf[TxnMarker])
            if (refersToSelf || refersToSubtypeOfTxnMarker) transform(tq"${RefTree(xqual, xname.toTermName)}.Txn[Tx, $a]")
            else super.transform(tree)
          case _ =>
            super.transform(tree)
        }
      }
      RewireTransformer.transform(tpt)
    }
    val stats1 = stats map {
      // this is a simplification, probably you'll also want to do recursive rewiring and whatnot
      // but I'm omitting that here to focus on the question at hand
      case q"$mods val $name: $tpt = $_" => q"$mods def $name: $tpt"
      case q"$mods var $name: $tpt = $_" => q"$mods def $name: Var[${rewire(tpt)}]"
      case stat => stat
    }
    val annottee1 = q"$mods trait $name[..$targs] extends ..${parents :+ tq"TxnMarker"} { ..$stats }"
    val companion = q"""
      object ${name.toTermName} {
        trait Txn[Tx, A] { ..$stats1 }
      }
    """
    c.Expr[Any](Block(List(annottee1, companion), Literal(Constant(()))))
  }
}

class txn extends StaticAnnotation {
  def macroTransform(annottees: Any*): Any = macro txnMacro.impl
}

Actually, while I was writing this macro, I realized that it is ill-equipped to deal with cyclic dependencies of @txn-annotated definitions. The info.baseClasses.contains(symbolOf[TxnMarker]) check essentially forces expansion of the class/trait referred to in info, so the macro is going to loop. This won't lead to a SOE or a freeze though - scalac will just produce a cyclic reference error and bail out.

At the moment I've no idea how to address this problem purely with macros. Maybe you could leave the code generation part in an annotation macro and then move the type transformation part into a fundep materializer. Oh right, it seems that it might work! Instead of generating

object Cell {
  trait Txn[-Tx, A] {
    def value: A
    def next: Var[Option[Cell.Txn[Tx, A]]]
  }
}
trait Cell[A] {
  val value: A
  var next: Option[Cell[A]]
}

You could actually generate this:

object Cell {
  trait Txn[-Tx, A] {
    def value: A
    def next[U](implicit ev: TxnTypeMapper[Option[Cell[A]], U]): U
  }
}
trait Cell[A] {
  val value: A
  var next: Option[Cell[A]]
}

And TxnTypeMapper[T, U] could be summoned by a fundep materializer macro that would do the Type => Type transformation using Type.map, and that one won't lead to cyclic reference errors, because by the time a materializer is invoked (during typer), all macro annotations will have already expanded. Unfortunately, I don't have time to elaborate at the moment, but this looks doable!!