Create or extend a companion object using macro annotation in a class

Question

Using the macro annotation macro of a Scala 2.10 / 2.11 macro, how can I add or extend a companion object to 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] = ??? } 

Thus, given

 trait Foo[A] 

next entry

 @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 the first naive attempt by simply adding def hasFoo = 33 :

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

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