feat: Phase 3 collection-element types + HOF awareness + ordered re-emit (jolt-d6u)

Extends the inference lattice with a parametric vector type {:vec ELEM} and
threads element types through the program:
- vector literals, conj/into, and range produce element-typed vectors;
- reduce/map/mapv/filter/filterv seed their closure's element (and reduce's
  accumulator) param, so a lookup inside the closure over a vector-of-structs
  specializes (the HOF-element-awareness piece);
- a var reference carries a VALUE type — a fn var is :truthy (non-nil, sealed
  root), a def var carries its inferred init type (e.g. a color table is
  {:vec :struct-map}); element-returning fns (rand-nth/first/nth/...) yield the
  collection's element type. These let the dynamically-built scene's sphere
  maps type as structs.
The inter-procedural fixpoint now also infers non-fn def value types, and the
recompile re-emits the WHOLE unit callee-first (reverse-topological) so a
caller re-embeds its recompiled, now-specialized callees and a call site
compiled after the pass links the whole chain.

Result on the ray tracer (no hints): the chain closes — hittables infers to
{:vec :struct-map}, hit-sphere's hittable param to :struct-map — and the render
goes 13.1s -> 12.8s. That is only ~3%, far short of the explicit hint's 1.22x.
The remaining gap is nested field access: a lookup RESULT like (:direction ray)
is :any, so (:r (:direction ray)) stays guarded, and the vec3 fns (called with
such values) can't be typed struct. The hint asserts the vec3 params directly
and propagates through inlining; matching it needs field-shape types
(ray.direction : vec3, vec3.r : number) — a structural extension (Phase 4).

Sound: a seeded full render produces an identical checksum (1915337);
conformance 335/335 x3 and the full jpm test pass; type-infer-phase3-test pins
the element-typing + HOF mechanism. Phase 2 (vector nth/count specialization)
was deprioritized — it is orthogonal to this benchmark.
This commit is contained in:
Yogthos 2026-06-13 06:53:21 -04:00
parent ea1d9a23e1
commit 09e5af02c9
3 changed files with 249 additions and 50 deletions

View file

@ -716,15 +716,25 @@
;; dynamic guard in place. Sound by construction: a concrete type is assigned
;; only when proven, so a wrong bare get is impossible.
;;
;; Lattice values: :struct-map (raw-get-safe), :phm-map, :vector, :set, :truthy
;; (a provably non-nil/non-false scalar — numbers, strings, keywords), :any (top).
(defn- join [a b] (if (= a b) a :any))
;; Lattice values: :struct-map (raw-get-safe), :phm-map, :set, :truthy (a
;; provably non-nil/non-false scalar), :any (top), and a PARAMETRIC vector type
;; {:vec ELEM} (jolt-d6u, Phase 3) carrying its element type so a reduce/map
;; closure over it can type its element param. {:vec ELEM} is a small struct, so
;; it compares by value on both the Clojure and the Janet (orchestrator) side.
(defn- velem [t] (get t :vec))
(defn- vec-type? [t] (some? (velem t)))
(defn- mk-vec [t] {:vec (if t t :any)})
(defn- join [a b]
(cond
(= a b) a
(and (vec-type? a) (vec-type? b)) (mk-vec (join (velem a) (velem b)))
:else :any))
(defn- struct-safe? [t] (= t :struct-map))
;; a value whose type guarantees it is neither nil nor false — the back end only
;; builds a struct (vs a phm) when every value is truthy, so a map literal is a
;; struct only when all its values have a truthy type.
;; struct only when all its values have a truthy type. Collections are non-nil.
(defn- truthy-type? [t]
(or (= t :truthy) (= t :struct-map) (= t :phm-map) (= t :vector) (= t :set)))
(or (= t :truthy) (= t :struct-map) (= t :phm-map) (= t :set) (vec-type? t)))
(def ^:private truthy-ret-fns
#{"+" "-" "*" "/" "inc" "dec" "mod" "rem" "quot" "min" "max" "abs"
@ -739,6 +749,14 @@
(def ^:private rtenv-box (atom {})) ;; "ns/name" -> inferred return type
(def ^:private calls-box (atom [])) ;; collected [ "ns/name" [arg-types...] ]
(def ^:private escapes-box (atom #{})) ;; var-keys used as a VALUE (not a call head)
;; jolt-d6u: a var reference's VALUE type — a fn var is :truthy (non-nil), a def
;; var carries its inferred init type (e.g. a color table -> {:vec :struct-map}).
;; The orchestrator populates this from sealed (opt-mode) cell roots + def inits.
(def ^:private vtype-box (atom {})) ;; "ns/name" -> value type
;; fns that RETURN an element of their (first) collection arg, so a lookup on the
;; result of (rand-nth coll-of-structs) etc. types as the element.
(def ^:private elem-fns #{"rand-nth" "first" "peek" "last" "nth" "fnext" "second"})
(defn- var-key [fnode] (str (get fnode :ns) "/" (get fnode :name)))
@ -750,14 +768,44 @@
(if r r (let [nm (and (= "clojure.core" (get fnode :ns)) (get fnode :name))]
(cond (nil? nm) :any
(contains? truthy-ret-fns nm) :truthy
(contains? vector-ret-fns nm) :vector
(contains? vector-ret-fns nm) (mk-vec :any)
:else :any))))
(= op :host) (let [nm (get fnode :name)]
(cond (contains? truthy-ret-fns nm) :truthy
(contains? vector-ret-fns nm) :vector
(contains? vector-ret-fns nm) (mk-vec :any)
:else :any))
:else :any)))
(declare infer)
;; HOFs that apply their fn arg to the ELEMENTS of a collection (jolt-d6u,
;; Phase 3). :epos is which param of the fn receives an element. reduce is
;; handled separately (its arity changes the coll position, and its closure
;; also takes an accumulator).
(def ^:private hof-table
{"map" {:epos 0} "mapv" {:epos 0} "filter" {:epos 0} "filterv" {:epos 0}
"keep" {:epos 0} "remove" {:epos 0} "run!" {:epos 0} "mapcat" {:epos 0}})
(defn- infer-fn-seeded
"Infer a fn-literal passed to a HOF, seeding the given params to element/accum
types (seeds: param-index -> type), other params :any, captured locals from
tenv. Returns [ret-type node'] ret is the lub of arity tail types, used to
type the HOF result (e.g. reduce's accumulator, mapv's element)."
[node seeds tenv]
(let [res (mapv (fn [a]
(let [params (get a :params)
pe (reduce (fn [e i]
(assoc e (nth params i)
(let [s (get seeds i)] (if s s :any))))
tenv (range (count params)))
pe (if (get a :rest) (assoc pe (get a :rest) :any) pe)
br (infer (get a :body) pe)]
[(nth br 0) (assoc a :body (nth br 1))]))
(get node :arities))
rets (mapv (fn [r] (nth r 0)) res)
ret (if (empty? rets) :any (reduce join (first rets) (rest rets)))]
[ret (assoc node :arities (mapv (fn [r] (nth r 1)) res))]))
(defn- infer
"Returns [type node'] the inferred type of node and node with struct-safe
:local references annotated :hint :struct. tenv maps in-scope local names to
@ -782,7 +830,10 @@
:struct-map :any)]
[t (assoc node :pairs (mapv (fn [r] [(nth r 0) (nth r 1)]) res))])
(= op :vector)
[:vector (assoc node :items (mapv (fn [x] (nth (infer x tenv) 1)) (get node :items)))]
(let [irs (mapv (fn [x] (infer x tenv)) (get node :items))
ets (mapv (fn [r] (nth r 0)) irs)
el (if (empty? ets) :any (reduce join (first ets) (rest ets)))]
[(mk-vec el) (assoc node :items (mapv (fn [r] (nth r 1)) irs))])
(= op :set)
[:set (assoc node :items (mapv (fn [x] (nth (infer x tenv) 1)) (get node :items)))]
(= op :if)
@ -799,17 +850,68 @@
[:any (assoc node :expr (nth (infer (get node :expr) tenv) 1))]
;; a :var reached HERE is in value position (an arg, a let init, ...), not
;; a call head — so the fn it names escapes and its params can't be inferred.
(= op :var) (do (swap! escapes-box conj (var-key node)) [:any node])
;; Its VALUE type comes from vtype-box (a fn is :truthy, a def carries its
;; inferred type); unknown -> :any.
(= op :var) (do (swap! escapes-box conj (var-key node))
[(let [vt (get @vtype-box (var-key node))] (if vt vt :any)) node])
(= op :invoke)
(let [fnode (get node :fn)
iscall-var (= :var (get fnode :op))
;; a :var in call-head position is a call, NOT an escape — so don't
;; route it through infer (which would record it as escaped).
fnode' (if iscall-var fnode (nth (infer fnode tenv) 1))
ares (mapv (fn [a] (infer a tenv)) (get node :args))]
(when iscall-var
(swap! calls-box conj [(var-key fnode) (mapv (fn [r] (nth r 0)) ares)]))
[(call-ret-type fnode) (assoc node :fn fnode' :args (mapv (fn [r] (nth r 1)) ares))])
cn (when (and iscall-var (= "clojure.core" (get fnode :ns))) (get fnode :name))
args (get node :args)
n (count args)]
(cond
;; reduce over a typed vector with a fn-literal (jolt-d6u): seed the
;; closure's accumulator (param 0) to the init type and its element
;; (param 1) to the vector's element type, so its body — and any calls
;; it makes — see those types.
(and (= cn "reduce") (>= n 2) (= :fn (get (nth args 0) :op)))
(let [three (>= n 3)
coll-r (infer (nth args (if three 2 1)) tenv)
init-r (when three (infer (nth args 1) tenv))
et (let [ct (nth coll-r 0)] (if (vec-type? ct) (velem ct) :any))
init-t (if init-r (nth init-r 0) :any)
fn-r (infer-fn-seeded (nth args 0) {0 init-t 1 et} tenv)]
[(join init-t (nth fn-r 0))
(assoc node :args (if three
[(nth fn-r 1) (nth init-r 1) (nth coll-r 1)]
[(nth fn-r 1) (nth coll-r 1)]))])
;; map/mapv/filter/... over a typed vector with a fn-literal: seed the
;; fn's element param; mapv/filterv produce a typed vector.
(and cn (get hof-table cn) (>= n 2) (= :fn (get (nth args 0) :op)))
(let [coll-r (infer (nth args 1) tenv)
et (let [ct (nth coll-r 0)] (if (vec-type? ct) (velem ct) :any))
fn-r (infer-fn-seeded (nth args 0) {(get (get hof-table cn) :epos) et} tenv)
rt (cond (= cn "mapv") (mk-vec (nth fn-r 0))
(= cn "filterv") (mk-vec et)
:else :any)]
[rt (assoc node :args [(nth fn-r 1) (nth coll-r 1)])])
;; conj/into: track the element type of a vector being grown.
(and (or (= cn "conj") (= cn "into")) (>= n 1))
(let [ares (mapv (fn [a] (infer a tenv)) args)
base (nth (nth ares 0) 0)
rest-ts (mapv (fn [r] (nth r 0)) (rest ares))
rt (cond
(and (= cn "conj") (vec-type? base))
(mk-vec (reduce join (velem base) rest-ts))
(and (= cn "into") (vec-type? base) (= 2 n) (vec-type? (nth rest-ts 0)))
(mk-vec (join (velem base) (velem (nth rest-ts 0))))
:else (call-ret-type fnode))]
[rt (assoc node :args (mapv (fn [r] (nth r 1)) ares))])
;; everything else: type args, collect the call (var callee), use the
;; declared/estimated return type. range produces a numeric vector.
:else
(let [fnode' (if iscall-var fnode (nth (infer fnode tenv) 1))
ares (mapv (fn [a] (infer a tenv)) args)]
(when iscall-var
(swap! calls-box conj [(var-key fnode) (mapv (fn [r] (nth r 0)) ares)]))
[(cond
(= cn "range") (mk-vec :truthy)
;; element-returning fn over a typed vector -> the element type
(and cn (contains? elem-fns cn) (> n 0))
(let [a0 (nth (nth ares 0) 0)] (if (vec-type? a0) (velem a0) :any))
:else (call-ret-type fnode))
(assoc node :fn fnode' :args (mapv (fn [r] (nth r 1)) ares))])))
(= op :let)
(let [res (reduce (fn [acc b]
(let [te (nth acc 0) binds (nth acc 1)
@ -854,6 +956,11 @@
type call results during the fixpoint."
[m] (reset! rtenv-box m))
(defn set-vtypes!
"Install var VALUE types (a map \"ns/name\" -> type): fn vars are :truthy
(non-nil), def vars carry their inferred init type (jolt-d6u)."
[m] (reset! vtype-box m))
(defn reset-escapes! [] (reset! escapes-box #{}))
(defn collected-escapes [] (vec @escapes-box))

View file

@ -72,16 +72,21 @@
(when (get (ctx :env) :inline?)
(def init (norm-node (node :init)))
(def meta (node :meta))
(when (and (= :fn (init :op))
(not (and meta (or (get meta :redef) (get meta :dynamic)))))
(def arities (vview (init :arities)))
(when (= 1 (length arities))
(def ar (norm-node (in arities 0)))
(unless (ar :rest)
(put cell :inline-ir {:params (ar :params) :body (ar :body)})
# jolt-767: stash the whole (post-pass) :def IR so the inter-procedural
# pass can re-infer its body with discovered param types and re-emit it.
(put cell :infer-ir node))))))
(def redefable (and meta (or (get meta :redef) (get meta :dynamic))))
(cond
redefable nil
(= :fn (init :op))
(let [arities (vview (init :arities))]
(when (= 1 (length arities))
(def ar (norm-node (in arities 0)))
(unless (ar :rest)
(put cell :inline-ir {:params (ar :params) :body (ar :body)})
# jolt-767: stash the whole (post-pass) :def IR so the inter-procedural
# pass can re-infer its body with discovered param types and re-emit it.
(put cell :infer-ir node))))
# a non-fn def: stash so the pass can infer its VALUE type (jolt-d6u), e.g.
# a color table used via rand-nth — its element type flows to lookups.
true (put cell :infer-ir node))))
# Var late-binding: reads go through `(var-get cell)` with the cell embedded as a
# constant, so compiled code sees redefinition (Janet early-binds plain symbols)
@ -759,12 +764,21 @@
# Recompiled bodies are semantically identical to the guarded ones, so this is
# correct regardless of recompile order; order only affects how far a direct-
# linked call propagates the faster callee.
(defn- itype-join [a b] (cond (nil? a) b (nil? b) a (= a b) a :any))
(defn- itype-join [a b]
(cond
(nil? a) b
(nil? b) a
(= a b) a
# compound vector types {:vec ELEM} join element-wise (jolt-d6u)
(and (struct? a) (struct? b) (in a :vec) (in b :vec))
(struct :vec (itype-join (in a :vec) (in b :vec)))
:any))
(defn infer-unit!
[ctx ns-name]
(def pns (ctx-find-ns ctx "jolt.passes"))
(def f-set-rtenv (and pns (ns-find pns "set-rtenv!")))
(def f-set-vtypes (and pns (ns-find pns "set-vtypes!")))
(def f-infer-body (and pns (ns-find pns "infer-body")))
(def f-reinfer (and pns (ns-find pns "reinfer-def")))
(def f-reset-esc (and pns (ns-find pns "reset-escapes!")))
@ -772,28 +786,34 @@
(def ns (ctx-find-ns ctx ns-name))
(def report @{})
(when (and ns f-set-rtenv f-infer-body f-reinfer f-reset-esc f-get-esc)
# gather single-fixed-arity fns with a stashed :def
# gather single-fixed-arity fns AND non-fn defs that stashed a :def IR
(def fns @[])
(def defs @[])
(def by-key @{})
(def vtypes @{}) # var VALUE types: fns -> :truthy (non-nil), defs -> inferred
(each nm (keys (ns :mappings))
(def v (get (ns :mappings) nm))
(when (and (table? v) (get v :infer-ir))
(def d (norm-node (get v :infer-ir)))
(def init (norm-node (d :init)))
(when (= :fn (init :op))
(def ars (vview (init :arities)))
(when (= 1 (length ars))
(def ar (norm-node (in ars 0)))
(unless (ar :rest)
(def pv (vview (ar :params)))
(def rec @{:key (string ns-name "/" nm) :cell v :def d
:params (ar :params) :body (ar :body)
:np (length pv) :pt (array/new-filled (length pv)) :ret nil})
(array/push fns rec)
(put by-key (rec :key) rec))))))
(when (> (length fns) 0)
(def key (string ns-name "/" nm))
(if (= :fn (init :op))
(let [ars (vview (init :arities))]
(when (= 1 (length ars))
(def ar (norm-node (in ars 0)))
(unless (ar :rest)
(def pv (vview (ar :params)))
(def rec @{:key key :cell v :def d :params (ar :params) :body (ar :body)
:np (length pv) :pt (array/new-filled (length pv)) :ret nil})
(array/push fns rec)
(put by-key key rec)
# a fn value is non-nil -> :truthy (sealed root in opt mode)
(put vtypes key :truthy))))
# non-fn def: its value type is inferred from its init (jolt-d6u)
(array/push defs @{:key key :init (d :init) :vt nil}))))
(when (or (> (length fns) 0) (> (length defs) 0))
((var-get f-reset-esc))
# --- param/return-type fixpoint (chaotic iteration to the LEAST fixpoint) ---
# --- param/return/value-type fixpoint (chaotic iteration to LEAST fixpoint) ---
# Param types are RECOMPUTED FRESH each iteration, not accumulated: :any is
# the lattice top, so a join with an early-iteration :any (a caller whose own
# params weren't typed yet) would poison the result permanently. Recomputing
@ -802,7 +822,8 @@
(var changed true) (var iter 0)
(while (and changed (< iter 16))
((var-get f-set-rtenv) prev-rt)
# type every body once under current param types; stash ret + calls
((var-get f-set-vtypes) vtypes)
# type every fn body once under current param types; stash ret + calls
(each f fns
(def tenv @{})
(def pv (vview (f :params)))
@ -810,6 +831,9 @@
(def res (vview ((var-get f-infer-body) (f :body) tenv)))
(put f :tret (in res 0))
(put f :tcalls (in res 2)))
# infer each def's VALUE type from its init
(each dv defs
(put dv :tvt (in (vview ((var-get f-infer-body) (dv :init) @{})) 0)))
# recompute param types FRESH (start at bottom = nil) from this round's calls
(def newpt @{})
(each f fns (put newpt (f :key) (array/new-filled (f :np))))
@ -832,24 +856,51 @@
(put f :pt np)
(put f :ret (f :tret))
(when (f :tret) (put nrt (f :key) (f :tret))))
(each dv defs
(when (not= (dv :tvt) (dv :vt)) (set changed true))
(put dv :vt (dv :tvt))
(when (dv :tvt) (put vtypes (dv :key) (dv :tvt))))
(set prev-rt nrt)
(++ iter))
# --- escaped fns: var used as a value -> params untrustworthy -> skip ---
(def esc @{})
(each k (vview ((var-get f-get-esc))) (put esc k true))
# --- re-infer + re-emit each fn with concrete param types seeded ---
(each f fns
# install the FINAL return + value types so reinfer-def sees them
(def final-rt @{})
(each f fns (when (f :ret) (put final-rt (f :key) (f :ret))))
((var-get f-set-rtenv) final-rt)
((var-get f-set-vtypes) vtypes)
# --- re-emit the WHOLE unit, callees first (jolt-d6u) -------------------
# Re-inference alone only rebinds a fn's own var, but the hot path runs
# through callee bodies INLINED / direct-linked into callers at first
# compile. Re-emitting in callee-first (reverse-topological) order makes
# each caller re-embed its now-recompiled callees, and re-infers its body
# (typing locals via return inference) — so the specialization propagates,
# and a call site compiled AFTER this pass (the -e entry) links the whole
# recompiled chain. Every fn is re-emitted, not just those with concrete
# params, so the embedding refreshes even where a fn gained no param type.
(def order @[])
(def seen @{})
(defn visit [k]
(unless (get seen k)
(put seen k true)
(def f (get by-key k))
(when f
(each c (vview (f :tcalls)) (visit (in (vview c) 0)))
(array/push order f))))
(each f fns (visit (f :key)))
(each f order
(put report (f :key) (f :pt))
(def ptmap @{})
# escaped fn: its param types are untrustworthy (callers not all visible),
# so re-emit it WITHOUT seeding params (still re-embeds recompiled callees).
(unless (get esc (f :key))
(def ptmap @{})
(var concrete false)
(def pv (vview (f :params)))
(for i 0 (f :np)
(def t (in (f :pt) i))
(when (and t (not= t :any)) (set concrete true) (put ptmap (in pv i) t)))
(when concrete
(def def2 ((var-get f-reinfer) (f :def) ptmap))
(protect (eval (emit-ir ctx def2) (ctx-janet-env ctx))))))))
(when (and t (not= t :any)) (put ptmap (in pv i) t))))
(def def2 ((var-get f-reinfer) (f :def) ptmap))
(protect (eval (emit-ir ctx def2) (ctx-janet-env ctx))))))
report)
(defn ensure-macros-compiled!

View file

@ -0,0 +1,41 @@
# Collection-element types + HOF awareness, Phase 3 (jolt-d6u). A vector carries
# its element type ({:vec ELEM}); a reduce/map/filter closure over it gets that
# element type on its element param. So a lookup inside a reduce closure over a
# vector-of-structs specializes — no hint — WHEN the element type is provable.
(import ../../src/jolt/api :as api)
(import ../../src/jolt/backend :as backend)
(import ../../src/jolt/types :as types)
(import ../../src/jolt/reader :as reader)
(print "Type inference Phase 3 (jolt-d6u)...")
(os/setenv "JOLT_DIRECT_LINK" "1")
(def ctx (api/init {:compile? true}))
(api/eval-string ctx "(ns p3)")
(def pns (types/ctx-find-ns ctx "jolt.passes"))
(def reinfer (types/var-get (types/ns-find pns "reinfer-def")))
# helper: analyze a defn, reinfer with seeded param types, count guards
(defn guards [src ptmap]
(def d (backend/analyze-form ctx (reader/parse-string src)))
(length (string/find-all ":jolt/type" (string/format "%p" (backend/emit-ir ctx (reinfer d ptmap))))))
# a reduce closure's element param gets the vector's element type
(def red "(defn f [coll] (reduce (fn [acc h] (+ acc (:r h))) 0 coll))")
(assert (= 0 (guards red @{"coll" {:vec :struct-map}})) "reduce element typed -> bare lookup in closure")
(assert (= 1 (guards red @{"coll" {:vec :any}})) "reduce over vector of unknown -> guard kept")
(assert (= 1 (guards red @{})) "untyped coll -> guard kept")
# mapv over a vector-of-structs types the closure element too
(def mp "(defn g [coll] (mapv (fn [h] (:r h)) coll))")
(assert (= 0 (guards mp @{"coll" {:vec :struct-map}})) "mapv element typed -> bare lookup")
(assert (= 1 (guards mp @{"coll" {:vec :any}})) "mapv over unknown element -> guard")
# element type is DERIVED, not just seeded: a vector literal of structs, reduced
(def derived "(defn h2 [] (reduce (fn [acc x] (+ acc (:r x))) 0 [{:r 1 :g 2} {:r 3 :g 4}]))")
(assert (= 0 (guards derived @{})) "vector literal of structs -> element struct -> bare lookup")
# correctness: the specialized closures compute the same
(assert (= 4 (api/eval-string ctx "((fn [coll] (reduce (fn [acc h] (+ acc (:r h))) 0 coll)) [{:r 1} {:r 3}])")) "reduce value")
(assert (= 4 (api/eval-string ctx "(reduce (fn [acc x] (+ acc (:r x))) 0 [{:r 1 :g 2} {:r 3 :g 4}])")) "derived value")
(print "Type inference Phase 3 passed!")