Hint-directed fast arithmetic: loop-carried variable typing (round 2)

A loop binding whose init is double and whose every recur arg stays double (a
bounded monotone fixpoint) is typed :double, so its arithmetic — and the recur
args feeding it — emit fl-ops. Chez can then keep the accumulator unboxed in a
float register across the loop.

Integer loop vars stay untyped: a bare integer init never seeds :long (same rule
as round 1), so a bignum-producing loop keeps arbitrary precision rather than
overflowing a fixnum. recur-kinds walks only tail position (if/do-ret/let-body),
stopping at nested loop/fn so a loop sees only its own recur.

A/B on a loop-carried double accumulator: 735ms generic -> 500ms typed (1.47x),
closing the gap to the JVM from ~3.3x to ~2.2x. The integer counter stays generic,
which is most of the residual.

jolt-core/jolt/passes/numeric.clj

@@ -52,6 +52,39 @@
 
 (declare an)
 
+;; The recur-arg kinds for the recurs targeting THIS loop level. recur only appears
+;; in tail position (an if branch, a do's ret, a let body), so descend only those;
+;; a nested loop/fn (and any non-tail child) owns its own recur and is skipped.
+(defn- recur-kinds [node tenv]
+  (let [op (get node :op)]
+    (cond
+      (= op :recur) [(mapv (fn [a] (nth (an a tenv) 0)) (get node :args))]
+      (= op :let) (recur-kinds (get node :body)
+                               (reduce (fn [te b] (assoc te (nth b 0) (nth (an (nth b 1) te) 0)))
+                                       tenv (get node :bindings)))
+      (= op :if) (concat (recur-kinds (get node :then) tenv) (recur-kinds (get node :else) tenv))
+      (= op :do) (recur-kinds (get node :ret) tenv)
+      :else [])))
+
+;; Loop-var kinds by bounded fixpoint. A var is :double only if its init is double
+;; AND every recur arg in that slot is double (under the current assumption) — a
+;; monotone demotion that stops at a fixpoint, bounded by the var count. Integers
+;; stay untyped (no :long from a bare init literal, so a bignum-producing loop keeps
+;; arbitrary precision). A :double loop var's init and recur args are all flonums,
+;; so no entry coercion is needed (unlike a fn param fed an arbitrary argument).
+(defn- loop-kinds [names ik body tenv]
+  (loop [cur (mapv (fn [k] (if (= k :double) :double nil)) ik) iter 0]
+    (if (> iter (count names))
+      cur
+      (let [te (reduce (fn [t i] (assoc t (nth names i) (nth cur i))) tenv (range (count names)))
+            rks (recur-kinds body te)
+            nxt (mapv (fn [j]
+                        (if (and (= (nth cur j) :double)
+                                 (every? (fn [rk] (= :double (nth rk j))) rks))
+                          :double nil))
+                      (range (count names)))]
+        (if (= nxt cur) cur (recur nxt (inc iter)))))))
+
 ;; Seed a fn arity's local env from its numeric param hints; an unhinted param
 ;; shadows any same-named outer local to nil.
 (defn- arity-env [tenv a]
@@ -109,11 +142,15 @@
             br (an (get node :body) (nth res 0))]
         [(nth br 0) (assoc node :bindings (nth res 1) :body (nth br 1))])
       (= op :loop)
-      ;; loop vars join across recur, untracked here, so they stay untyped; still
+      ;; inits evaluate in the OUTER env; loop vars get their fixpoint kinds for the body.
-      ;; descend to specialize any non-loop arithmetic in the inits/body.
+      (let [binds (get node :bindings)
-      [nil (assoc node
+            names (mapv (fn [b] (nth b 0)) binds)
-                  :bindings (mapv (fn [b] [(nth b 0) (nth (an (nth b 1) tenv) 1)]) (get node :bindings))
+            ik (mapv (fn [b] (nth (an (nth b 1) tenv) 0)) binds)
-                  :body (nth (an (get node :body) tenv) 1))]
+            lk (loop-kinds names ik (get node :body) tenv)
+            te (reduce (fn [t i] (assoc t (nth names i) (nth lk i))) tenv (range (count names)))]
+        [nil (assoc node
+                    :bindings (mapv (fn [b] [(nth b 0) (nth (an (nth b 1) tenv) 1)]) binds)
+                    :body (nth (an (get node :body) te) 1))])
       (= op :if)
       (let [tr (an (get node :test) tenv)
             thn (an (get node :then) tenv)

test/chez/numeric-test.ss

@@ -56,6 +56,15 @@
 (let ((e (emitf "u" "(fn* ([^double x] (let [d (* x x)] (+ d 1.0))))")))
   (ok "let-bound double propagates (fl* then fl+)" (and (has? e "(fl*") (has? e "(fl+"))))
 
+;; --- loop-carried variable typing (round 2) ---
+;; a double accumulator types via fixpoint, so its recur arithmetic is fl-ops.
+(let ((e (emitf "u" "(fn* ([] (loop [acc 0.0 i 0] (if (< i 5) (recur (+ acc 1.5) (inc i)) acc))))")))
+  (ok "loop double accumulator lowers (+ acc 1.5) to fl+" (has? e "(fl+")))
+;; an integer accumulator stays generic — a bignum-producing loop keeps arbitrary
+;; precision (no fx* overflow).
+(let ((e (emitf "u" "(fn* ([] (loop [acc 1 i 1] (if (< i 25) (recur (* acc i) (inc i)) acc))))")))
+  (ok "loop integer accumulator is NOT fx-specialized" (not (has? e "(fx*"))))
+
 ;; --- soundness: un-hinted / integer-literal code stays generic ---
 (let ((e (emitf "u" "(fn* ([a b] (+ a b)))")))
   (ok "un-hinted + stays generic (no fl/fx)" (and (not (has? e "(fl+")) (not (has? e "(fx+")))))
@@ -79,6 +88,10 @@
 (ok "long unary negate"        (= -5 (jnum->exact (ev "((fn* ([^long a] (- a))) 5)"))))
 (ok "long quot 7/2 = 3"        (= 3  (jnum->exact (ev "((fn* ([^long a ^long b] (quot a b))) 7 2)"))))
 (ok "double + int literal = 4.5" (= 9 (jnum->exact (ev "((fn* ([^double x] (* (+ x 1) 2))) 3.5)"))))
+(ok "loop double accumulator: 10*1.5 = 15"
+    (= 15 (jnum->exact (ev "((fn* ([] (loop [acc 0.0 i 0] (if (< i 10) (recur (+ acc 1.5) (inc i)) acc)))))"))))
+(ok "loop integer factorial stays exact (bignum preserved)"
+    (jolt-truthy? (ev "(< 1000000000000000000000 ((fn* ([] (loop [acc 1 i 1] (if (< i 25) (recur (* acc i) (inc i)) acc)))) ))")))
 
 (printf "~a/~a passed~n" (- total fails) total)
 (exit (if (zero? fails) 0 1))