diff --git a/jolt-core/jolt/passes.clj b/jolt-core/jolt/passes.clj index 075a0d3..9b3af8f 100644 --- a/jolt-core/jolt/passes.clj +++ b/jolt-core/jolt/passes.clj @@ -17,7 +17,7 @@ [jolt.passes.fold :refer [const-fold]] [jolt.passes.inline :refer [inline-node flatten-lets scalar-replace dirty set-rec-shapes!]] [jolt.passes.types :refer [run-inference - check-form infer-body reinfer-def + check-form infer-body reinfer-def phint-seed set-rtenv! set-vtypes! join-types set-record-shapes! set-map-shapes! set-protocol-methods! reset-escapes! collected-escapes diff --git a/jolt-core/jolt/passes/types.clj b/jolt-core/jolt/passes/types.clj index 1ae345c..6027331 100644 --- a/jolt-core/jolt/passes/types.clj +++ b/jolt-core/jolt/passes/types.clj @@ -791,6 +791,24 @@ (get fnode :arities)))) def-node))) +(defn phint-seed + "Positional declared-hint type seeds for a fn arity. Given the param-name + vector and the arity's :phints (a seq of [name ctor-key] pairs), return a + vector parallel to params whose slot i is the resolved record TYPE of that + param's ^Record hint (via the record-shapes registry), or nil. The + whole-program fixpoint seeds these as a param-type FLOOR so a declared hint + propagates to a fn's callees DURING inference — not only at the final re-emit + (reinfer-def). Without it a hinted param with no callers stays :any through the + fixpoint, so a field read off it (e.g. (:origin ^Ray r)) never tells a shared + callee its arg is a Vec3 (jolt-3ko)." + [params phints] + (let [m (reduce (fn [acc pr] (assoc acc (nth pr 0) (nth pr 1))) {} phints)] + (mapv (fn [nm] + (let [ck (get m nm) + e (and ck (get @record-shapes-box ck))] + (when e (record-type-from-entry e type-depth)))) + params))) + ;; Piggyback checking (jolt audit). In direct-link mode infer-top already runs ;; one inference pass for specialization; turning checking? on during it makes ;; the success checker nearly free there (no extra traversal — just the diff --git a/src/jolt/backend.janet b/src/jolt/backend.janet index 5025e3d..6656fc3 100644 --- a/src/jolt/backend.janet +++ b/src/jolt/backend.janet @@ -1023,6 +1023,7 @@ (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-set-pmethods (and pns (ns-find pns "set-protocol-methods!"))) # jolt-41m + (def f-phint-seed (and pns (ns-find pns "phint-seed"))) # jolt-3ko (def f-get-esc (and pns (ns-find pns "collected-escapes"))) (def report @{}) (when (and f-set-rtenv f-set-vtypes f-join f-infer-body f-reinfer f-reset-esc f-get-esc) @@ -1048,7 +1049,8 @@ (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}) + :phints (ar :phints) :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) @@ -1060,6 +1062,13 @@ # jolt-t34: feed record-ctor shapes + the map-shaping flag to the inference (when f-set-rshapes ((var-get f-set-rshapes) (or (get (ctx :env) :record-shapes) @{}))) (when f-set-mshapes ((var-get f-set-mshapes) (get (ctx :env) :map-shapes?))) + # jolt-3ko: resolve each fn's declared ^Record param hints to positional + # type seeds (needs the registry above). Seeded as a param-type FLOOR in the + # fixpoint so a hinted param propagates its (field-read) types to callees + # during inference, not only at the final re-emit. + (when f-phint-seed + (each f fns + (put f :phint-types (vview ((var-get f-phint-seed) (f :params) (or (f :phints) [])))))) # jolt-41m: feed the protocol-method registry for devirtualization (when f-set-pmethods ((var-get f-set-pmethods) (or (get (ctx :env) :protocol-methods) @{}))) # --- param/return/value-type fixpoint (chaotic iteration to LEAST fixpoint) --- @@ -1083,9 +1092,17 @@ # infer each def's VALUE type from its init (each dv defs (put dv :tvt (in (vview ((var-get f-infer-body) (dv :init) @{})) 0))) - # recompute param types FRESH (start at bottom = nil) from this round's calls + # recompute param types FRESH (start at bottom) from this round's calls. + # Bottom is nil, EXCEPT a declared ^Record hint seeds its slot as a floor + # (jolt-3ko) — a fixed declared type can't poison the fixpoint the way an + # early-iteration :any would, and it lets a hinted param propagate to its + # callees even when it has no callers of its own. (def newpt @{}) - (each f fns (put newpt (f :key) (array/new-filled (f :np)))) + (each f fns + (def npa (array/new-filled (f :np))) + (def phts (f :phint-types)) + (when phts (for i 0 (f :np) (when (in phts i) (put npa i (in phts i))))) + (put newpt (f :key) npa)) (each f fns (each c (vview (f :tcalls)) (def cv (vview c)) diff --git a/test/integration/phint-propagation-test.janet b/test/integration/phint-propagation-test.janet new file mode 100644 index 0000000..abba90b --- /dev/null +++ b/test/integration/phint-propagation-test.janet @@ -0,0 +1,77 @@ +# Declared record param hints propagate THROUGH the whole-program fixpoint +# (jolt-3ko). A ^Record param's field reads type to the field's record type, and +# when such a field-read value is passed to a SHARED helper (no hints, inferred +# from call sites), the helper's params must pick up that record type so its own +# field reads bare-index. Before the fix, ^-hints were applied only at the final +# re-emit (reinfer-def), not during the fixpoint, so a hinted param with no +# callers stayed :any during inference and never propagated to its callees — +# exactly why the ray tracer's vec ops (called with (:origin ray) etc.) stayed +# unproven even under whole-program optimization. +(import ../../src/jolt/api :as api) +(import ../../src/jolt/backend :as backend) +(import ../../src/jolt/types :as types) + +(print "phint propagation through the fixpoint (jolt-3ko)...") + +(def dir (string (or (os/getenv "TMPDIR") "/tmp") "/jolt-phint-prop")) +(os/mkdir dir) +(spit (string dir "/pp.clj") + (string + "(ns pp)\n" + "(defrecord Vec3 [r g b])\n" + "(defrecord Ray [^Vec3 origin ^Vec3 direction])\n" + # shared ops, NO hints — must be inferred from the call sites below. A + # `loop` makes them inline-INELIGIBLE so they stay real calls (a small fn + # would be spliced into the caller, where the caller's hint already proves + # it); the point here is propagation to a separate callee's PARAMS. + "(defn dot [a b]\n" + " (loop [i 0 s 0.0]\n" + " (if (< i 1) (recur (inc i) (+ (+ (* (:r a) (:r b)) (* (:g a) (:g b))) (* (:b a) (:b b)))) s)))\n" + "(defn add [a b]\n" + " (loop [i 0 s nil]\n" + " (if (< i 1) (recur (inc i) (->Vec3 (+ (:r a) (:r b)) (+ (:g a) (:g b)) (+ (:b a) (:b b)))) s)))\n" + # callers pass field-read Vec3s from a ^Ray param into the shared ops + "(defn use-dot [^Ray r] (dot (:origin r) (:direction r)))\n" + "(defn use-add [^Ray r] (add (:origin r) (:direction r)))\n")) + +(os/setenv "JOLT_DIRECT_LINK" "1") +(os/setenv "JOLT_WHOLE_PROGRAM" "1") +(os/setenv "JOLT_PATH" dir) +(def ctx (api/init {:compile? true})) +(api/eval-string ctx "(require '[pp])") +(def report (backend/infer-program! ctx)) + +(def pns (types/ctx-find-ns ctx "jolt.passes")) +(def reinfer (types/var-get (types/ns-find pns "reinfer-def"))) +(def ppns (types/ctx-find-ns ctx "pp")) +(defn ptmap-for [fname params] + (def pts (get report (string "pp/" fname))) + (def m @{}) (when pts (var i 0) (each p params (put m p (get pts i)) (++ i))) m) +(defn guards [fname params] + (def cell (get (ppns :mappings) fname)) + (length (string/find-all ":jolt/type" + (string/format "%p" (backend/emit-ir ctx (reinfer (get cell :infer-ir) (ptmap-for fname params))))))) + +(var fails 0) +(defn check [label got expected] + (if (= got expected) (print " ok " label) + (do (++ fails) (printf " FAIL %s: want %p got %p" label expected got)))) + +# the shared ops' params get the Vec3 record type from the field-read call args, +# so all their field reads bare-index (no :jolt/type guard) +(check "dot params typed Vec3 -> reads bare" (guards "dot" ["a" "b"]) 0) +(check "add params typed Vec3 -> reads bare" (guards "add" ["a" "b"]) 0) +# the hinted callers were already fine (re-emit applies the phint) +(check "use-dot hinted reads bare" (guards "use-dot" ["r"]) 0) + +# the report shows the shared ops' params as a Vec3 struct (has :shape / :type) +(def dot-a (get (get report "pp/dot") 0)) +(check "dot param a is a struct type" (truthy? (and dot-a (or (get dot-a :shape) (get dot-a :struct)))) true) + +# correctness: results are unchanged +(check "dot computes" (api/eval-string ctx "(pp/dot (pp/->Vec3 1 2 3) (pp/->Vec3 4 5 6))") 32) +(check "use-dot computes" + (api/eval-string ctx "(pp/use-dot (pp/->Ray (pp/->Vec3 1 2 3) (pp/->Vec3 4 5 6)))") 32) + +(if (> fails 0) (do (printf "phint-propagation: %d FAILED" fails) (os/exit 1)) + (print "phint propagation (jolt-3ko) passed!"))