Overriding class method in compiler plugin

Question
1

I’m working on a compiler plugin to generate a “clone” method that just needs to be able to create a fresh instance of a class with the same parameters as a given instance.

My plugin transform method is as follows:

def isBundle(sym: Symbol): Boolean =
  sym.fullName == "<package>.Bundle" || sym.parentSymbolsIterator.exists(isBundle)

override def transform(tree: Tree): Tree = tree match {
  case bundle: ClassDef if isBundle(bundle.symbol) =>
    bundle match {
      case q"""$mods class $tpname[..$tparams] $ctorMods(...$paramss)
                   extends { ..$earlydefns } with ..$parents { $self =>
                     ..$stats
            }""" =>
        println(s"========== Matched on $tpname ==========")
        println(bundle)
        // We know the types but the compiler doesn't
        val paramssTyped = paramss.asInstanceOf[List[List[ValDef]]]
        val tparamsTyped = tparams.asInstanceOf[List[TypeDef]]

        val genMethod = {
          // Get refs to the Trees
          val paramssRefs = paramssTyped.map(_.map(_.name))
          val tparamsRefs = tparamsTyped.map { t => tq"${t.name}" }
          val impl = q"override def cloneType = new $tpname[..$tparamsRefs](...$paramssRefs)"
          println(impl)
          // This fails
          localTyper.typed(impl)
        }

        val ClassDef(mods0, name0, tparams0, template@Template(parents0, self0, body0)) = bundle
        val templateResult = treeCopy.Template(template, parents0, self0, body0 :+ genMethod)
        val result = treeCopy.ClassDef(bundle, mods0, name0, tparams0, templateResult)
        // If I comment out the typing of impl, this returns but doesn't seem to fill in the types for cloneType
        localTyper.typed(result)

      // Some classes don't match the quasiquotes above
      case _ => super.transform(tree)
    }
  case _ => super.transform(tree)
}

When I try testing this on a relatively simple example it fails:

class MyBundle(x: Int) extends Bundle {
  val foo = Output(UInt(x.W))
}

This gives:

========== Matched on MyBundle ==========
class MyBundle extends chisel3.Bundle {
  <paramaccessor> private[this] val x: Int = _;
  def <init>(x: Int): MyBundle = {
    MyBundle.super.<init>();
    ()
  };
  private[this] val foo: chisel3.UInt = chisel3.Output.apply[chisel3.UInt](chisel3.`package`.UInt.apply(chisel3.`package`.fromIntToWidth(MyBundle.this.x).W));
  <stable> <accessor> def foo: chisel3.UInt = MyBundle.this.foo
}
override def cloneType = new MyBundle(x)
[error] ## Exception when compiling 1 sources to .../target/scala-2.12/classes
[error] scala.reflect.internal.Types$TypeError: not found: value x

I’ve also tried running the typer on the full class after adding the method to the class body, but it doesn’t seem like the typer recurses (the tree it returns doesn’t have any types filled in).

I have looked at examples like data-class (https://github.com/alexarchambault/data-class/blob/797de5dd0cd5cf8804bca2088ceaa1c7b96f6815/src/main/scala/dataclass/Macros.scala) but I suspect that macros don’t need to invoke the typer themselves since I don’t see any use of it there.

Any tips or examples of generating a method in a compiler plugin? Thanks!

1 Like
2

I think I’m making progress by using showRaw on the DefDef when I implemented it manually and comparing it to the DefDef when I’m building it in genMethod. I’ve also been using localTyper.typed on smaller subsets of the DefDef. Now I have:

val genMethod = {
  val `this` = This(bundle.impl.symbol.owner)
  val paramssRefs = paramssTyped.map(_.map(p => Select(`this`, p.name)))
  val tparamsRefs = tparamsTyped.map { t => tq"${t.name}" }
  val thisDotType = TypeTree().setOriginal(SingletonTypeTree(`this`))
  val impl = q"override def cloneType = new ${bundle.symbol}[..$tparamsRefs](...$paramssRefs).asInstanceOf[$thisDotType]"
  println(impl)
  // Still fails
  localTyper.typed(impl)
}

val `this` = This(bundle.impl.symbol.owner) seems to be right, but Select(`this`, p.name) (where p is a ValDef from the paramss) errors with:

scala.reflect.internal.Types$TypeError: value x is not a member of MyBundle

I’m trying out gen.mkAttributedSelect, but the paramss ValDefs don’t seem to have valid .symbols.

3

Today’s update (in case anyone comes across this stream of consciousness someday):

  • Quasiquotes (at least unapply form) should not be used in plugins because the Trees that come out don’t have symbols. For example, if you want to get a legal reference to a class parameter, you need the right ValDef object with its .symbol field correctly populated, and you don’t get that from the quasiquote unapply I used in my original question.
  • gen.mk* methods are useful and seem to build the right things so long as you create them from the right Trees.
  • Reiterating that showRaw, especially with the extra printIds argument set are great. Just writing out the method I want to generate and comparing the individual pieces as I build them up seems to be moving things in the right direction.
1 Like
4

A final update because I was ultimately able to get things to work! Some tips in addition to the ones from my previous update:

  • Symbols are important, pay close attention to how they are created and that they have the right flags and info (which seems to be a Type).

In my case, for adding a method, it was important to create the new method symbol from the symbol of the class to which I’m adding the method (myClass.symbol.newMethod(...) returns a MethodSymbol).

  • Do as little as possible in the compiler plugin (or put another way, do anything you can at runtime).

The method I’m trying to override has the signature def cloneType: this.type. I found it very difficult to properly override a method of type this.type, and actually had similar issues with any narrowing of return types, so the trick is to just not do that. I ended up doing something like this:

 // The supertype of the classes for which I generate the method
trait Bundle {
  // The actual method I override
  protected def _cloneTypeImpl: Bundle = ???
  // Eventually it will be final, but my users currently override it
  def cloneType: this.type = _cloneTypeImpl.asInstanceOf[this.type]
}

See how _cloneTypeImpl just returns the type of the superclass? This avoids this whole issue with JVM bridge methods that Scalac has to generate whenever you have return type narrowing. I know the erasure phase of scalac generates these bridges, but I just could not figure out how to get it to work for my generated methods. So the trick is to just avoid the issue :slight_smile:

My final transformer is as follows:

  private class MyTypingTransformer(unit: CompilationUnit) extends TypingTransformer(unit) {

    val bundleTpe = localTyper.typed(tq"chisel3.Bundle", nsc.Mode.TYPEmode).tpe

    def isBundle(sym: Symbol): Boolean =
      sym.fullName == "chisel3.Bundle" || sym.parentSymbolsIterator.exists(isBundle)

    def isCloneMethod(name: String, defdef: DefDef): Boolean =
      defdef.name.decodedName.toString == name && defdef.tparams.isEmpty && defdef.vparamss.isEmpty

    def getParamsAndConstructor(body: List[Tree]): (Option[DefDef], Seq[Symbol]) = {
      val paramAccessors = mutable.ListBuffer[Symbol]()
      var primaryConstructor: Option[DefDef] = None
      body.foreach {
        case acc if acc.symbol.isParamAccessor =>
          paramAccessors += acc.symbol
        case con: DefDef if con.symbol.isPrimaryConstructor =>
          primaryConstructor = Some(con)
        case d: DefDef if isCloneMethod("cloneType", d) =>
          global.reporter.warning(d.pos, "The Chisel compiler plugin now implements cloneType")
        case d: DefDef if isCloneMethod("_cloneTypeImpl", d) =>
          global.globalError(d.pos, "Users should never implement _cloneTypeImpl")
        case _ =>
      }
      (primaryConstructor, paramAccessors.toList)
    }

    override def transform(tree: Tree): Tree = tree match {

      case bundle: ClassDef if isBundle(bundle.symbol) && !bundle.mods.hasFlag(Flag.ABSTRACT) =>

        val (con, params) = getParamsAndConstructor(bundle.impl.body)
        if (con.isEmpty) {
          global.reporter.warning(bundle.pos, "Unable to determine primary constructor!")
          return super.transform(tree)
        }
        val constructor = con.get

        val thiz = gen.mkAttributedThis(bundle.symbol)

        // Create a this.<ref> for each field matching order of constructor arguments
        // List of Lists because we can have multiple parameter lists
        val conArgs: List[List[RefTree]] =
          constructor.vparamss.map(_.map { vp =>
            params.collectFirst {
              case p if vp.name == p.name => gen.mkAttributedSelect(thiz, p)
            }.get
          })

        val ttpe = Ident(bundle.symbol)
        val neww = localTyper.typed(New(ttpe, conArgs))

        // Create the symbol for the method and have it be associated with the Bundle class
        val cloneTypeSym =  bundle.symbol.newMethod(TermName("_cloneTypeImpl"), bundle.symbol.pos.focus, Flag.OVERRIDE)
        // Handwritten cloneTypes don't have the Method flag set, unclear if it matters
        cloneTypeSym.resetFlag(Flags.METHOD)
        // Need to set the type to chisel3.Bundle for the override to work
        cloneTypeSym.setInfo(NullaryMethodType(bundleTpe))

        val defdef = localTyper.typed(DefDef(cloneTypeSym, neww))

        val withMethod = deriveClassDef(bundle) { t =>
          deriveTemplate(t)(_ :+ defdef)
        }
        localTyper.typed(withMethod)

      case _ => super.transform(tree)
    }
  }

I hope this is useful to someone someday.

1 Like