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

feat: Phase 1 inter-procedural collection-type inference (jolt-767)

Browse source 
Closed-world (optimization mode): after a unit loads, infer-unit! runs a
whole-unit fixpoint over the call graph and recompiles. A fn's param types are
the lub of its in-unit call-site arg types; its return type is the lub of its
tail positions; iterated to a least fixpoint. Param types are RECOMPUTED FRESH
each iteration (not accumulated) because :any is the lattice top — joining an
early-iteration :any would poison the result permanently. Closures inherit the
enclosing tenv so captured locals keep their types (their own params shadow to
:any). A fn whose var escapes as a VALUE keeps :any params (its callers aren't
all visible). Each fn is then re-inferred with its param types seeded and
re-emitted; recompiled bodies are semantically identical, so correctness holds
regardless of order. Sound under source distribution + whole-program compile
(the consumer compiles all call sites together).

Plumbing: the portable pass (jolt.passes) gained inter-procedural primitives —
set-rtenv!, infer-body (types a body, collects its call sites), reinfer-def
(seeds param types), and escape tracking. The back end stashes each
single-fixed-arity defn's :def IR (:infer-ir); the evaluator triggers
infer-unit! after a unit loads (via an env hook, opt mode only).

Result and honest finding: the fixpoint correctly types scalar-flowing params
(ray-cast/hit-all/hit-sphere all get the ray param as :struct-map, no hint),
but the ray tracer does NOT speed up — its dominant lookups are on `hittable`,
the element of the `hittables` vector threaded through `reduce`, which stays
:any. Typing it needs collection-element types (vector<struct>) plus HOF-element
awareness (knowing reduce applies the closure to elements), which is beyond
inter-procedural param inference. The explicit ^:struct hint reaches it (it
types the reduce closure param directly), which is why the hinted run is 1.22x.

Verified: conformance 335/335 x3, full jpm test; new type-infer-phase1-test
pins the fixpoint, the escape gate, the seeded re-inference, and correctness.
This commit is contained in:
Yogthos 2026-06-13 04:45:13 -04:00
parent 3c20383851
commit ea1d9a23e1
5 changed files with 251 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

83
jolt-core/jolt/passes.clj
View file

@ -731,15 +731,31 @@
"bit-and" "bit-or" "bit-xor" "count"})
(def ^:private vector-ret-fns #{"vec" "vector" "mapv" "filterv" "subvec"})
;; Inter-procedural state (jolt-767, Phase 1). The Janet orchestrator (backend
;; infer-unit!) drives a whole-unit fixpoint: before typing a fn body it installs
;; the current return-type estimates of all unit fns here, and after typing it
;; reads back the call sites this body made (callee + inferred arg types) to
;; propagate into callee param types. Both are plain module state, like `dirty`.
(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)
(defn- var-key [fnode] (str (get fnode :ns) "/" (get fnode :name)))
(defn- call-ret-type [fnode]
(let [nm (cond
(and (= :var (get fnode :op)) (= "clojure.core" (get fnode :ns))) (get fnode :name)
(= :host (get fnode :op)) (get fnode :name)
:else nil)]
(let [op (get fnode :op)]
(cond
(nil? nm) :any
;; a user fn whose return type the fixpoint has estimated
(= op :var) (let [r (get @rtenv-box (var-key fnode))]
(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
:else :any))))
(= op :host) (let [nm (get fnode :name)]
(cond (contains? truthy-ret-fns nm) :truthy
(contains? vector-ret-fns nm) :vector
:else :any))
:else :any)))
(defn- infer
@ -781,10 +797,19 @@
[(nth r 0) (assoc node :statements stmts :ret (nth r 1))])
(= op :throw)
[: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])
(= op :invoke)
(let [fr (infer (get node :fn) tenv)
args (mapv (fn [a] (nth (infer a tenv) 1)) (get node :args))]
[(call-ret-type (get node :fn)) (assoc node :fn (nth fr 1) :args args)])
(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))])
(= op :let)
(let [res (reduce (fn [acc b]
(let [te (nth acc 0) binds (nth acc 1)
@ -803,7 +828,14 @@
(= op :recur)
[:any (assoc node :args (mapv (fn [a] (nth (infer a tenv) 1)) (get node :args)))]
(= op :fn)
[:any (assoc node :arities (mapv (fn [a] (assoc a :body (nth (infer (get a :body) {}) 1)))
;; a closure inherits the enclosing tenv so CAPTURED locals keep their
;; types (e.g. a reduce closure that calls (f captured-struct ...)); its own
;; params/rest shadow to :any (unknown until Phase 1 types them via callers).
[:any (assoc node :arities
(mapv (fn [a]
(let [pe (reduce (fn [e p] (assoc e p :any)) tenv (get a :params))
pe (if (get a :rest) (assoc pe (get a :rest) :any) pe)]
(assoc a :body (nth (infer (get a :body) pe) 1))))
(get node :arities)))]
(= op :def)
[:any (assoc node :init (nth (infer (get node :init) tenv) 1))]
@ -816,6 +848,39 @@
(defn- infer-top [node] (nth (infer node {}) 1))
;; --- Inter-procedural driver API (jolt-767) consumed by the back end --------
(defn set-rtenv!
"Install the current return-type estimates (a map \"ns/name\" -> type) used to
type call results during the fixpoint."
[m] (reset! rtenv-box m))
(defn reset-escapes! [] (reset! escapes-box #{}))
(defn collected-escapes [] (vec @escapes-box))
(defn infer-body
"Type `body` under tenv (local-name -> type). Returns [ret-type node' calls],
where calls is the [[\"ns/name\" [arg-types...]] ...] this body invokes (for
propagating into callee param types). Also accumulates escapes (read with
collected-escapes after a full sweep)."
[body tenv]
(reset! calls-box [])
(let [r (infer body tenv)]
[(nth r 0) (nth r 1) @calls-box]))
(defn reinfer-def
"Re-run inference on a stashed :def's fn arity bodies with param types seeded
(ptmap: param-name -> type), returning the def with annotated bodies. The back
end emits the result directly (no further passes), so the param-typed lookups
keep their specialization. Used by the inter-procedural recompile."
[def-node ptmap]
(let [fnode (get def-node :init)]
(if (= :fn (get fnode :op))
(assoc def-node :init
(assoc fnode :arities
(mapv (fn [a] (assoc a :body (nth (infer (get a :body) ptmap) 1)))
(get fnode :arities))))
def-node)))
(defn run-passes
"All passes, in order. The back end applies this to every analyzed form. When
inlining is enabled for the unit (user code under direct-linking, jolt-87f),

4
src/jolt/api.janet
View file

@ -189,6 +189,10 @@
# — that would be circular — so it reads this hook). Without it, required
# namespaces ran interpreted-only.
(put (ctx :env) :toplevel-eval eval-toplevel)
# Inter-procedural type-inference hook (jolt-767): the evaluator calls this
# after a unit finishes loading (optimization mode only). Installed here to
# avoid an evaluator->backend circular import.
(put (ctx :env) :infer-unit! backend/infer-unit!)
# Stateful primitives as ctx-capturing clojure.core fns (protocol-dispatch,
# register-method, …) — so the protocol macros compile to plain invokes. Must
# precede the overlay (its defprotocol/extend-type expansions call these).

109
src/jolt/backend.janet
View file

@ -78,7 +78,10 @@
(when (= 1 (length arities))
(def ar (norm-node (in arities 0)))
(unless (ar :rest)
(put cell :inline-ir {:params (ar :params) :body (ar :body)}))))))
(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))))))
# 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)
@ -745,6 +748,110 @@
(++ n))))
n)
# Inter-procedural collection-type inference + recompile (jolt-767, Phase 1),
# closed-world / optimization mode. After a unit loads, every single-fixed-arity
# fn stashed a post-pass :def IR (:infer-ir). We:
# 1. run a whole-unit fixpoint: a fn's param types = lub of its in-unit
# call-site arg types (computed by jolt.passes/infer-body); a fn whose var
# escapes as a VALUE keeps :any params (its callers aren't all visible).
# 2. re-infer + re-emit each fn body with its param types seeded, so
# param-dependent lookups specialize (drop the :jolt/type guard).
# 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 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-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!")))
(def f-get-esc (and pns (ns-find pns "collected-escapes")))
(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
(def fns @[])
(def by-key @{})
(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)
((var-get f-reset-esc))
# --- param/return-type fixpoint (chaotic iteration to the 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
# from the current state lets a param refine as its callers' types improve.
(var prev-rt @{})
(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
(each f fns
(def tenv @{})
(def pv (vview (f :params)))
(for i 0 (f :np) (when (in (f :pt) i) (put tenv (in pv i) (in (f :pt) i))))
(def res (vview ((var-get f-infer-body) (f :body) tenv)))
(put f :tret (in res 0))
(put f :tcalls (in res 2)))
# 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))))
(each f fns
(each c (vview (f :tcalls))
(def cv (vview c))
(def npa (get newpt (in cv 0)))
(when npa
(def callee (get by-key (in cv 0)))
(def ats (vview (in cv 1)))
(def lim (min (length ats) (callee :np)))
(for i 0 lim (put npa i (itype-join (in npa i) (in ats i)))))))
# commit + detect change
(set changed false)
(def nrt @{})
(each f fns
(def np (get newpt (f :key)))
(for i 0 (f :np) (when (not= (in np i) (in (f :pt) i)) (set changed true)))
(when (not= (f :tret) (f :ret)) (set changed true))
(put f :pt np)
(put f :ret (f :tret))
(when (f :tret) (put nrt (f :key) (f :tret))))
(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
(put report (f :key) (f :pt))
(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))))))))
report)
(defn ensure-macros-compiled!
"Called once the overlay is fully loaded (api/load-core-overlay!): ensure the
analyzer is built, then run the staged macro-recompile pass so the early

8
src/jolt/evaluator.janet
View file

@ -421,6 +421,14 @@
(if path
(load-ns-source ctx (slurp path) path)
(load-ns-source ctx embedded (string ns-name " (stdlib)")))
# Inter-procedural collection-type inference (jolt-767): once the whole
# unit is loaded, run the closed-world fixpoint + recompile so param-
# dependent lookups specialize. Only in optimization mode; best-effort
# (a failure here must not break loading). Hook installed by the api to
# avoid an evaluator->backend circular import.
(when (get (ctx :env) :inline?)
(when-let [iu (get (ctx :env) :infer-unit!)]
(protect (iu ctx ns-name))))
# Record load order for tooling (uberscript): a dependency finishes
# loading before its requirer, so this is topological. Skip the
# baked-in stdlib — it's part of the runtime, not something to bundle.

54
test/integration/type-infer-phase1-test.janet Normal file
View file

@ -0,0 +1,54 @@
# Inter-procedural collection-type inference, Phase 1 (jolt-767): closed-world.
# A whole-unit fixpoint propagates collection types through the call graph — a
# fn's param types become the lub of its in-unit call-site arg types — so a
# param that always receives a struct map gets typed and its lookups specialize,
# with no hint. Fns whose var escapes as a value keep :any params (their callers
# aren't all visible). Sound under source distribution + whole-program compile.
(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 1 (jolt-767)...")
(os/setenv "JOLT_DIRECT_LINK" "1")
(def ctx (api/init {:compile? true}))
(api/eval-string ctx "(ns p1)")
# closed-world unit. mk is small (inlined away). rd is RECURSIVE, so it survives
# inlining and is called via its var — exactly the shape (big/recursive fn with
# escaping-from-the-caller params) that inter-procedural inference targets. Its
# param v flows from mk's struct-map literal (after mk inlines into drv).
(each s ["(defn mk [a b] {:r a :g b})"
"(defn rd [v n] (if (< n 1) (:r v) (rd v (dec n))))"
"(defn drv [] (rd (mk 1 2) 3))"
# esc's var is used as a VALUE (passed to mapv) -> params must stay :any
"(defn esc [w] (:r w))"
"(defn use-esc [xs] (mapv esc xs))"]
(api/eval-string ctx s))
(def report (backend/infer-unit! ctx "p1"))
# --- the fixpoint computed the right param types -----------------------------
# rd's param v flows from mk's struct-map result (mk inlines to a struct literal
# in drv) and stays struct across the recursive self-call -> :struct-map
(assert (= :struct-map (in (get report "p1/rd") 0)) (string "rd param v: " (in (get report "p1/rd") 0)))
# esc escaped (passed to mapv) -> param stays unknown (:any / nil), NOT struct
(assert (not= :struct-map (in (get report "p1/esc") 0)) "escaping fn param not inferred struct")
# --- the seeded re-inference drops the guard for a struct param --------------
# (on a FRESH analysis, since infer-unit! re-stashes the already-specialized body)
(def pns (types/ctx-find-ns ctx "jolt.passes"))
(def reinfer (types/ns-find pns "reinfer-def"))
(def rd-def (backend/analyze-form ctx (reader/parse-string "(defn rdx [v n] (if (< n 1) (:r v) (rdx v (dec n))))")))
(defn guards-seeded [ptmap]
(length (string/find-all ":jolt/type" (string/format "%p" (backend/emit-ir ctx ((types/var-get reinfer) rd-def ptmap))))))
(assert (= 0 (guards-seeded @{"v" :struct-map})) "struct param -> bare lookup")
(assert (= 1 (guards-seeded @{})) "no param type -> guard kept")
# --- correctness: recompiled unit still computes the same --------------------
(assert (= 1 (api/eval-string ctx "(p1/drv)")) "drv correct after recompile")
(assert (= 7 (api/eval-string ctx "(p1/rd {:r 7 :g 8} 0)")) "rd correct on a struct")
(assert (= nil (api/eval-string ctx "(p1/rd (hash-map :r nil) 0)")) "rd correct on a phm (key present, nil)")
(assert (deep= [1 1] (api/normalize-pvecs (api/eval-string ctx "(p1/use-esc [{:r 1} {:r 1}])"))) "escaping fn still correct")
(print "Type inference Phase 1 passed!")