Scala compiler plugin to rewrite method calls

Question

I'm trying to write a Scala compiler plugin that rewrites method calls from instantiating Set to instantiations of LinkedHashSet. Unfortunately I can't find any working example that does this already. The following code fails with no-symbol does not have an owner:

object DemoErasureComponent extends PluginComponent with TypingTransformers with Transform {
val global: DemoPlugin.this.global.type = DemoPlugin.this.global

import global._

override val runsAfter = List("erasure")

val phaseName = "rewrite-sets"

def newTransformer(unit: CompilationUnit) = new SetTransformer(unit)

class SetTransformer(unit: CompilationUnit) extends TypingTransformer(unit) {

  override def transform(tree: Tree): Tree = tree match {
    case a@Apply(r@Select(rcvr@Select(predef, set), name), args) if name.toString == "Set" =>
      localTyper.typed(treeCopy.Apply(tree, Ident(newTermName("LinkedHashSet")), args))

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

For the record, I've found these resources so far:

• Scalaxy compiler plugin: https://github.com/ochafik/Scalaxy
• Boxer compiler plugin example: https://github.com/retronym/boxer
• outdated compiler plugin inside the Scala compiler: http://lampsvn.epfl.ch/trac/scala/browser/scala/trunk/docs/examples/plugintemplate/src/plugintemplate/TemplateTransformComponent.scala

Answer 1

localTyper.typed(treeCopy.Apply(tree, Ident(newTermName("LinkedHashSet")), args))

Here, you are creating a new Apply node via a tree copy, which will copy the type and symbol from tree.

When you typecheck this node, the typer will not recurse into its children as it is already typed, so the Ident will pass through without a type and symbol, which will likely lead to a crash in the code generation phase.

Instead of creating an unattributed Ident and typechecking it, it is more customary to create a fully attrtibuted reference with of the utility methods in TreeGen.

gen.mkAttributedRef(typeOf[scala.collection.immutable.LinkedHashSet].typeSymbol)

The case is also pretty suspicious. You should never need to compare strings like that. It is always better to compare Symbols.

Furthermore, you have to be aware of the tree shapes at the phase you're inserting your compiler plugin. Notice below how after typer, the tree is expanded to include the call to the apply method, and that after erasure, the variable arguments are wrapped in single argument.

% qscalac -Xprint:parser,typer,erasure sandbox/test.scala
[[syntax trees at end of                    parser]] // test.scala
package <empty> {
  object Test extends scala.AnyRef {
    def <init>() = {
      super.<init>();
      ()
    };
    Set(1, 2, 3)
  }
}

[[syntax trees at end of                     typer]] // test.scala
package <empty> {
  object Test extends scala.AnyRef {
    def <init>(): Test.type = {
      Test.super.<init>();
      ()
    };
    scala.this.Predef.Set.apply[Int](1, 2, 3)
  }
}

[[syntax trees at end of                   erasure]] // test.scala
package <empty> {
  object Test extends Object {
    def <init>(): Test.type = {
      Test.super.<init>();
      ()
    };
    scala.this.Predef.Set().apply(scala.this.Predef.wrapIntArray(Array[Int]{1, 2, 3}))
  }
}

Tracking down non-determinism that is caused by making assumptions about the iteration order of HashMaps can be a bit of a nightmare. But I'd caution against this sort of rewriting. If you want a policy within your system to say that non-determistic Sets and Maps should not be used, removing direct usages of them is insufficient. Every call to .toSet or toMap will end up using them anyway.

You might be better served to instrument the bytecode of the standard library (or use a patched standard library) in a test mode to catch all instantiations of these structures (perhaps logging a stack trace). Or, as a more fine grained alternative, find calls to place like HashTrieSet#foreach (although, you'll need to filter out benign usages like the call to foreach within collection.immutable.HashSet(1, 2, 3).count(_ => true).