simon/jolt
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

backend: lower tail loop/recur to while + state vars, not a per-iter closure (jolt-v28u) (#151)

Browse source 
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 <yogthos@gmail.com>
This commit is contained in:
Dmitri Sotnikov 2026-06-16 21:07:06 +00:00 • committed by GitHub
parent 6cace3db90
commit e2b607f0f3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 82 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

81
src/jolt/backend.janet
View file

@ -208,27 +208,76 @@
(array/push call ['if ['> ['length jargs] nfixed] ['tuple/slice jargs nfixed]])) (array/push call ['if ['> ['length jargs] nfixed] ['tuple/slice jargs nfixed]]))
(tuple/slice call)) (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 <body>))) ; 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] (defn- emit-loop [ctx node]
(def L (symbol (node :recur-name))) (def recur-name (node :recur-name))
(def params @[]) (def syms @[]) # loop binding names (source symbols), read by the body
# Initial inits bind SEQUENTIALLY (a later init can reference an earlier binding, (def state @[]) # fresh vars carrying each binding across iterations
# like let / Clojure's loop) — emit them in a Janet `let`, then enter the recur (def inits @[])
# target L with those values, rather than computing all inits in the outer scope.
(def let-binds @[])
(each pair (vview (node :bindings)) (each pair (vview (node :bindings))
(def p (vview pair)) (def p (vview pair))
(def sym (symbol (in p 0))) (array/push syms (symbol (in p 0)))
(array/push params sym) (array/push state (gsym))
(array/push let-binds sym) (array/push inits (emit ctx (in p 1))))
(array/push let-binds (emit ctx (in p 1)))) (def flag (gsym))
['do (def result (gsym))
['var L nil] # Register this loop so its recurs lower to state-var sets (emit-recur looks it
['set L ['fn (tuple/slice params) (emit ctx (node :body))]] # up by recur-name); restore any prior binding after the body (nil removes it).
['let (tuple/slice let-binds) (tuple/slice (array/concat @[L] params))]]) (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] (defn- emit-recur [ctx node]
(tuple/slice (array/concat @[(symbol (node :recur-name))] (def recur-name (node :recur-name))
(map |(emit ctx $) (vview (node :args)))))) (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 # 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. # live-root call the :var emit uses, for the ns get/set helpers below.

17
test/spec/control-flow-spec.janet
View file

@ -60,6 +60,23 @@
["when-some binds zero" "1" "(when-some [x 0] (inc x))"] ["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))"]) ["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" (defspec "control / iteration"
["dotimes side-effect" "5" "(let [a (atom 0)] (dotimes [i 5] (swap! a inc)) @a)"] ["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)"] ["while" "5" "(let [a (atom 0)] (while (< @a 5) (swap! a inc)) @a)"]