Automatic tail-call optimization (auto-TCO) in StaticOptimizer

[He-Pin]

Motivation

Jsonnet programs often use recursive functions that are naturally tail-recursive (the recursive call is the last expression in a branch). Without tail-call optimization (TCO), these hit StackOverflowError on deep recursion. Currently, sjsonnet supports TCO via explicit TailCall nodes, but users must manually structure their code to trigger it.

Key Design Decision

Implement automatic TCO detection in the StaticOptimizer. During the optimization pass, analyze function bodies to identify tail-position calls and automatically rewrite them to use TailCall nodes. This happens transparently — no source code changes needed.

Modification

sjsonnet/src/sjsonnet/StaticOptimizer.scala:

• Added tail-call analysis pass that walks function bodies
• Identifies self-recursive calls in tail position (last expression in if/else branches, local bindings)
• Rewrites qualifying calls from Apply* to the TailCall variant

sjsonnet/src/sjsonnet/Expr.scala:

• Added new Expr node types for automatic tail calls
• Support for TCO metadata in the AST

sjsonnet/src/sjsonnet/ExprTransform.scala:

• Updated expression transformer to handle new TCO node types

sjsonnet/src/sjsonnet/Evaluator.scala:

• Evaluation support for the auto-TCO nodes
• Trampoline handling for the new tail-call variants

sjsonnet/src/sjsonnet/Val.scala:

• Function value support for TCO metadata

Tests:

• new_test_suite/auto_tco.jsonnet — verifies deep recursion (>10000 depth) works with auto-TCO
• new_test_suite/auto_tco.jsonnet.golden — expected output

Benchmark Results

Expected Impact

• tail_call benchmark: direct improvement (deeper recursion without stack overflow)
• inheritance_function_recursion: may benefit from TCO on recursive inheritance patterns
• realistic2: depends on recursion depth in the workload

JMH — Pending Data

Analysis

• Correctness: Only rewrites calls that are provably in tail position. Non-tail calls are left unchanged.
• Scope: Handles if/else branches, local bindings, and direct self-recursion. Mutual recursion is out of scope.
• Safety: The trampoline mechanism (TailCall.resolve) is already battle-tested in the existing manual TCO path.
• Interaction with existing TCO: Complements the existing manual TailCall infrastructure. The static optimizer detects cases that users would otherwise need to manually optimize.

References

• Existing TailCall mechanism in sjsonnet/src/sjsonnet/Val.scala
• jrsonnet implements TCO at the evaluator level for self-recursive functions
• go-jsonnet supports TCO for tail-position calls

Result

Automatic tail-call optimization for self-recursive Jsonnet functions. Eliminates StackOverflowError on deep recursion without source changes. Draft PR pending benchmark data.