Home Kotlin/Wasm Tail Calls: What We Can and Can't Optimize Today
Post
Cancel

Kotlin/Wasm Tail Calls: What We Can and Can't Optimize Today

Posted Jun 24 Updated Aug 12 9 min read

Patterns that are optimized

Static dispatch

PR #2 makes any direct function call in tail position emit return_call instead of call.

1
2
3
4
5
6
7
8
9

fun isEven(n: Int): Boolean {
    if (n == 0) return true
    return isOdd(n - 1)   // emits return_call to isOdd
}

fun isOdd(n: Int): Boolean {
    if (n == 0) return false
    return isEven(n - 1)  // emits return_call to isEven
}

tailrec rewrites direct self-recursion into a loop. Mutual recursion between isEven and isOdd overflows at depth ~10K on V8 without tail calls but runs to 1M in constant stack with them.

The implementation also covers self-recursion without the tailrec annotation.

1
2
3
4

fun sumTo(n: Int, acc: Int = 0): Int {
    if (n == 0) return acc
    return sumTo(n - 1, acc + n)  // emits return_call (self)
}

TailrecLowering rewrites tailrec fun sumTo(...) into a do-while loop, which is ~20% faster. The compiler leaves tailrec functions alone and only emits native tail calls for unmarked ones.

Virtual dispatch

PR #3 extends tail call emission to virtual and interface dispatch. Calling an open or abstract method in tail position emits return_call_ref via the vtable.

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

sealed class Expr {
    abstract fun eval(env: Env): Value
}

class If(val cond: Expr, val then: Expr, val else_: Expr) : Expr() {
    override fun eval(env: Env): Value {
        return if (cond.eval(env).toBool())
            then.eval(env)   // return_call_ref via vtable
        else
            else_.eval(env)  // return_call_ref via vtable
    }
}

A chain of 10K if expressions overflows without tail calls but stays flat with them, because whichever branch If.eval picks is a tail call.

Interface dispatch

Interface method calls in tail position emit return_call_ref via the itable.

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

interface Processor {
    fun process(data: Data): Result
}

class Pipeline(val stages: List<Processor>) {
    fun run(data: Data, index: Int = 0): Result {
        if (index == stages.lastIndex) return stages[index].process(data)
        val intermediate = stages[index].process(data)
        return run(intermediate.toData(), index + 1)  // return_call (self, static)
    }
}

class ValidateProcessor(val next: Processor) : Processor {
    override fun process(data: Data): Result {
        validate(data)
        return next.process(data)  // return_call_ref via itable
    }
}

Function references

Calling a value of function type is an indirect call, even when it reads like a plain call.

1
2
3
4
5
6
7
8
9

fun doubled(x: Int) = x * 2

fun applyTransform(value: Int, transform: (Int) -> Int): Int {
    return transform(value)  // indirect: transform is a runtime value
}

fun main() {
    println(applyTransform(21, ::doubled))  // 42
}

For function references like ::doubled, the compiler generates a reference class whose invoke bridge dispatches through callRef, a typed call_ref on a funcref. PR #5 emits return_call_ref for that bridge in tail position.

Lambda dispatch

Lambda closures go through FunctionN<R>.invoke(). Because R is a type parameter, Kotlin erases the invoke method’s return type to Any? at the Wasm level. This creates two obstacles for tail calls.

  1. GenericReturnTypeLowering sees a mismatch between the erased return type Any? and the call-site type, wrapping the invoke in a cast that demotes it from tail position.

  2. Even without the cast, the Wasm-level signature mismatch between (ref null $kotlin.Any) and the caller’s actual return type would cause the validator to reject return_call_ref.

1
2
3
4

fun runThen(n: Int, action: () -> Unit) {
    if (n == 0) return action()      // in tail position, but NOT emitted as a tail call
    return runThen(n - 1, action)    // emitted (static self-dispatch)
}

runThen returns void at the Wasm level while invoke returns (ref null $kotlin.Any). Both obstacles apply here.

Lambda-to-lambda dispatch

1
2
3

fun cps(n: Int, k: (Int) -> Int): Int =
    if (n == 0) k(0)
    else cps(n - 1) { x: Int -> k(x + 1) }

When one lambda’s invoke calls another lambda’s invoke, both methods return Any?. The Wasm return types match, so obstacle 2 does not apply.

Obstacle 1 still produces an intermediate IR pattern. The cast generates an unbox of the inner invoke result from Any? to Int, and the autoboxing phase then re-boxes it back to Any? because the outer invoke also returns Any?. The round-trip is a no-op at the Wasm level, but it buries the inner invoke call inside wrapper IR and prevents tail call marking.

PR #17 adds a lowering pass that removes these box(unbox(e)) round-trips before tail call marking.

General case

When a non-lambda function calls a lambda in tail position, both obstacles remain. The caller’s concrete Wasm return type (i32, void, etc.) does not match invoke’s (ref null $kotlin.Any), so return_call_ref cannot be emitted. I haven’t come up with an idea to solve this.

Beyond tail position

The patterns above all require the call to be in tail position. Many recursive calls are not in tail positions.

1
2
3

return Cons(head, self(tail))       // constructor wrap
return 1 + self(n - 1)              // arithmetic
return process(self(subproblem))    // arbitrary computation

“Tail Recursion Modulo Context” (Leijen and Lorenzen’s, POPL 2023) provides a framework for these cases. The PR series implements three of these instantiations.

Constructor contexts

When a constructor wraps the recursive result, the compiler allocates it with a null placeholder, tail-calls the recursion and patches the result in afterward. This is destination-passing style, the same transform OCaml 4.14 shipped as [@tail_mod_cons].

1
2
3
4

fun replicate(n: Int, x: Int): IList<Int> = when {
    n <= 0 -> IList.Nil
    else -> IList.Cons(x, replicate(n - 1, x))  // context: Cons(x, ·)
}

See: Kotlin/Wasm: Tail Modulo Cons Lowering

Monoid contexts

When an associative operation wraps the recursive result, the compiler passes an accumulator parameter.

1
2
3
4
5
6
7
8
9
10
11

fun countUp(n: Int): Int {
    if (n == 0) return 0
    return 1 + countUp(n - 1)  // context: 1 + ·
}

// Generated by the accumulator lowering
private fun countUp$accum(n: Int, acc: Int): Int {
    if (n == 0) return 0 + acc
    return countUp$accum(n - 1, 1 + acc)  // self-call in tail position
}
fun countUp(n: Int): Int = countUp$accum(n, 0)

Currently, my PR handles commutative operators on Int and Long with +, *, and, or, and xor. It also handles String.plus, which is associative but not commutative, by preserving operand order for one-sided patterns like return str + self(args).

My PR does not cover the full general monoid case, where operands appear on both sides of the recursive call.

CPS (Continuation-Passing Style)

Continuation-Passing Style is the most general instantiation. It can transform any recursive pattern by representing the context as a continuation.

1
2
3
4
5
6
7
8
9
10
11

fun transform(n: Int): Int {
    if (n == 0) return 1
    return transform(n - 1) * 2 + 1  // context: · * 2 + 1
}

// Textbook CPS
fun transform_cps(n: Int, k: (Int) -> Int): Int {
    if (n == 0) return k(1)
    return transform_cps(n - 1) { x -> k(x * 2 + 1) }  // tail call
}
fun transform(n: Int): Int = transform_cps(n) { it }

The self-call sits in tail position. Rather than allocating a closure for each continuation, the compiler defunctionalizes the continuation into a typed heap frame and a trampoline loop that processes frames iteratively. This avoids closure allocation and lambda dispatch.

This transformation is being implemented in this PR (work in progress).

Each recursive call allocates one heap frame, which can make performance worse than native recursion. Constructor and accumulator transforms avoid this cost for their respective shapes, so the compiler prefers them over CPS when the pattern fits.

Kotlin’s DeepRecursiveFunction covers the same class of patterns at the library level through the coroutine machinery. The suspend/resume protocol per call can make it slower than native recursion on Kotlin/Wasm. Using it also requires rewriting the recursive function into the DeepRecursiveFunction { ... } form and replacing every recursive call with callRecursive.

1
2
3
4
5
6
7
8
9
10
11

// Before: plain recursion
fun sumList(node: ListNode?): Int {
    if (node == null) return 0
    return node.value + sumList(node.next)
}

// After: DeepRecursiveFunction rewrite
val sumList = DeepRecursiveFunction<ListNode?, Int> { node ->
    if (node == null) 0
    else node.value + callRecursive(node.next)
}

The CPS lowering automates this rewrite at the compiler level. The developer writes plain recursion and the compiler generates the heap-frame version.

Others

The paper also describes following two patterns, but they are not implemented yet. These are covered by CPS.

  • exponent contexts: the context is repeated application of the same function
  • semiring contexts: combine two monoid operators with a distributivity law, such as return x + 31 * hash(xs)

Real World Source Code Survey

I surveyed 18 Kotlin multiplatform libraries targeting wasmJs, selected from ~180 GitHub hits for wasmJs() filename:build.gradle.kts by domain likelihood of recursive logic. I classified every recursive call by which compiler transform could optimize it and verified each hit by hand.

Surveyed repositories

As a result, I found that

  • Genuine tail calls, construttor/accumulator patterns are almost not found in real Kotlin codebases as far as I searched.
  • Most recursion needs CPS
    • 13 of 18 repositories contain recursive functions. The dominant pattern is a recursive call inside a loop body or a higher-order function like forEach, map, or any.

How library authors avoid recursion

I found that some library authors have engineered the recursion away.

  • JetBrains/markdown parses block structure with no recursion at all. The parser core, MarkerProcessor, maintains an explicit stack of open blocks and pushes and pops it in a flat scan loop.

  • The kudzu parser combinator library makes every parser a DeepRecursiveFunction by declaration, moving recursion to the heap via the coroutine machinery.

  • Some libraries still crash on deep input.

    • Apollo’s GraphQL parser overflows around nesting depth 2,000
    • Stdlib regex matcher recurses per input character with an open crash report (KT-63689).

Libraries that rewrote recursion by hand paid a cost in development effort and, in kudzu’s case, runtime performance. Libraries that did not rewrite it crash on deep input. The compiler lowerings can automate these rewrites.

Technology, Compiler
Kotlin WebAssembly Compiler GSoC
This post is licensed under CC BY 4.0 by the author.
Recent Update
Contents