whole-program closed-world inference, opt-in (jolt-t34)

JOLT_WHOLE_PROGRAM (requires direct-linking) defers the per-namespace inference
and runs ONE fixpoint over every user unit at once, so param types propagate
across namespace boundaries — a non-inlined fn's record params get proven from
its callers in another unit, which the per-ns pass can't see. Sound only under
the closed-world assumption (no later eval/redefinition) the flag asserts; slow,
memory-heavy builds are the documented trade-off (the reason it's opt-in).

infer-unit! now takes one ns-name OR a list; infer-program! gathers all recorded
user namespaces and runs the existing fixpoint over the union (re-emit was already
ns-agnostic — keyed by var-key, callee-first). The evaluator defers + records each
unit under the flag; run-main triggers infer-program! after all requires, before
-main. Off by default — per-ns behaviour unchanged.

Measured: a recursive (non-inlined) cross-ns record reader runs 1.66x faster
(8.9s -> 5.3s) — params proven -> bare-index reads. NOTE: small accessor fns are
INLINED cross-ns and records carry GLOBAL declared shapes, so most record reads
are already proven without this pass; the win is for non-inlined hot fns, and it's
the foundation for future whole-program work (devirtualization, unboxing).

Adds whole-program-test (subprocess soundness: per-ns and whole-program produce
identical results on a cross-ns record program).

src/jolt/api.janet

@@ -193,6 +193,7 @@
     # 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!)
+    (put (ctx :env) :infer-program! backend/infer-program!)   # jolt-t34 whole-program
     # 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).
@@ -233,6 +234,12 @@
       (and (get (ctx :env) :direct-linking?) (not (os/getenv "JOLT_NO_SHAPE"))))
     (put (ctx :env) :map-shapes?
       (and (os/getenv "JOLT_SHAPE") (not (os/getenv "JOLT_NO_SHAPE"))))
+    # Whole-program (Stalin) mode (jolt-t34): opt-in, closed-world. Defers the
+    # per-ns inference and runs one fixpoint over all units at the end (main, or a
+    # harness calling infer-program!). Needs direct-linking (the closed-world
+    # assumption); slow/memory-heavy builds are the documented trade-off.
+    (put (ctx :env) :whole-program?
+      (and (os/getenv "JOLT_WHOLE_PROGRAM") (get (ctx :env) :direct-linking?)))
     ctx))
 
 # --- Context snapshot/fork (cheap isolated copies) --------------------------

src/jolt/backend.janet

@@ -937,7 +937,12 @@
     :any))
 
 (defn infer-unit!
-  [ctx ns-name]
+  # ns-names-arg is one ns-name (per-namespace pass, the default) or a LIST of
+  # ns-names (whole-program pass, jolt-t34): gathering all units into ONE fixpoint
+  # propagates param types across namespace boundaries, which the per-ns pass
+  # can't (a fn's callers in another ns aren't visible when its own ns is typed).
+  [ctx ns-names-arg]
+  (def ns-names (if (indexed? ns-names-arg) ns-names-arg [ns-names-arg]))
   (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!")))
@@ -948,34 +953,37 @@
   (def f-set-rshapes (and pns (ns-find pns "set-record-shapes!")))   # jolt-t34
   (def f-set-mshapes (and pns (ns-find pns "set-map-shapes!")))      # jolt-t34
   (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-set-vtypes f-join f-infer-body f-reinfer f-reset-esc f-get-esc)
-    # gather single-fixed-arity fns AND non-fn defs that stashed a :def IR
+  (when (and f-set-rtenv f-set-vtypes f-join f-infer-body f-reinfer f-reset-esc f-get-esc)
+    # gather single-fixed-arity fns AND non-fn defs that stashed a :def IR, across
+    # every ns in ns-names (one ns for the per-unit pass, all for whole-program)
     (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)))
-        (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}))))
+    (each ns-name ns-names
+      (def ns (ctx-find-ns ctx ns-name))
+      (when ns
+        (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)))
+            (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))
       # jolt-t34: feed record-ctor shapes + the map-shaping flag to the inference
@@ -1071,6 +1079,19 @@
         (protect (eval (emit-ir ctx def2) (ctx-janet-env ctx))))))
   report)
 
+(defn infer-program!
+  "Whole-program closed-world pass (jolt-t34, opt-in JOLT_WHOLE_PROGRAM): run ONE
+  inference fixpoint over every user namespace at once, so param types propagate
+  across namespace boundaries (the per-ns pass can't see a fn's callers in other
+  units). Sound only under the closed-world assumption — direct-linking, no later
+  eval/redefinition — which the flag asserts. The recorded ns list is in load
+  (topological) order; the fixpoint is order-independent but re-emit is callee-
+  first regardless."
+  [ctx]
+  (def nses (get (ctx :env) :inferred-nses))
+  (when (and nses (> (length nses) 0))
+    (infer-unit! ctx nses)))
+
 (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

@@ -455,8 +455,15 @@
             # (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))))
+              (if (get (ctx :env) :whole-program?)
+                # whole-program (jolt-t34): defer — record the ns and run ONE
+                # fixpoint over all units later (the closed-world pass sees every
+                # caller, so cross-ns param types propagate)
+                (let [lst (or (get (ctx :env) :inferred-nses)
+                              (let [a @[]] (put (ctx :env) :inferred-nses a) a))]
+                  (array/push lst ns-name))
+                (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.

src/jolt/main.janet

@@ -430,6 +430,10 @@
   (try
     (do
       (load-string ctx (string "(require '[" ns-name "])"))
+      # whole-program (jolt-t34): every unit is loaded now — run the one closed-
+      # world fixpoint over all of them before -main, so cross-ns types propagate
+      (when (get (ctx :env) :whole-program?)
+        (when-let [ip (get (ctx :env) :infer-program!)] (protect (ip ctx))))
       (load-string ctx (string "(apply " ns-name "/-main *command-line-args*)")))
     ([err fib] (report-error err fib) (os/exit 1))))
 
@@ -514,6 +518,8 @@
     (and (get (ctx :env) :direct-linking?) (not (os/getenv "JOLT_NO_SHAPE"))))
   (put (ctx :env) :map-shapes?
     (and (os/getenv "JOLT_SHAPE") (not (os/getenv "JOLT_NO_SHAPE"))))
+  (put (ctx :env) :whole-program?
+    (and (os/getenv "JOLT_WHOLE_PROGRAM") (get (ctx :env) :direct-linking?)))
   (cond
     (empty? argv) (run-repl)
     (help-flags (argv 0)) (print-help)

test/integration/whole-program-test.janet

@@ -0,0 +1,44 @@
+# Whole-program (Stalin) mode (jolt-t34, opt-in JOLT_WHOLE_PROGRAM): one closed-
+# world inference fixpoint over ALL user namespaces, so param types propagate
+# across ns boundaries (a non-inlined fn's record params get proven from its
+# callers in another unit). This must be SOUND — same results as the per-ns
+# pass — which is what this test guards, by running a cross-namespace record
+# program both ways through the built binary and comparing output. Skips cleanly
+# if build/jolt is absent (source-only test run).
+(def jolt "build/jolt")
+
+(defn- run [env-extra]
+  (def dir (string (or (os/getenv "TMPDIR") "/tmp") "/jolt-wp-test"))
+  (os/mkdir dir)
+  (spit (string dir "/wputil.clj")
+        (string "(ns wputil)\n"
+                "(defrecord V [x y z])\n"
+                # recursive => never inlined; params proven only whole-program
+                "(defn dot [a b n]\n"
+                "  (if (<= n 0) 0.0\n"
+                "    (+ (* (:x a) (:x b)) (* (:y a) (:y b)) (* (:z a) (:z b)) (dot a b (dec n)))))\n"))
+  (spit (string dir "/wpmain.clj")
+        (string "(ns wpmain (:require [wputil :as v]))\n"
+                "(defn -main []\n"
+                "  (loop [i 0 acc 0.0]\n"
+                "    (if (< i 1000)\n"
+                "      (let [a (v/->V (double i) 2.0 3.0) b (v/->V 1.0 (double i) 0.5)]\n"
+                "        (recur (inc i) (+ acc (v/dot a b 2))))\n"
+                "      (println \"sum\" acc))))\n"))
+  (def out (string dir "/out.txt"))
+  (def jbin (string (os/cwd) "/" jolt))
+  (def cmd (string env-extra "JOLT_DIRECT_LINK=1 JOLT_PATH=" dir " " jbin
+                   " -m wpmain > " out " 2>&1"))
+  (os/execute ["sh" "-c" cmd] :p)
+  (string/trimr (slurp out)))
+
+(if (not (os/stat jolt))
+  (print "whole-program: SKIP (no build/jolt — run from source)")
+  (let [per-ns (run "")
+        whole  (run "JOLT_WHOLE_PROGRAM=1 ")]
+    (printf "  per-ns:        %s" per-ns)
+    (printf "  whole-program: %s" whole)
+    (if (and (= per-ns whole) (string/has-prefix? "sum" per-ns))
+      (print "whole-program: results match — sound")
+      (do (printf "whole-program: MISMATCH per-ns=%q whole=%q" per-ns whole)
+          (os/exit 1)))))