How to transform and apply a partial function using Scala Macros?

Question

I want to implement a Scala macro which takes a partial function, performs some transformations on the patterns of the function, and then applies it to a given expression.

To do so, I started with the following code:

def myMatchImpl[A: c.WeakTypeTag, B: c.WeakTypeTag](c: Context)(expr: c.Expr[A])(patterns: c.Expr[PartialFunction[A, B]]): c.Expr[B] = {
  import c.universe._

  /*
   * Deconstruct the partial function and select the relevant case definitions.
   * 
   * A partial, anonymus function will be translated into a new class of the following form:
   * 
   * { @SerialVersionUID(0) final <synthetic> class $anonfun extends scala.runtime.AbstractPartialFunction[A,B] with Serializable {
   * 
   *     def <init>(): anonymous class $anonfun = ...
   *
   *     final override def applyOrElse[...](x1: ..., default: ...): ... = ... match {
   *       case ... => ...
   *       case (defaultCase$ @ _) => default.apply(x1)
   *     }
   *
   *     def isDefined ...
   *   }
   *   new $anonfun()
   * }: PartialFunction[A,B]
   *
   */
  val Typed(Block(List(ClassDef(a, b, x, Template(d, e, List(f, DefDef(g, h, i, j, k, Match(l, allCaseDefs)), m)))), n), o) = patterns.tree

  /* Perform transformation on all cases */
  val transformedCaseDefs: List[CaseDef] = allCaseDefs map {
    case caseDef => caseDef // This code will perform the desired transformations, now it's just identity
  }

  /* Construct anonymus partial function with transformed case patterns */
  val result = Typed(Block(List(ClassDef(a, b, x, Template(d, e, List(f, DefDef(g, h, i, j, k, Match(l, transformedCaseDefs)), m)))), n), o)
  // println(show(result))

  c.Expr[B](q"$result($expr)")
}

I deconstruct the partial function, select the case definitions of the applyOrElse function, perform the desired transformation on each definition, and put everything back together. The macro is invoked like this:

def myMatch[A, B](expr: A)(patterns: PartialFunction[A, B]): B = macro myMatchImpl[A,B]

Unfortunately, this doesn't work as expected. Using the macro in a simple example

def test(x: Option[Int]) = myMatch(x){
  case Some(n) => n
  case None    => 0
}

results in the following error message:

object creation impossible, since method isDefinedAt in trait PartialFunction of type (x: Option[Int])Boolean is not defined

This is a bit confusing, since printing the generated partial function yields

({
  final <synthetic> class $anonfun extends scala.runtime.AbstractPartialFunction[Option[Int],Int] with Serializable {
    def <init>(): anonymous class $anonfun = {
      $anonfun.super.<init>();
      ()
    };
    final override def applyOrElse[A1 <: Option[Int], B1 >: Int](x2: A1, default: A1 => B1): B1 = ((x2.asInstanceOf[Option[Int]]: Option[Int]): Option[Int] @unchecked) match {
      case (x: Int)Some[Int]((n @ _)) => n
      case scala.None => 0
    };
    final def isDefinedAt(x2: Option[Int]): Boolean = ((x2.asInstanceOf[Option[Int]]: Option[Int]): Option[Int] @unchecked) match {
      case (x: Int)Some[Int]((n @ _)) => true
      case scala.None => true
      case (defaultCase$ @ _) => false
    }
  };
  new $anonfun()
}: PartialFunction[Option[Int],Int])

which clearly defines the isDefinedAt method.

Does anybody has an idea, what's the problem here and how to do this right?

Answer 1

The new reflection API provides a Transformer class that's specifically designed to help with this kind of tree transformation:

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

def myMatchImpl[A: c.WeakTypeTag, B: c.WeakTypeTag](c: Context)(
  expr: c.Expr[A]
)(
  patterns: c.Expr[PartialFunction[A, B]]
): c.Expr[B] = {
  import c.universe._

  val transformer = new Transformer {
    override def transformCaseDefs(trees: List[CaseDef]) = trees.map {
      case caseDef => caseDef
    }
  }

  c.Expr[B](q"${transformer.transform(patterns.tree)}($expr)")
}

def myMatch[A, B](expr: A)(patterns: PartialFunction[A, B]): B =
  macro myMatchImpl[A,B]

def test(x: Option[Int]) = myMatch(x) {
  case Some(n) => n
  case None    => 0
}

You may need some additional machinery to make sure that the transformation is only being applied to the case lists you want it to be applied to, but in general this approach is going to be more robust than transforming the tree manually.

I'm still curious about why your version doesn't work, though, and if you have the time it might be worth putting together a reduced example for another question here.