Create or extend a companion object, using a macro

问题:

Using a Scala 2.10/2.11 macro paradise annotation macro, how can I add or extend the companion object of an annotated class? Skeleton:

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

class foo extends StaticAnnotation {
  def macroTransform(annottees: Any*) = macro fooMacro.impl
}
object fooMacro {
  def impl(c: Context)(annottees: c.Expr[Any]*): c.Expr[Any] = ???
}

Such that, given

trait Foo[A]

the following input

@foo class Bar

object Baz {
  def baz = 33
}
@foo class Baz

will be expanded as:

object Bar {
  implicit def hasFoo: Foo[Bar] = ???
}
class Bar

object Baz {
  def baz = 33

  implicit def hasFoo: Foo[Baz] = ???
}
class Baz

---

Here is a first naive attempt, just adding def hasFoo = 33 for now:

def impl(c: Context)(annottees: c.Expr[Any]*): c.Expr[Any] = {
  import c.universe._
  val inputs : List[Tree] = annottees.map(_.tree)(collection.breakOut)
  val outputs: List[Tree] = inputs match {
    case (cd @ ClassDef(_, cName, _, _)) :: tail =>
      val mod0: ModuleDef = tail match {
        case (md @ ModuleDef(_, mName, _)) :: Nil
             if cName.decoded == mName.decoded => md
        case Nil =>
          val cMod = cd.mods
          var mModF = NoFlags
          if (cMod hasFlag Flag.PRIVATE  ) mModF |= Flag.PRIVATE
          if (cMod hasFlag Flag.PROTECTED) mModF |= Flag.PROTECTED
          if (cMod hasFlag Flag.LOCAL    ) mModF |= Flag.LOCAL
          val mMod = Modifiers(mModF, cMod.privateWithin, Nil)
          // or should we have parents = List(AnyRef) and body = List(DefDef(???))
          val mTemp = Template(parents = Nil, self = noSelfType, body = Nil)
          val mName = TermName(cName.decoded) // or encoded?
          ModuleDef(mMod, mName, mTemp)
        case _ => c.abort(c.enclosingPosition, "Expected a companion object")
      }
      val Template(mTempParents, mTempSelf, mTempBody0) = mod0.impl
      val fooDef = DefDef(NoMods, TermName("hasFoo"), Nil, Nil, 
        TypeTree(typeOf[Int]), Literal(Constant(33)))
      val mTempBody1 = fooDef :: mTempBody0
      val mTemp1 = Template(mTempParents, mTempSelf, mTempBody1)
      val mod1 = ModuleDef(mod0.mods, mod0.name, mTemp1)
      cd :: mod1 :: Nil

    case _ => c.abort(c.enclosingPosition, "Must annotate a class or trait")
  }
  c.Expr[Any](Block(outputs, Literal(Constant(()))))
}

This works when the companion object already exists:

object Foo
@mkCompanion class Foo

assert(Foo.hasFoo == 33)

But not when it is created:

@mkCompanion class Foo

[error] no constructor in template: impl = Object {
[error]   def hasFoo(): Int = 33
[error] }

So I still need to figure out how to provide the module constructor...

回答1:

Here is my current solution:

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

trait Foo[A]

class mkCompanion extends StaticAnnotation {
  def macroTransform(annottees: Any*) = macro mkCompanionMacro.impl
}
object mkCompanionMacro {
  def impl(c: Context)(annottees: c.Expr[Any]*): c.Expr[Any] = {
    import c.universe._

    val inputs : List[Tree] = annottees.map(_.tree)(collection.breakOut)
    val outputs: List[Tree] = inputs match {
      case (cd @ ClassDef(_, cName, _, _)) :: tail =>
        val mod0: ModuleDef = tail match {
          case (md @ ModuleDef(_, mName, mTemp)) :: Nil 
               if cName.decoded == mName.decoded => md

          case Nil =>
            val cMod  = cd.mods
            var mModF = NoFlags
            if (cMod hasFlag Flag.PRIVATE  ) mModF |= Flag.PRIVATE
            if (cMod hasFlag Flag.PROTECTED) mModF |= Flag.PROTECTED
            if (cMod hasFlag Flag.LOCAL    ) mModF |= Flag.LOCAL
            val mMod = Modifiers(mModF, cMod.privateWithin, Nil)

            // XXX TODO: isn't there a shortcut for creating the constructor?
            val mkSuperSelect = Select(Super(This(tpnme.EMPTY), tpnme.EMPTY), 
                                       nme.CONSTRUCTOR)
            val superCall     = Apply(mkSuperSelect, Nil)
            val constr        = DefDef(NoMods, nme.CONSTRUCTOR, Nil, List(Nil), 
              TypeTree(), Block(List(superCall), Literal(Constant())))

            val mTemp = Template(parents = List(TypeTree(typeOf[AnyRef])), 
              self = noSelfType, body = constr :: Nil)
            val mName = TermName(cName.decoded) // or encoded?

            ModuleDef(mMod, mName, mTemp)

          case _ => c.abort(c.enclosingPosition, "Expected a companion object")
        }

        val Template(mTempParents, mTempSelf, mTempBody0) = mod0.impl

        // cf. http://stackoverflow.com/questions/21044957/type-of-a-macro-annottee
        val cTpe        = Ident(TypeName(cd.name.decoded))
        val fooName     = TermName("hasFoo")
        val fooDef      = q"implicit def $fooName: Foo[$cTpe] = ???"
        val mTempBody1  = fooDef :: mTempBody0
        val mTemp1      = Template(mTempParents, mTempSelf, mTempBody1)
        val mod1        = ModuleDef(mod0.mods, mod0.name, mTemp1)

        cd :: mod1 :: Nil

      case _ => c.abort(c.enclosingPosition, "Must annotate a class or trait")
    }

    c.Expr[Any](Block(outputs, Literal(Constant(()))))
  }
}

Test:

object Bar
@mkCompanion class Bar

@mkCompanion class Baz

implicitly[Foo[Bar]]
implicitly[Foo[Baz]]