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

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.

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
+  (let [op (get fnode :op)]
-             (and (= :var (get fnode :op)) (= "clojure.core" (get fnode :ns))) (get fnode :name)
-             (= :host (get fnode :op)) (get fnode :name)
-             :else nil)]
     (cond
-      (nil? nm) :any
+      ;; a user fn whose return type the fixpoint has estimated
-      (contains? truthy-ret-fns nm) :truthy
+      (= op :var) (let [r (get @rtenv-box (var-key fnode))]
-      (contains? vector-ret-fns nm) :vector
+                    (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)
+      (let [fnode (get node :fn)
-            args (mapv (fn [a] (nth (infer a tenv) 1)) (get node :args))]
+            iscall-var (= :var (get fnode :op))
-        [(call-ret-type (get node :fn)) (assoc node :fn (nth fr 1) :args args)])
+            ;; 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)))
+      ;; a closure inherits the enclosing tenv so CAPTURED locals keep their
-                                       (get node :arities)))]
+      ;; 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),

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).

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

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.

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!")