From e2b607f0f3cdc68bb47b2547c9859d17e958d495 Mon Sep 17 00:00:00 2001 From: Dmitri Sotnikov Date: Tue, 16 Jun 2026 21:07:06 +0000 Subject: [PATCH] backend: lower tail loop/recur to while + state vars, not a per-iter closure (jolt-v28u) (#151) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit emit-loop compiled every loop/recur to a self-recursive local closure called once per iteration — relying on Janet TCO for stack safety but paying a fn frame + arg bind each iteration. The jolt-5vsp spike localized the whole ~1.43x jolt-over-hand-Janet gap on compute loops to exactly this. Lower instead to a Janet `while` + state vars: the loop bindings become vars carried across iterations, a recur writes them and raises a continue flag, and a non-recur tail value falls out through a result var. recur-name routing in emit-recur picks the while-set lowering for loops and leaves the fn-arity self-call path untouched. The one subtlety is closure capture: Janet closures capture vars BY REFERENCE, so a closure built in the body over a shared mutable loop var would see the final value ([3 3 3]) instead of its iteration's ([0 1 2]). Each iteration rebinds the loop names into a fresh immutable `let` before running the body, which restores per-iteration capture. recur reads those immutable bindings and writes the state vars, so cross-referencing args (swap, fib) need no temps. mandelbrot 218 -> 164 ms (~11.2x JVM, from 15x). fib is unaffected — it's fn-arity recursion, not a loop. Regression spec in control-flow-spec covers closure capture, no-clobber recur, nested loops, sequential init, recur through let, and that fn-arity recur still works. Gate green (conformance x3, full suite). Note: validating this after a rebuild needs JOLT_NO_DEPS_CACHE=1 — the deps-image cache keys on the version string, not build identity, so it served stale codegen (filed separately). Co-authored-by: Yogthos --- src/jolt/backend.janet | 81 +++++++++++++++++++++++++------ test/spec/control-flow-spec.janet | 17 +++++++ 2 files changed, 82 insertions(+), 16 deletions(-) diff --git a/src/jolt/backend.janet b/src/jolt/backend.janet index 94358d2..cac8ee9 100644 --- a/src/jolt/backend.janet +++ b/src/jolt/backend.janet @@ -208,27 +208,76 @@ (array/push call ['if ['> ['length jargs] nfixed] ['tuple/slice jargs nfixed]])) (tuple/slice call)) +# A tail loop lowers to a Janet `while` + state vars (jolt-v28u), NOT a +# self-recursive closure called once per iteration — that paid a fn frame + arg +# bind every iteration, which the jolt-5vsp spike localized as the ENTIRE ~1.43x +# jolt-over-hand-Janet gap on compute loops. Shape: +# (let [b0 init0] (let [b1 init1] ... ; sequential inits (later sees earlier) +# (do (var $b0 b0) (var $b1 b1) ... ; state carried across iterations +# (var flag true) (var result nil) +# (while flag +# (set flag false) +# (let [b0 $b0 b1 $b1 ...] ; FRESH immutable per-iteration binding +# (set result ))) ; a tail recur sets the $vars + flag +# result))) +# The per-iteration `let` is load-bearing, not cosmetic: a closure built in the +# body must capture THIS iteration's value (Clojure semantics), and Janet closures +# capture vars BY REFERENCE — close over a shared mutable var and every closure +# sees the final value ([3 3 3]). Rebinding into an immutable `let` each iteration +# restores per-iteration capture ([0 1 2]). recur reads the immutable bi and writes +# $bi, so the read/write sets are disjoint and cross-referencing args (swap, fib) +# need no temps. (defn- emit-loop [ctx node] - (def L (symbol (node :recur-name))) - (def params @[]) - # Initial inits bind SEQUENTIALLY (a later init can reference an earlier binding, - # like let / Clojure's loop) — emit them in a Janet `let`, then enter the recur - # target L with those values, rather than computing all inits in the outer scope. - (def let-binds @[]) + (def recur-name (node :recur-name)) + (def syms @[]) # loop binding names (source symbols), read by the body + (def state @[]) # fresh vars carrying each binding across iterations + (def inits @[]) (each pair (vview (node :bindings)) (def p (vview pair)) - (def sym (symbol (in p 0))) - (array/push params sym) - (array/push let-binds sym) - (array/push let-binds (emit ctx (in p 1)))) - ['do - ['var L nil] - ['set L ['fn (tuple/slice params) (emit ctx (node :body))]] - ['let (tuple/slice let-binds) (tuple/slice (array/concat @[L] params))]]) + (array/push syms (symbol (in p 0))) + (array/push state (gsym)) + (array/push inits (emit ctx (in p 1)))) + (def flag (gsym)) + (def result (gsym)) + # Register this loop so its recurs lower to state-var sets (emit-recur looks it + # up by recur-name); restore any prior binding after the body (nil removes it). + (def frames (or (get (ctx :env) :loop-frames) + (let [t @{}] (put (ctx :env) :loop-frames t) t))) + (def prev (get frames recur-name)) + (put frames recur-name {:state state :flag flag}) + (def body (emit ctx (node :body))) + (put frames recur-name prev) + (def per-iter @[]) + (for i 0 (length syms) (array/push per-iter (syms i) (state i))) + (def block @['do]) + (for i 0 (length syms) (array/push block ['var (state i) (syms i)])) + (array/push block ['var flag true] ['var result nil]) + (array/push block ['while flag + ['set flag false] + ['let (tuple/slice per-iter) ['set result body]]]) + (array/push block result) + (var form (tuple/slice block)) + (for ri 0 (length syms) + (def i (- (length syms) 1 ri)) + (set form ['let [(syms i) (inits i)] form])) + form) (defn- emit-recur [ctx node] - (tuple/slice (array/concat @[(symbol (node :recur-name))] - (map |(emit ctx $) (vview (node :args)))))) + (def recur-name (node :recur-name)) + (def frame (get (get (ctx :env) :loop-frames @{}) recur-name)) + (if frame + # while-lowered loop: write each state var from its recur arg (args read the + # immutable iteration bindings — disjoint from the $vars, so no clobber), then + # raise the continue flag. The `do` yields the flag value, discarded in tail. + (let [args (map |(emit ctx $) (vview (node :args))) + block @['do]] + (for i 0 (length (frame :state)) + (array/push block ['set ((frame :state) i) (args i)])) + (array/push block ['set (frame :flag) true]) + (tuple/slice block)) + # fn-arity recur: a self-call the Janet runtime tail-calls (unchanged path). + (tuple/slice (array/concat @[(symbol recur-name)] + (map |(emit ctx $) (vview (node :args))))))) # Var-cell read of a clojure.core fn, as ((in 'cell :root) args...) — the same # live-root call the :var emit uses, for the ns get/set helpers below. diff --git a/test/spec/control-flow-spec.janet b/test/spec/control-flow-spec.janet index ed21b78..59e3ccf 100644 --- a/test/spec/control-flow-spec.janet +++ b/test/spec/control-flow-spec.janet @@ -60,6 +60,23 @@ ["when-some binds zero" "1" "(when-some [x 0] (inc x))"] ["if-let evals test once" "1" "(let [c (atom 0)] (if-let [v (do (swap! c inc) :v)] @c :none))"]) +# Regression: loop/recur lowering (jolt-v28u). The backend lowers tail loop/recur +# to a Janet while + state vars with a fresh per-iteration `let` rebinding of the +# loop names. The let rebinding is load-bearing: a closure created in the body +# must capture THAT iteration's value (Clojure semantics), not a shared mutable +# var — Janet closures capture vars by reference, so a naive while+var would give +# [3 3 3]. recur reads the (immutable) iteration bindings and writes the state +# vars, so cross-referencing args don't clobber (swap, fib). +(defspec "control / loop lowering (jolt-v28u)" + ["closure captures per-iter binding" "[0 1 2]" + "(mapv (fn [g] (g)) (loop [i 0 fs []] (if (< i 3) (recur (inc i) (conj fs (fn [] i))) fs)))"] + ["fib via loop" "55" "(loop [a 0 b 1 i 0] (if (= i 10) a (recur b (+ a b) (inc i))))"] + ["recur args no clobber" "[2 1]" "(loop [a 1 b 2 n 0] (if (= n 1) [a b] (recur b a (inc n))))"] + ["nested loops" "9" "(loop [i 0 s 0] (if (= i 3) s (recur (inc i) (loop [j 0 t s] (if (= j 3) t (recur (inc j) (inc t)))))))"] + ["loop sequential init" "12" "(loop [a 1 b (+ a 10)] (+ a b))"] + ["recur through let" "6" "(loop [i 0 acc 0] (let [x (* i 2)] (if (< i 3) (recur (inc i) (+ acc x)) acc)))"] + ["fn-arity recur intact" "15" "((fn f [n acc] (if (zero? n) acc (recur (dec n) (+ acc n)))) 5 0)"]) + (defspec "control / iteration" ["dotimes side-effect" "5" "(let [a (atom 0)] (dotimes [i 5] (swap! a inc)) @a)"] ["while" "5" "(let [a (atom 0)] (while (< @a 5) (swap! a inc)) @a)"]