Skip to content

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

Draft
He-Pin wants to merge 1 commit intodatabricks:masterfrom
He-Pin:perf/auto-tco
Draft

Automatic tail-call optimization (auto-TCO) in StaticOptimizer#694
He-Pin wants to merge 1 commit intodatabricks:masterfrom
He-Pin:perf/auto-tco

Conversation

@He-Pin
Copy link
Copy Markdown
Contributor

@He-Pin He-Pin commented Apr 6, 2026

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.

@He-Pin He-Pin force-pushed the perf/auto-tco branch 3 times, most recently from 7adb221 to da38f7c Compare April 9, 2026 03:02
@He-Pin He-Pin force-pushed the perf/auto-tco branch 2 times, most recently from 3532dbf to 2755bea Compare April 10, 2026 09:29
@He-Pin He-Pin closed this Apr 10, 2026
@He-Pin He-Pin reopened this Apr 10, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant