Create or extend a companion object, using a macro annotation on the class

Question

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...

Yadirayaeger

The subsequent problem is getting the type of cd – Yadirayaeger

Yeah, you need to create the constructor manually. – Monarchy

Answer 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. https://mcmap.net/q/1626210/-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]]

Yadirayaeger

Have you looked at quasiquotes? They do most of the hard work of creating internal tree representations for you. I, for one, have long stopped constructed trees manually and do that only when hacking the compiler, where we can't use qqs yet :) – Monarchy

@EugeneBurmako - yes I meant to. I haven't found the time yet to deep dive so to understand how they actually work in detail. I'm not that uncomfortable with trees, though. – Yadirayaeger