Skip to content

WIP: Proper tco optimization#121

Draft
CertainLach wants to merge 3 commits intomasterfrom
feature/proper-tco
Draft

WIP: Proper tco optimization#121
CertainLach wants to merge 3 commits intomasterfrom
feature/proper-tco

Conversation

@CertainLach
Copy link
Copy Markdown
Owner

@CertainLach CertainLach commented Jun 18, 2023

C++ jsonnet does have proper TCO optimization.
However, I haven't used this approach since it is hard to maintain in some cases, especially when it is used in most of the code.
But now I got an idea: Why not implement a hybrid approach where TCO exists only for the main interpreter loop and not in the rest of the codebase, i.e. builtins?
In this PR I explore this possibility, it may also allow for async builtins/std.parallelMap/huge optimizations in recursive functions/better GC root management strategies.

Stack traces and some other things are broken for now, not ready for use.

@CertainLach
Copy link
Copy Markdown
Owner Author

Why not implement a hybrid approach where TCO exists only for the main interpreter loop and not in the rest of the codebase, i.e. builtins?

Builtins can be implemented too, I.e by utilizing generators/async-await

In cpp-jsonnet, there is a PR to implement manifestification inside of interpreter loop, and it looks like a cool idea
google/jsonnet#1142

@CertainLach CertainLach mentioned this pull request Apr 26, 2024
pub(crate) struct TcVM {
pub(crate) vals: Fifo<Val>,
pub(crate) ctxs: Fifo<Context>,
pub(crate) apply: Fifo<TailCallApply>,
Copy link
Copy Markdown
Owner Author

@CertainLach CertainLach Apr 27, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right here I have 3 stacks, where TailCallApply is an enum with items of varying length, that contains references to other stacks.
This enum value is pretty big, with no means to shrink without moving some things to allocation.
rsjsonnet did funny thing here, where it has 9 stacks instead (https://github.com/eduardosm/rsjsonnet/blob/main/rsjsonnet-lang/src/program/eval/mod.rs#L32-L40), and keeps allocations even, because all the variable-size data is in the other stacks. This is pretty verbose code, with some implicit relations...

But I think it is possible to extend this approach using proc macro:

#[derive(FifoData)]
enum Op {
  A { a: u32, b: u32 },
  B { v: String },
  C { a: Context, b: Val, c: u32, d: String }
  // There should be a way to opt-out from moving data to stacks,
  // the remaining data might be attached to... OpTag enum?
}

Which would expand into

enum OpTag {
  A,
  B,
  C,
}

struct OpFifo {
  ops: Vec<OpTag>,
  u32s: Vec<u32>,
  strings: Vec<String>,
  // Where Val can also be inlined to some extent together with its children, so this proc macro
  // should allow composition of its elements
  vals: Vec<Val>,
  contexts: Vec<Context>,
}

impl Op {
  fn pop(fifo: &mut OpFifo) -> Self {
    match fifo.ops.pop() {
      OpTag::A => Self::A { b: fifo.u32s.pop(), a: fifo.u32s.pop() },
      OpTag::B => Self::B { v: fifo.strings.pop() },
      OpTag::C => Self::C { d: fifo.strings.pop(), c: fifo.u32s.pop(), b: fifo.vals.pop() },
    }
  }
  fn push(self, fifo: &mut OpFifo) {
    match self {
      Self::A {a, b} => { fifo.ops.push(OpTag::A); fifo.u32s.push(a); fifo.u32s.push(b); },
      _ => todo!(),
    }
  }
}

This structure is kinda dumb for the provided enum, but might make sense for the whole jrsonnet?

stephenamar-db pushed a commit to databricks/sjsonnet that referenced this pull request Mar 3, 2026
Motivation:
Add partial TCO.

refs:CertainLach/jrsonnet#121

```jsonnet
local sum(n, accum=0) =
  if n <= 0 then accum
  else sum(n - 1, accum + n) tailstrict;
 
local sz = 10000;
std.assertEqual(sum(sz), sz * (sz + 1) / 2)
```
<img width="967" height="267" alt="image"
src="https://github.com/user-attachments/assets/7879e8d4-a603-48fd-89c3-d7d763b3d6b3"
/>


can work now and evaluate to `true` but was:

<img width="1648" height="419" alt="image"
src="https://github.com/user-attachments/assets/b2c8b9eb-7522-48fa-86a7-9d24b85b125f"
/>

on main

```
125] Benchmark                                                                  (path)  Mode  Cnt   Score   Error  Units
125] RegressionBenchmark.main             bench/resources/bug_suite/assertions.jsonnet  avgt        0.434          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.01.jsonnet  avgt        0.076          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.02.jsonnet  avgt       41.836          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.03.jsonnet  avgt       13.310          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.04.jsonnet  avgt       32.440          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.06.jsonnet  avgt        0.453          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.07.jsonnet  avgt        2.910          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.08.jsonnet  avgt        0.060          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.09.jsonnet  avgt        0.067          ms/op
125] RegressionBenchmark.main         bench/resources/cpp_suite/gen_big_object.jsonnet  avgt        0.998          ms/op
125] RegressionBenchmark.main      bench/resources/cpp_suite/large_string_join.jsonnet  avgt        2.000          ms/op
125] RegressionBenchmark.main  bench/resources/cpp_suite/large_string_template.jsonnet  avgt        2.384          ms/op
125] RegressionBenchmark.main             bench/resources/cpp_suite/realistic1.jsonnet  avgt        3.287          ms/op
125] RegressionBenchmark.main             bench/resources/cpp_suite/realistic2.jsonnet  avgt       75.854          ms/op
125] RegressionBenchmark.main                  bench/resources/go_suite/base64.jsonnet  avgt        0.846          ms/op
125] RegressionBenchmark.main            bench/resources/go_suite/base64Decode.jsonnet  avgt        0.650          ms/op
125] RegressionBenchmark.main       bench/resources/go_suite/base64DecodeBytes.jsonnet  avgt        9.567          ms/op
125] RegressionBenchmark.main       bench/resources/go_suite/base64_byte_array.jsonnet  avgt        1.495          ms/op
125] RegressionBenchmark.main              bench/resources/go_suite/comparison.jsonnet  avgt       23.110          ms/op
125] RegressionBenchmark.main             bench/resources/go_suite/comparison2.jsonnet  avgt       80.098          ms/op
125] RegressionBenchmark.main        bench/resources/go_suite/escapeStringJson.jsonnet  avgt        0.051          ms/op
125] RegressionBenchmark.main                   bench/resources/go_suite/foldl.jsonnet  avgt        9.389          ms/op
125] RegressionBenchmark.main             bench/resources/go_suite/lstripChars.jsonnet  avgt        0.652          ms/op
125] RegressionBenchmark.main          bench/resources/go_suite/manifestJsonEx.jsonnet  avgt        0.073          ms/op
125] RegressionBenchmark.main          bench/resources/go_suite/manifestTomlEx.jsonnet  avgt        0.090          ms/op
125] RegressionBenchmark.main         bench/resources/go_suite/manifestYamlDoc.jsonnet  avgt        0.076          ms/op
125] RegressionBenchmark.main                  bench/resources/go_suite/member.jsonnet  avgt        0.744          ms/op
125] RegressionBenchmark.main                bench/resources/go_suite/parseInt.jsonnet  avgt        0.053          ms/op
125] RegressionBenchmark.main                 bench/resources/go_suite/reverse.jsonnet  avgt       11.664          ms/op
125] RegressionBenchmark.main             bench/resources/go_suite/rstripChars.jsonnet  avgt        0.655          ms/op
125] RegressionBenchmark.main              bench/resources/go_suite/stripChars.jsonnet  avgt        0.635          ms/op
125] RegressionBenchmark.main                  bench/resources/go_suite/substr.jsonnet  avgt        0.164          ms/op
125/125, SUCCESS] ./mill bench.runRegressions 399s

```


current:

```
125] Benchmark                                                                  (path)  Mode  Cnt   Score   Error  Units
125] RegressionBenchmark.main             bench/resources/bug_suite/assertions.jsonnet  avgt        0.324          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.01.jsonnet  avgt        0.072          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.02.jsonnet  avgt       41.850          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.03.jsonnet  avgt       12.604          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.04.jsonnet  avgt       32.272          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.06.jsonnet  avgt        0.449          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.07.jsonnet  avgt        3.592          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.08.jsonnet  avgt        0.059          ms/op
125] RegressionBenchmark.main               bench/resources/cpp_suite/bench.09.jsonnet  avgt        0.067          ms/op
125] RegressionBenchmark.main         bench/resources/cpp_suite/gen_big_object.jsonnet  avgt        1.012          ms/op
125] RegressionBenchmark.main      bench/resources/cpp_suite/large_string_join.jsonnet  avgt        2.212          ms/op
125] RegressionBenchmark.main  bench/resources/cpp_suite/large_string_template.jsonnet  avgt        2.375          ms/op
125] RegressionBenchmark.main             bench/resources/cpp_suite/realistic1.jsonnet  avgt        3.277          ms/op
125] RegressionBenchmark.main             bench/resources/cpp_suite/realistic2.jsonnet  avgt       73.374          ms/op
125] RegressionBenchmark.main                  bench/resources/go_suite/base64.jsonnet  avgt        0.834          ms/op
125] RegressionBenchmark.main            bench/resources/go_suite/base64Decode.jsonnet  avgt        0.639          ms/op
125] RegressionBenchmark.main       bench/resources/go_suite/base64DecodeBytes.jsonnet  avgt        9.436          ms/op
125] RegressionBenchmark.main       bench/resources/go_suite/base64_byte_array.jsonnet  avgt        1.502          ms/op
125] RegressionBenchmark.main              bench/resources/go_suite/comparison.jsonnet  avgt       22.254          ms/op
125] RegressionBenchmark.main             bench/resources/go_suite/comparison2.jsonnet  avgt       77.183          ms/op
125] RegressionBenchmark.main        bench/resources/go_suite/escapeStringJson.jsonnet  avgt        0.051          ms/op
125] RegressionBenchmark.main                   bench/resources/go_suite/foldl.jsonnet  avgt        9.352          ms/op
125] RegressionBenchmark.main             bench/resources/go_suite/lstripChars.jsonnet  avgt        0.647          ms/op
125] RegressionBenchmark.main          bench/resources/go_suite/manifestJsonEx.jsonnet  avgt        0.073          ms/op
125] RegressionBenchmark.main          bench/resources/go_suite/manifestTomlEx.jsonnet  avgt        0.088          ms/op
125] RegressionBenchmark.main         bench/resources/go_suite/manifestYamlDoc.jsonnet  avgt        0.076          ms/op
125] RegressionBenchmark.main                  bench/resources/go_suite/member.jsonnet  avgt        0.752          ms/op
125] RegressionBenchmark.main                bench/resources/go_suite/parseInt.jsonnet  avgt        0.053          ms/op
125] RegressionBenchmark.main                 bench/resources/go_suite/reverse.jsonnet  avgt       11.094          ms/op
125] RegressionBenchmark.main             bench/resources/go_suite/rstripChars.jsonnet  avgt        0.655          ms/op
125] RegressionBenchmark.main              bench/resources/go_suite/stripChars.jsonnet  avgt        0.637          ms/op
125] RegressionBenchmark.main                  bench/resources/go_suite/substr.jsonnet  avgt        0.167          ms/op
125/125, SUCCESS] ./mill bench.runRegressions 399s
```

Another case is :
```
local sum = function(n, acc=0)
    if n == 0 then
      acc
    else
      sum(n - 1, acc + n) tailstrict; 

sum(100000) 

```
@CertainLach CertainLach marked this pull request as draft April 6, 2026 03:47
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