Varargs update support in Scala

Scala translates the following assignment call into calls to the update method.

With the following signature of update,

def update(a: Int, b: Int, c: Int, value: Boolean): Unit

One is able to write

obj(1, 2, 3) = true

and it is translated to

obj.update(1, 2, 3, true)

This is all great. However, there's one drawback.

What if I want the parameter list in the syntactic sugar to be a vararg?

obj(1, 2, 3) = true
obj(2, 3, 4, 5) = false

The ideal way to support this is to allow currying in the update method

def update(dims: Int*)(value: Boolean): Unit

Bad news, scala doesn't support this.
Good news, Macro to the rescue!

Interface design

Programmers are lazy.

So we only add an annotation to the update method definition.

@CurriedUpdate
def update[T](a: Int, b: Int, cs: Int*)(value: T): Unit = ???

Implementation

The method definition is first renamed to another fixed name, and we append an another delegate method, which matches exactly the signature needed by Scala compiler to expand the syntactic sugar.

How about the types? We cannot express the same types in a non-curried update method, because if we can, we won't be having this problem.

A possible and somehow brutal way is to define the parameters as Any*, and to then pass the type checking, use a lot of asInstanceOf to cast the params.

Seems working, right? We can do better.

Delegating macro

We define this delegating implementation as another macro, then the types don't matter!

def updateDelegate(c: Context)(args: c.Expr[Any]*): c.Expr[Any] = ???

Let's first write this macro!

It should dispatch the args to the renamed original update method, reordered. Remember we've renamed the original annotated method to another name? Yeah, that one. (Ideally there's a better way, if Scala's macro system allows to pass static parameters to macro implementation, so that we don't need to rename to a fixed name)

We split the values into two parts - the last one (which is the target value), and the others, and call the original method with two argument list. Simple enough!

def updateDelegate(c: Context)(args: c.Expr[Any]*): c.Expr[Any] = {
  import c.universe._

  c.Expr[Unit] {
    Apply(
      Apply(
        Select(c.prefix.tree, TermName(updateRenamed)),
        values.take(values.size - 1).map(_.tree).toList
      ),
      List(values.last.tree)
    )
  }
}

Macro annotation

Now let's write the macro that calls the delegating macro!

The main part - the resulting tree, is quite easy,

c.Expr[Any] {
  q"""
    import scala.language.experimental.macros

    ${renameMethod(methodDef, updateRenamed)}
    def update(values: Any*): Unit = macro CurriedUpdate.updateMacroDispatcher
  """
}

renameMethod is a utility method which renames a DefDef (method definition AST node).

Summary

Basically, we're good to go!

To wrap it up,

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

class CurriedUpdate extends StaticAnnotation {
  def macroTransform(annottees: Any*): Any = macro CurriedUpdate.curry
}

object CurriedUpdate {

  private val updateRenamed = "updateR"

  def curry(c: Context)(annottees: c.Expr[Any]*): c.Expr[Any] = {
    import c.universe._

    def renameMethod(method: DefDef, newName: String): DefDef = {
      DefDef(
        method.mods, TermName(newName), method.tparams,
        method.vparamss,
        method.tpt, method.rhs)
    }

    val methodDef = annottees(0).tree match {
      case m: DefDef => m
      case _ => c.abort(c.enclosingPosition, "CurriedUpdate can only be applied on method")
    }
    if (methodDef.name.decodedName.toString != "update") {
      c.abort(c.enclosingPosition, "CurriedUpdate can only be applied to the update method")
    }
    if (methodDef.vparamss.size != 2) {
      c.abort(c.enclosingPosition, "Curried update must have two argument list")
    }
    if (methodDef.vparamss(0).size == 0) {
      c.abort(c.enclosingPosition, "The first argument list must not be empty")
    }
    if (methodDef.vparamss(1).size != 1) {
      c.abort(c.enclosingPosition, "The second argument list must have only one element")
    }

    c.Expr[Any] {
      q"""
        import scala.language.experimental.macros

        ${renameMethod(methodDef, updateRenamed)}
        def update(values: Any*): Unit = macro CurriedUpdate.updateMacroDispatcher
      """
    }
  }

  def updateMacroDispatcher(c: Context)(values: c.Expr[Any]*): c.Expr[Unit] = {
    import c.universe._

    c.Expr[Unit] {
      Apply(
        Apply(
          Select(c.prefix.tree, TermName(updateRenamed)),
          values.take(values.size - 1).map(_.tree).toList
        ),
        List(values.last.tree)
      )
    }
  }
}

The extra code are for annotation target shape checking.

Time to test.

scala> :pa
// Entering paste mode (ctrl-D to finish)

class Test {
    @CurriedUpdate
    def update(a: Int, b: Int, cs: Int*)(value: Boolean) = {
        println(s"update $a $b ${cs mkString " "} to $value")
    }
}

defined class Test

scala> val t = new Test
t: Test = Test@3f2eb247

scala> t(1, 2, 3, 4) = false
update 1 2 3 4 to false

scala> t(1) = false
<console>:10: error: not enough arguments for method updateR: (a: Int, b: Int, c
s: Int*)(value: Boolean)Unit.
Unspecified value parameters b, cs.
              t(1) = false
                   ^


scala> t(1, 2) = true
update 1 2  to true

Seems working! Yo!