Bartosz Bąbol

Software engineering

Scala Macros: Part II- Macro Annotations

I. Introduction. What are macro annotations?

With macro annotations you can annotate any definition with something that Scala recognizes as a macro and it will give you ability to modify arbitrarily this definition. Personally speaking, it is my favorite type of macro with many cool use cases.

Code is on my github, branch part-2. Let’s start!

II. Setup

Setup is the same as in part 1

III. Example no.2: Benchmark

In part 1 we created pretty much useless macro, just to become familiar with new quasiquotes api and macro project structure. Now let’s try to create something more useful. Imagine following problem. We want to benchmark methods. We would like to check how much time takes method to execute. This is our method. As simple as it could be. Type parameter is a fake, just to make this method look more fancy.

1
2
3
4
5
6

object Test{
    def testMethod[String]: Double = {
        val x = 2.0 + 2.0
        Math.pow(x, x)
    }
}

Now how to measure running time of body of this method? This is one possibility:

1
2
3
4
5
6
7
8
9
10
11
12

object Test{
    def testMethod[String]: Double = {
        val start = System.nanoTime()
        val result = {
            val x = 2.0 + 2.0
            Math.pow(x, x)
        }
        val end = System.nanoTime()
        println("testMethod elapsed time: " + (end - start) + "ns")
        result
    }
}

So we wrap body of function with time snapshots then we assign result of #testMethod() to val result, then we println time difference and return result. The problem with it is we had to touch #testMethod() and modify it. Better solution would be to not touch code of #testMethod() at all. Moreover it is boilerplate code, not interesting for developer at all.

In Test.scala you will find one possible solution could be, for example:

Test.scala
1
2
3
4
5
6
7
8
9
10
11
12
13

object Test{
  @Benchmark
  def testMethod[String]: Double = {
    val x = 2.0 + 2.0
    Math.pow(x, x)
  }

  @Benchmark
  def methodWithArguments(a: Double, b: Double) = {
    val c = Math.pow(a, b)
    c > a+b
  }
}

And we want that result of above code to be the same as previously. We can easily do this with macro annotations.

Look at Main.scala, there’s usage of our methods:

Main.scala
1
2
3
4

object Main extends App{
  Test.testMethod
  //...
}

Now check the implementation of @Benchmark

Benchmark.scala
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

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

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

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

    val result = {
      annottees.map(_.tree).toList match {
        case q"$mods def $methodName[..$tpes](...$args): $returnType = { ..$body }" :: Nil => {
          q"""$mods def $methodName[..$tpes](...$args): $returnType =  {
            val start = System.nanoTime()
            val result = {..$body}
            val end = System.nanoTime()
            println(${methodName.toString} + " elapsed time: " + (end - start) + "ns")
            result
          }"""
        }
        case _ => c.abort(c.enclosingPosition, "Annotation @Benchmark can be used only with methods")
      }
    }
    c.Expr[Any](result)
  }
}

There are few differences between def macros and macro annotations when looking at their implementations.

  1. First of all class Benchmark has to extend StaticAnnotation trait.

  2. Second difference is we need to implement macroTransform method which take annottees: c.Expr[Any]* as argument. You might think about Expr as wrapper around AST.

Let’s focus on implementation:

Benchmark.scala 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

...
    val result = {
      annottees.map(_.tree).toList match {
        case q"$mods def $methodName[..$tpes](...$args): $returnType = { ..$body }" :: Nil => {
          q"""$mods def $methodName[..$tpes](...$args): $returnType =  {
            val start = System.nanoTime()
            val result = {..$body}
            val end = System.nanoTime()
            println(${methodName.toString} + " elapsed time: " + (end - start) + "ns")
            result
          }"""
        }
        case _ => c.abort(c.enclosingPosition, "Annotation @Benchmark can be used only with methods")
      }
    }
    c.Expr[Any](result)
  }
}

What is going on here?

  1. Annotees is Seq of annotated definitions. We want to have them in form of AST this is why we map over this Seq and return new Seq with AST nodes.

  2. Next outstanding feature of quasiquotes is extractor pattern and this is why we can use pattern matching over Tree node. And you will find in docs syntax summary which pattern corresponds to which scala construction. In our case we will annotate methods so we used syntax for extracting method tokens.

Look at this weird syntax:

Benchmark.scala 
1

$methodName[..$tpes](...$args)

Ok, so we are extracting methodName but what are those “..” and “…” signs means?

Let’s start with ..$- this pattern expects List[universe.Tree]. And this is nice because our annotated method could take many type parameters.

And what is …$- this pattern expects List[List[universe.Tree]]. This is becase our method can take many parameters sets, so it could look like this one:

1
2
3

private def foo[A, B, C](a:A ,b:B)(c: C): A = {
//body
}

In conclusion we are matching each annottee to this method pattern, we extract from annottee(method in our case) all potentially useful elements like method name or parameter list etc. And we are returning new collection of modified AST’s. In this example we return the same method signature with modified body. You see that inside quasiquote we are doing what we actually expected. So we wrap code with those timestamps, and println time difference between endTime and startTime.

If we would apply @Benchmark to some other definition like class or val or whatever, then preventing from match error, I’ve added default case:

1

    case _ => c.abort(c.enclosingPosition, "Annotation @Benchmark can be used only with methods")

IV. Example no.3: Talking Animal

Macros can be treated as some kind of magic. Previous examples have shown that they can modify AST and slightly change behaviour of your code. But they can do much more than that. Let’s look at Main.scala:

Main.scala
1
2
3
4
5

object Main extends App{
  //...

  Dog("Szarik").sayHello
}

Something is wrong here, I personally use Intellij Idea and it doesn’t see this method sayHello on object Dog. But this code compiles properly. To solve this mystery open Animal.scala:

Animal.scala
1
2
3
4
5
6

trait Animal{
  val name: String
}

@TalkingAnimalSpell
case class Dog(name: String) extends Animal

We have got simple case class Dog which extends Animal trait. It doesn’t have method sayHello implemented, but there is @TalkingAnimalSpell annotation. So let’s look at implementation of it:

TalkingAnimalSpell.scala
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

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

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

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

    val result = {
      annottees.map(_.tree).toList match {
        case q"$mods class $tpname[..$tparams] $ctorMods(...$paramss) extends Animal with ..$parents { $self => ..$stats }" :: Nil => {
          val animalType = tpname.toString()
          q"""$mods class $tpname[..$tparams] $ctorMods(...$paramss) extends Animal with ..$parents{
            def sayHello: Unit = {
              println("Hello I'm " + $animalType + " and my name is " + name)
            }
          }"""
        }
        case _ => c.abort(c.enclosingPosition, "Annotation @TalkingAnimal can be used only with case classes which extends Animal trait")
      }
    }
    c.Expr[Any](result)
  }
}

Difference between this and previous benchmark example is in pattern matching. Case condition is different. In this TalkingAnimalSpell macro we want to annotate classes not methods like in benchmark. In quasisquotes syntax summary you will find pattern for classes. And because I want this annotation to work only with classes which extends Animal trait I specified it explicitly in pattern case.

What this macro does is returning the same object but with added new method sayHello which println the name argument. Writting this blog post I’ve noticed a bug. Does it return exactly the same object? Modify a little bit Dog class:

TalkingAnimalSpell.scala
1
2
3
4
5
6

@TalkingAnimalSpell
case class Dog(name: String) extends Animal {
    def apport: Unit = {
        println("Apporting...")
    }
}

And invoke it in Main.scala

TalkingAnimalSpell.scala
1
2
3
4
5

object Main extends App{
  //...

  Dog("Szarik").apport
}

And you get compilation error :) Everything looks fine, method is implemented Idea can see it but you’ve got compilation error. What is wrong? Clearly @TalkingAnimalSpell is really some kind of magic, look at its implementation again:

We are returning class with the same signature, but we are missing existing body of annotated class. Let’s add missing body:

TalkingAnimalSpell.scala
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

    //...
    val result = {
      annottees.map(_.tree).toList match {
        case q"$mods class $tpname[..$tparams] $ctorMods(...$paramss) extends Animal with ..$parents { $self => ..$stats }" :: Nil => {
          val animalType = tpname.toString()
          q"""$mods class $tpname[..$tparams] $ctorMods(...$paramss) extends Animal with ..$parents{
            $self => ..$stats
            def sayHello: Unit = {
              println("Hello I'm " + $animalType + " and my name is " + name)
            }
          }"""
        }
        case _ => c.abort(c.enclosingPosition, "Annotation @TalkingAnimal can be used only with case classes which extends Animal trait")
      }
    }
    c.Expr[Any](result)
  }
}

And run code again, everything should compile fine. So as you see you can make arbitrary change with you annotated code. With great power comes great responsibility. I encourage you to play a little bit with this example.

V. Homework

Maybe some kind of good homework exercises would be

1) Imagine that you are funny software developer who likes to make jokes. Change implementation and returning type of

1

def apport

2) change implementation of sayHello to println each animal attributes. For example, you’ve got this class with @TalkingAnimalSpell annotation:

Animal.scala
1
2

@TalkingAnimalSpell
case class Cat(name: String, favoriteFood: String, race: String, color: String) extends Animal

And invoking sayHello on this object:

Main.scala
1
2
3
4

object Main extends App{
  //...
  Cat("Tom", "Whiskas", "persian", "black").sayHello
}

should result with this println:

1

Hello I'm Cat and my name is Tom my favorite food is whiskas, my race is persian, my color is black.

3) Change @Benchmark annotation. It should be possible to annotate scala object and it will add this benchmark code to each non private method inside this object. So this code, after compilation:

1
2
3
4
5
6
7
8
9
10
11
12
13

object Test{
  @Benchmark
  def testMethod[String]: Double = {
    val x = 2.0 + 2.0
    Math.pow(x, x)
  }

  @Benchmark
  def methodWithArguments(a: Double, b: Double) = {
    val c = Math.pow(a, b)
    c > a+b
  }
}

should work the same as this code:

1
2
3
4
5
6
7
8
9
10
11
12

@Benchmark
object Test{
  def testMethod[String]: Double = {
    val x = 2.0 + 2.0
    Math.pow(x, x)
  }

  def methodWithArguments(a: Double, b: Double) = {
    val c = Math.pow(a, b)
    c > a+b
  }
}

VI. Summary of part II

In this post, we explored another type of macro: macro annotations. But there are still some question marks. One of them is highlighted by last example with @TalkingAnimalSpell and exercise no. 1 from your homework. Intellij Idea coding assistance is based on static code analysis and it is not aware of AST changes, so there is lack of support for macros.

You probably feel disappointed now. What if your macro generates methods that are invisible for your IDE? Your code will be highlighted with red , you wouldn’t know returning type of generated methods, moreover somebody will use macros to change existing code etc. Without documentation this macro could be more confusing than useful. Thankfully Intellij Idea has API for writing plugins to support macros. In Part 3 of this blog series, we will create plugin for Intellij Idea to add support to our @TalkingAnimalSpell macro annotation. See you in Part 3!

Comments