Splicing a passed function body into a macro-rewritten expression

Question

I was playing with Scala 2.11's new macro features. I wanted to see if I could do the following rewrite:

forRange(0 to 10) { i => println(i) }

// into

val iter = (0 to 10).iterator
while (iter.hasNext) {
  val i = iter.next
  println(i)
}

I think I got fairly close with this macro:

def _forRange[A](c: BlackboxContext)(range: c.Expr[Range])(func: c.Expr[Int => A]): c.Expr[Unit] = {
  import c.universe._

  val tree = func.tree match {
    case q"($i: $t) => $body" => q"""
        val iter = ${range}.iterator
        while (iter.hasNext) {
          val $i = iter.next
          $body
        }
      """
    case _ => q""
  }

  c.Expr(tree)
}

This produces the following output when called as forRange(0 to 10) { i => println(i) } (at least, it's what the show function gives me on the resultant tree):

{
  val iter = scala.this.Predef.intWrapper(0).to(10).iterator;
  while$1(){
    if (iter.hasNext)
      {
        {
          val i = iter.next;
          scala.this.Predef.println(i)
        };
        while$1()
      }
    else
      ()
  }
}

That looks like it should work, but there's a conflict between my manually defined val i and the i referenced in the spliced-in function body. I get the following error:

ReplGlobal.abort: symbol value i does not exist in$line38.$read$$iw$$iw$$iw$$iw$$iw$$iw$$iw$$iw. error: symbol value i does not exist in scala.reflect.internal.FatalError: symbol value i does not exist in $line38.$read$$iw$$iw$$iw$$iw$$iw$$iw$$iw$$iw.

And then a rather large stack trace, resulting in an "Abandoned crashed session" notification.

I can't tell if this is a problem with my logic (you simply can't splice in a function body that references a closed-over variable), or if it's a bug with the new implementation. The error reporting certainly could be better. It may be exacerbated by the fact that I'm running this on the Repl.

Is it possible to pull apart a function, separating the body from the closed-over terms, and rewrite it in order to splice the logic directly into a resulting tree?

Answer 1

When in doubt, resetAllAttrs:

import scala.language.experimental.macros
import scala.reflect.macros.BlackboxContext

def _forRange[A](c: BlackboxContext)(range: c.Expr[Range])(
  func: c.Expr[Int => A]
): c.Expr[Unit] = {
  import c.universe._

  val tree = func.tree match {
    case q"($i: $t) => $body" => q"""
        val iter = ${range}.iterator
        while (iter.hasNext) {
          val $i = iter.next
          ${c.resetAllAttrs(body)} // The only line I've changed.
        }
      """
    case _ => q""
  }

  c.Expr(tree)
}

And then:

scala> def forRange[A](range: Range)(func: Int => A) = macro _forRange[A]
defined term macro forRange: [A](range: Range)(func: Int => A)Unit

scala> forRange(0 to 10) { i => println(i) }
0
1
2
3
4
5
6
7
8
9
10

In general, when you're grabbing a tree from one place and plopping it somewhere else, it's likely going to be necessary to use resetAllAttrs to get all the symbols right.

Answer 2

Oscar Boykin pointed out on Twitter that my previous answer no longer works, and it wasn't a very complete answer anyway—it addresses the problem pointed out by the OP on Scala 2.10, but it's not careful about hygiene—if you wrote iter => println(iter) you'd get a compile-time failure, for example.

A better implementation for 2.11 would use a Transformer to rewrite the tree after un-typechecking it:

import scala.language.experimental.macros
import scala.reflect.macros.blackbox.Context

def _forRange[A](c: Context)(r: c.Expr[Range])(f: c.Expr[Int => A]): c.Tree = {
  import c.universe._

  f.tree match {
    case q"($i: $_) => $body" =>
      val newName = TermName(c.freshName())
      val transformer = new Transformer {
        override def transform(tree: Tree): Tree = tree match {
          case Ident(`i`) => Ident(newName)
          case other => super.transform(other)
        }
      }

      q"""
        val iter = ${r.tree}.iterator
        while (iter.hasNext) {
          val $newName = iter.next
          ${ transformer.transform(c.untypecheck(body)) }
        }
      """
  }
}

def forRange[A](r: Range)(f: Int => A): Unit = macro _forRange[A]

Which works like this:

scala> forRange(0 to 10)((i: Int) => println(i))
0
1
2
3
4
5
6
7
8
9
10

Now it doesn't matter what variable name we use in our function literal, since it'll just be replaced with a fresh variable anyway.