Tail Recursion Optimization with Groovy’s AST Transformations

The JVM is notorious for not being able to optimze tail recursive functions. Tail recursion is a special case of recursion which can be changed to iteration. In simple terms: If the call to itself is the last thing a function does before returning, then it’s tail recursive. Changing the actual execution model from recursive to iterative saves the runtime machine from putting the just finished funtion’s context on the stack. Among other advantages this saves you from getting the infamous java.lang.StackOverflowError when feeding large numbers into a recursive function in Java.

Let’s look at a classical example: Calculating the factorial of a number. A recursive – but not tail recursive! – solution is this:

int factiorial(int number) {
     if (number == 1)
          return 1;
     return number * factorial(number - 1);
}

We can see here that tail recursion is not about the recursive call being the last thing in the function’s source code but in the function’s control flow!

For many recursive problems there’s an easy way to transform it to a tail recursive solution. Often the trick is to introduce an aggregator, i.e. an additional parameter that takes the intermediate result of a calculation. Doing that our function becomes:

int factorial(int number, int aggregator) {
    if (number == 1)
        return aggregator;
    return factorial(number - 1, number * aggregator);
}

Using this function requires a start value for the aggregator, in this case factorial(10, 1) would evaluate to 3628800 which sound reasonable. So how can we translate this into an iterative version? Let’s see:

int iFactorial(int number, int aggregator = 1) {
    int _number = number;
    int _aggregator = aggregator;
    while(true) {
        if (number == 1)
            return _aggregator;
        _aggregator = _number * _aggregator;
        _number = _number - 1;
    }
}

The trick is to wrap our function into an endless loop, save the function args into local temps, replace all usages of the args with temps and replace recursive calls by setting the temps to the recursive call args. Sounds easy, eh? If the JVM won’t do that for usautomaticall why not build a recursion to iteration translator ourselves? Since Groovy is my favourite language on the JVM it seemed the perfect time and target to give my first shot at AST transformations. AST transformations are @nnotations that will be invoked at compile time to let you analyze and manipulate a programs abstract syntax tree. AST transformations can be used on different levels, in my case a local transformation seemed to suffice and I started my dive into the world of ASTNodes, expressions, statements and other weired things…

Actually, I don’t want to bother you with the details of how to test-drive AST transformations; I’m still at the very beginning of figuring things out. What I want to show you is my first prototype: tailrec 0.1. Look what we can do now:

import groovyx.transform.*
@TailRecursive
int factorial(int number, int aggregator = 1) {
    if (number == 1)
        return aggregator;
    return factorial(number - 1, number * aggregator);
}

assert factorial(10, 1) == 3628800
assert factorial(10) == 3628800

The code runs as is in Groovy’s console as soon as you add tailrec jar file to your classpath. As you can see, when making use of Groovy’s default arguments we even get rid of the clumsy aggregator initialization. To prove that tailrec really does what it promises try the next example with and without @TailRecursive:

import groovyx.transform.*
@TailRecursive //remove line to get StackOverflowError
def countDown(long from) {
    if (from == 0)
        return 0
    countDown(from - 1)
}

countDown(1000000)

So far tailrec only handles non static functions, i.e. instance methods that have some real (non-void) return type. Moreover, there are a few cases I can think of that tailrec should and could handle correctly but does not (think: ternary operator in return statement). And there are a few difficult problems where I have no current idea how to solve; consider for example the use of recursive calls within closures.

The current version is a proof of concept. I want you to experiment with it and give me your feedback. In my opinion tailrec should become part of standard GDK if it’s worth doing it at all.

Update:
The code has now moved to github. I’m actually planning to enhance and stabilize it since the groovy team encouraged me to. It meanwhile supports static functions as well and has many more tests to earn your trust.

Update 2:
@TailRecursive has become part of core Groovy since version 2.3