diff --git a/src/jolt/api.janet b/src/jolt/api.janet index dd6fa75..a1bf9a4 100644 --- a/src/jolt/api.janet +++ b/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) -------------------------- diff --git a/src/jolt/backend.janet b/src/jolt/backend.janet index b8e8d8f..d70ffed 100644 --- a/src/jolt/backend.janet +++ b/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 diff --git a/src/jolt/evaluator.janet b/src/jolt/evaluator.janet index cfe20c6..7cd9cfe 100644 --- a/src/jolt/evaluator.janet +++ b/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. diff --git a/src/jolt/main.janet b/src/jolt/main.janet index fdce089..c97c56e 100644 --- a/src/jolt/main.janet +++ b/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) diff --git a/test/integration/whole-program-test.janet b/test/integration/whole-program-test.janet new file mode 100644 index 0000000..ba8aea9 --- /dev/null +++ b/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)))))