Blog

Compile-Time Evaluation in Scala with macros

27 Mar, 2016
Xebia Background Header Wave

Many ‘compiled’ languages used to have a strict separation between what happens at ‘compile-time’ and what happens at ‘run-time’. This distinction is starting to fade: JIT compilation moves more of the compile phase to run-time, while conversely various kinds of optimizations do ‘run-time’ work at compile time. Powerful type systems allow the expression things previously only checked at run time, especially with the recent renaissance of dependent types popularized by Idris.
This post will show a very simple example of compile-time evaluation in Scala: we’ll write a regular ‘factorial’ function, and use macros to apply it (to constants) at compile time.

Our factorial

Let’s start with our factorial function:
[code language="scala"]
def normalFactorial(n: Int): Int =
if (n == 0) 1
else n * factorial(n – 1)
[/code]
This is fairly uneventful, but especially notice that this is a plain old Scala function, nothing fancy.

Defining a macro function

When defining a macro function, we give a signature in regular Scala syntax. The implementation consists of the special ‘macro’ keyword and a reference to the macro implementation.
[code language="scala"]
import scala.language.experimental.macros
def factorial(n: Int): Int = macro factorial_impl
import scala.reflect.macros.blackbox.Context
def factorial_impl(c: Context)(n: c.Expr[Int]): c.Expr[Int] = ???
[/code]
The macro implementation, then, is also a regular Scala function. The difference is that instead of ‘normal’ values, it receives and produces AST’s.

Implementing the macro function

To implement our compile-time factorial, we must unwrap the AST, apply our function, and produce a new AST that will be placed at the call site. Because the AST is expressed in regular Scala code, we can pattern-match on it, for example recognizing a literal Int constant:
[code language="scala"]
import scala.reflect.macros.blackbox.Context
def factorialimpl(c: Context)(n: c.Expr[Int]): c.Expr[Int] = {
import c.universe.

n match {
case Expr(Literal(Constant(nValue: Int))) =>
val result = normalFactorial(nValue)
c.Expr(Literal(Constant(result)))
case _ =>
println("Yow!")
???
}
}
[/code]
Keep in mind this code is executed at compile time: even our cheeky little message will be printed when the compiler encounters a non-matching AST, and the compilation will fail due to the NotImplementedError.

The final result

Combining all the previous steps, our macro-based compile-time factorial implementation now looks like this:
[code language="scala"]
object CompileTimeFactorial {
import scala.language.experimental.macros
// This function exposed to consumers has a normal Scala type:
def factorial(n: Int): Int =
// but it is implemented as a macro:
macro CompileTimeFactorial.factorial_impl
import scala.reflect.macros.blackbox.Context
// The macro implementation will receive a ‘Context’ and
// the AST’s of the parameters passed to it:
def factorialimpl(c: Context)(n: c.Expr[Int]): c.Expr[Int] = {
import c.universe.

// We can pattern-match on the AST:
n match {
case Expr(Literal(Constant(nValue: Int))) =>
// We perform the calculation:
val result = normalFactorial(nValue)
// And produce an AST for the result of the computation:
c.Expr(Literal(Constant(result)))
case other =>
// Yes, this will be printed at compile time:
println("Yow!")
???
}
}
// The actual implementation is regular old-fashioned scala code:
private def normalFactorial(n: Int): Int =
if (n == 0) 1
else n * normalFactorial(n – 1)
}
[/code]
For those of you following along at home, save this implementation in CompileTimeFactorial.scala.

Using the macro

Using the macro is as simple as calling a regular Scala function:
[code language="scala"]
import CompileTimeFactorial._
object Test extends App {
println(factorial(10))
// When uncommented, this will produce an error at compile-time, as we
// only implemented a case for an Int literal, not a variable:
// val n = 10
// println(factorial(n))
}
[/code]
There are a couple of caveats: to be able to evaluate the macro when compiling Test.scala above, CompileTimeFactorial.scala must have been previously compiled. Compiling both Test.scala and CompileTimeFactorial.scala in the same scalac invocation will not work reliably.
Also, this example illustrates that while factorial has a regular Scala type definition, the fact that macro’s perform AST manipulations does leak through. The bottom invocation, while semantically trivially equivalent to the upper one, will not compile: we only implemented our macro for Int literals, not for variables.

Did it work?

Indeed our program produces the correct output:
[code]
$ scalac CompileTimeFactorial.scala
$ scalac Test.scala
$ scala Test
3628800
$
[/code]
But can we prove this value was indeed calculated at compile time? Turns out we can, using the javap disassembler that comes with the JVM.
I’ll spare you the complete output, but the relevant bit is:
[code]
$ javap -c Test$
….
Code:
0: getstatic #61 // Field scala/Predef$.MODULE$:Lscala/Predef$;
3: ldc #62 // int 3628800
5: invokestatic #68 // Method scala/runtime/BoxesRunTime.boxToInteger:(I)Ljava/lang/Integer;
8: invokevirtual #72 // Method scala/Predef$.println:(Ljava/lang/Object;)V
11: return
[/code]
Showing the constant integer 3628800 is loaded, boxed and printed.

Conclusions

Of course this is a toy example, meant to illustrate the concept of Scala macro’s. When doing more serious AST matching and manipulation you’ll certainly want to look at quasiquotes, and the extra power of macro bundles.
This kind of metaprogramming has turned out to be hard to support across versions of the compiler. This lead to the scala.meta initiative to decouple metaprogramming from the compiler ⊢ though it will probably be a while before they get to macros. In the mean time, macro-compat might make it easier to write macro’s that support multiple compiler versions.
Applications of macro’s include:

It seems to be early days for metaprogramming in Scala still, but the potential is amazing.

Questions?

Get in touch with us to learn more about the subject and related solutions

Explore related posts