From ea1d9a23e1c977eefbc9ca3dc2c0d4460b297748 Mon Sep 17 00:00:00 2001 From: Yogthos Date: Sat, 13 Jun 2026 04:45:13 -0400 Subject: [PATCH] feat: Phase 1 inter-procedural collection-type inference (jolt-767) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) 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. --- jolt-core/jolt/passes.clj | 89 ++++++++++++-- src/jolt/api.janet | 4 + src/jolt/backend.janet | 109 +++++++++++++++++- src/jolt/evaluator.janet | 8 ++ test/integration/type-infer-phase1-test.janet | 54 +++++++++ 5 files changed, 251 insertions(+), 13 deletions(-) create mode 100644 test/integration/type-infer-phase1-test.janet diff --git a/jolt-core/jolt/passes.clj b/jolt-core/jolt/passes.clj index 5e73c06..eecba2d 100644 --- a/jolt-core/jolt/passes.clj +++ b/jolt-core/jolt/passes.clj @@ -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 - (contains? truthy-ret-fns nm) :truthy - (contains? vector-ret-fns nm) :vector + ;; 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,8 +828,15 @@ (= 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))) - (get node :arities)))] + ;; 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))] (= op :try) @@ -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), diff --git a/src/jolt/api.janet b/src/jolt/api.janet index cb84cb0..ed1d6b8 100644 --- a/src/jolt/api.janet +++ b/src/jolt/api.janet @@ -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). diff --git a/src/jolt/backend.janet b/src/jolt/backend.janet index 04c324f..96386a6 100644 --- a/src/jolt/backend.janet +++ b/src/jolt/backend.janet @@ -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 diff --git a/src/jolt/evaluator.janet b/src/jolt/evaluator.janet index 096501e..694b783 100644 --- a/src/jolt/evaluator.janet +++ b/src/jolt/evaluator.janet @@ -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. diff --git a/test/integration/type-infer-phase1-test.janet b/test/integration/type-infer-phase1-test.janet new file mode 100644 index 0000000..31735f8 --- /dev/null +++ b/test/integration/type-infer-phase1-test.janet @@ -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!")