diff --git a/jolt-core/jolt/passes.clj b/jolt-core/jolt/passes.clj index eecba2d..037f3ab 100644 --- a/jolt-core/jolt/passes.clj +++ b/jolt-core/jolt/passes.clj @@ -716,15 +716,25 @@ ;; dynamic guard in place. Sound by construction: a concrete type is assigned ;; only when proven, so a wrong bare get is impossible. ;; -;; Lattice values: :struct-map (raw-get-safe), :phm-map, :vector, :set, :truthy -;; (a provably non-nil/non-false scalar — numbers, strings, keywords), :any (top). -(defn- join [a b] (if (= a b) a :any)) +;; Lattice values: :struct-map (raw-get-safe), :phm-map, :set, :truthy (a +;; provably non-nil/non-false scalar), :any (top), and a PARAMETRIC vector type +;; {:vec ELEM} (jolt-d6u, Phase 3) carrying its element type so a reduce/map +;; closure over it can type its element param. {:vec ELEM} is a small struct, so +;; it compares by value on both the Clojure and the Janet (orchestrator) side. +(defn- velem [t] (get t :vec)) +(defn- vec-type? [t] (some? (velem t))) +(defn- mk-vec [t] {:vec (if t t :any)}) +(defn- join [a b] + (cond + (= a b) a + (and (vec-type? a) (vec-type? b)) (mk-vec (join (velem a) (velem b))) + :else :any)) (defn- struct-safe? [t] (= t :struct-map)) ;; a value whose type guarantees it is neither nil nor false — the back end only ;; builds a struct (vs a phm) when every value is truthy, so a map literal is a -;; struct only when all its values have a truthy type. +;; struct only when all its values have a truthy type. Collections are non-nil. (defn- truthy-type? [t] - (or (= t :truthy) (= t :struct-map) (= t :phm-map) (= t :vector) (= t :set))) + (or (= t :truthy) (= t :struct-map) (= t :phm-map) (= t :set) (vec-type? t))) (def ^:private truthy-ret-fns #{"+" "-" "*" "/" "inc" "dec" "mod" "rem" "quot" "min" "max" "abs" @@ -739,6 +749,14 @@ (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) +;; jolt-d6u: a var reference's VALUE type — a fn var is :truthy (non-nil), a def +;; var carries its inferred init type (e.g. a color table -> {:vec :struct-map}). +;; The orchestrator populates this from sealed (opt-mode) cell roots + def inits. +(def ^:private vtype-box (atom {})) ;; "ns/name" -> value type + +;; fns that RETURN an element of their (first) collection arg, so a lookup on the +;; result of (rand-nth coll-of-structs) etc. types as the element. +(def ^:private elem-fns #{"rand-nth" "first" "peek" "last" "nth" "fnext" "second"}) (defn- var-key [fnode] (str (get fnode :ns) "/" (get fnode :name))) @@ -750,14 +768,44 @@ (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 + (contains? vector-ret-fns nm) (mk-vec :any) :else :any)))) (= op :host) (let [nm (get fnode :name)] (cond (contains? truthy-ret-fns nm) :truthy - (contains? vector-ret-fns nm) :vector + (contains? vector-ret-fns nm) (mk-vec :any) :else :any)) :else :any))) +(declare infer) + +;; HOFs that apply their fn arg to the ELEMENTS of a collection (jolt-d6u, +;; Phase 3). :epos is which param of the fn receives an element. reduce is +;; handled separately (its arity changes the coll position, and its closure +;; also takes an accumulator). +(def ^:private hof-table + {"map" {:epos 0} "mapv" {:epos 0} "filter" {:epos 0} "filterv" {:epos 0} + "keep" {:epos 0} "remove" {:epos 0} "run!" {:epos 0} "mapcat" {:epos 0}}) + +(defn- infer-fn-seeded + "Infer a fn-literal passed to a HOF, seeding the given params to element/accum + types (seeds: param-index -> type), other params :any, captured locals from + tenv. Returns [ret-type node'] — ret is the lub of arity tail types, used to + type the HOF result (e.g. reduce's accumulator, mapv's element)." + [node seeds tenv] + (let [res (mapv (fn [a] + (let [params (get a :params) + pe (reduce (fn [e i] + (assoc e (nth params i) + (let [s (get seeds i)] (if s s :any)))) + tenv (range (count params))) + pe (if (get a :rest) (assoc pe (get a :rest) :any) pe) + br (infer (get a :body) pe)] + [(nth br 0) (assoc a :body (nth br 1))])) + (get node :arities)) + rets (mapv (fn [r] (nth r 0)) res) + ret (if (empty? rets) :any (reduce join (first rets) (rest rets)))] + [ret (assoc node :arities (mapv (fn [r] (nth r 1)) res))])) + (defn- infer "Returns [type node'] — the inferred type of node and node with struct-safe :local references annotated :hint :struct. tenv maps in-scope local names to @@ -782,7 +830,10 @@ :struct-map :any)] [t (assoc node :pairs (mapv (fn [r] [(nth r 0) (nth r 1)]) res))]) (= op :vector) - [:vector (assoc node :items (mapv (fn [x] (nth (infer x tenv) 1)) (get node :items)))] + (let [irs (mapv (fn [x] (infer x tenv)) (get node :items)) + ets (mapv (fn [r] (nth r 0)) irs) + el (if (empty? ets) :any (reduce join (first ets) (rest ets)))] + [(mk-vec el) (assoc node :items (mapv (fn [r] (nth r 1)) irs))]) (= op :set) [:set (assoc node :items (mapv (fn [x] (nth (infer x tenv) 1)) (get node :items)))] (= op :if) @@ -799,17 +850,68 @@ [: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]) + ;; Its VALUE type comes from vtype-box (a fn is :truthy, a def carries its + ;; inferred type); unknown -> :any. + (= op :var) (do (swap! escapes-box conj (var-key node)) + [(let [vt (get @vtype-box (var-key node))] (if vt vt :any)) node]) (= op :invoke) (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))]) + cn (when (and iscall-var (= "clojure.core" (get fnode :ns))) (get fnode :name)) + args (get node :args) + n (count args)] + (cond + ;; reduce over a typed vector with a fn-literal (jolt-d6u): seed the + ;; closure's accumulator (param 0) to the init type and its element + ;; (param 1) to the vector's element type, so its body — and any calls + ;; it makes — see those types. + (and (= cn "reduce") (>= n 2) (= :fn (get (nth args 0) :op))) + (let [three (>= n 3) + coll-r (infer (nth args (if three 2 1)) tenv) + init-r (when three (infer (nth args 1) tenv)) + et (let [ct (nth coll-r 0)] (if (vec-type? ct) (velem ct) :any)) + init-t (if init-r (nth init-r 0) :any) + fn-r (infer-fn-seeded (nth args 0) {0 init-t 1 et} tenv)] + [(join init-t (nth fn-r 0)) + (assoc node :args (if three + [(nth fn-r 1) (nth init-r 1) (nth coll-r 1)] + [(nth fn-r 1) (nth coll-r 1)]))]) + ;; map/mapv/filter/... over a typed vector with a fn-literal: seed the + ;; fn's element param; mapv/filterv produce a typed vector. + (and cn (get hof-table cn) (>= n 2) (= :fn (get (nth args 0) :op))) + (let [coll-r (infer (nth args 1) tenv) + et (let [ct (nth coll-r 0)] (if (vec-type? ct) (velem ct) :any)) + fn-r (infer-fn-seeded (nth args 0) {(get (get hof-table cn) :epos) et} tenv) + rt (cond (= cn "mapv") (mk-vec (nth fn-r 0)) + (= cn "filterv") (mk-vec et) + :else :any)] + [rt (assoc node :args [(nth fn-r 1) (nth coll-r 1)])]) + ;; conj/into: track the element type of a vector being grown. + (and (or (= cn "conj") (= cn "into")) (>= n 1)) + (let [ares (mapv (fn [a] (infer a tenv)) args) + base (nth (nth ares 0) 0) + rest-ts (mapv (fn [r] (nth r 0)) (rest ares)) + rt (cond + (and (= cn "conj") (vec-type? base)) + (mk-vec (reduce join (velem base) rest-ts)) + (and (= cn "into") (vec-type? base) (= 2 n) (vec-type? (nth rest-ts 0))) + (mk-vec (join (velem base) (velem (nth rest-ts 0)))) + :else (call-ret-type fnode))] + [rt (assoc node :args (mapv (fn [r] (nth r 1)) ares))]) + ;; everything else: type args, collect the call (var callee), use the + ;; declared/estimated return type. range produces a numeric vector. + :else + (let [fnode' (if iscall-var fnode (nth (infer fnode tenv) 1)) + ares (mapv (fn [a] (infer a tenv)) args)] + (when iscall-var + (swap! calls-box conj [(var-key fnode) (mapv (fn [r] (nth r 0)) ares)])) + [(cond + (= cn "range") (mk-vec :truthy) + ;; element-returning fn over a typed vector -> the element type + (and cn (contains? elem-fns cn) (> n 0)) + (let [a0 (nth (nth ares 0) 0)] (if (vec-type? a0) (velem a0) :any)) + :else (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) @@ -854,6 +956,11 @@ type call results during the fixpoint." [m] (reset! rtenv-box m)) +(defn set-vtypes! + "Install var VALUE types (a map \"ns/name\" -> type): fn vars are :truthy + (non-nil), def vars carry their inferred init type (jolt-d6u)." + [m] (reset! vtype-box m)) + (defn reset-escapes! [] (reset! escapes-box #{})) (defn collected-escapes [] (vec @escapes-box)) diff --git a/src/jolt/backend.janet b/src/jolt/backend.janet index 96386a6..13c33aa 100644 --- a/src/jolt/backend.janet +++ b/src/jolt/backend.janet @@ -72,16 +72,21 @@ (when (get (ctx :env) :inline?) (def init (norm-node (node :init))) (def meta (node :meta)) - (when (and (= :fn (init :op)) - (not (and meta (or (get meta :redef) (get meta :dynamic))))) - (def arities (vview (init :arities))) - (when (= 1 (length arities)) - (def ar (norm-node (in arities 0))) - (unless (ar :rest) - (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)))))) + (def redefable (and meta (or (get meta :redef) (get meta :dynamic)))) + (cond + redefable nil + (= :fn (init :op)) + (let [arities (vview (init :arities))] + (when (= 1 (length arities)) + (def ar (norm-node (in arities 0))) + (unless (ar :rest) + (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)))) + # a non-fn def: stash so the pass can infer its VALUE type (jolt-d6u), e.g. + # a color table used via rand-nth — its element type flows to lookups. + true (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) @@ -759,12 +764,21 @@ # 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- itype-join [a b] + (cond + (nil? a) b + (nil? b) a + (= a b) a + # compound vector types {:vec ELEM} join element-wise (jolt-d6u) + (and (struct? a) (struct? b) (in a :vec) (in b :vec)) + (struct :vec (itype-join (in a :vec) (in b :vec))) + :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-set-vtypes (and pns (ns-find pns "set-vtypes!"))) (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!"))) @@ -772,28 +786,34 @@ (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 + # gather single-fixed-arity fns AND non-fn defs that stashed a :def IR (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))) - (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) + (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)) - # --- param/return-type fixpoint (chaotic iteration to the LEAST fixpoint) --- + # --- param/return/value-type fixpoint (chaotic iteration to 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 @@ -802,7 +822,8 @@ (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 + ((var-get f-set-vtypes) vtypes) + # type every fn body once under current param types; stash ret + calls (each f fns (def tenv @{}) (def pv (vview (f :params))) @@ -810,6 +831,9 @@ (def res (vview ((var-get f-infer-body) (f :body) tenv))) (put f :tret (in res 0)) (put f :tcalls (in res 2))) + # 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 (def newpt @{}) (each f fns (put newpt (f :key) (array/new-filled (f :np)))) @@ -832,24 +856,51 @@ (put f :pt np) (put f :ret (f :tret)) (when (f :tret) (put nrt (f :key) (f :tret)))) + (each dv defs + (when (not= (dv :tvt) (dv :vt)) (set changed true)) + (put dv :vt (dv :tvt)) + (when (dv :tvt) (put vtypes (dv :key) (dv :tvt)))) (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 + # install the FINAL return + value types so reinfer-def sees them + (def final-rt @{}) + (each f fns (when (f :ret) (put final-rt (f :key) (f :ret)))) + ((var-get f-set-rtenv) final-rt) + ((var-get f-set-vtypes) vtypes) + # --- re-emit the WHOLE unit, callees first (jolt-d6u) ------------------- + # Re-inference alone only rebinds a fn's own var, but the hot path runs + # through callee bodies INLINED / direct-linked into callers at first + # compile. Re-emitting in callee-first (reverse-topological) order makes + # each caller re-embed its now-recompiled callees, and re-infers its body + # (typing locals via return inference) — so the specialization propagates, + # and a call site compiled AFTER this pass (the -e entry) links the whole + # recompiled chain. Every fn is re-emitted, not just those with concrete + # params, so the embedding refreshes even where a fn gained no param type. + (def order @[]) + (def seen @{}) + (defn visit [k] + (unless (get seen k) + (put seen k true) + (def f (get by-key k)) + (when f + (each c (vview (f :tcalls)) (visit (in (vview c) 0))) + (array/push order f)))) + (each f fns (visit (f :key))) + (each f order (put report (f :key) (f :pt)) + (def ptmap @{}) + # escaped fn: its param types are untrustworthy (callers not all visible), + # so re-emit it WITHOUT seeding params (still re-embeds recompiled callees). (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)))))))) + (when (and t (not= t :any)) (put ptmap (in pv i) t)))) + (def def2 ((var-get f-reinfer) (f :def) ptmap)) + (protect (eval (emit-ir ctx def2) (ctx-janet-env ctx)))))) report) (defn ensure-macros-compiled! diff --git a/test/integration/type-infer-phase3-test.janet b/test/integration/type-infer-phase3-test.janet new file mode 100644 index 0000000..89ecfc8 --- /dev/null +++ b/test/integration/type-infer-phase3-test.janet @@ -0,0 +1,41 @@ +# Collection-element types + HOF awareness, Phase 3 (jolt-d6u). A vector carries +# its element type ({:vec ELEM}); a reduce/map/filter closure over it gets that +# element type on its element param. So a lookup inside a reduce closure over a +# vector-of-structs specializes — no hint — WHEN the element type is provable. +(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 3 (jolt-d6u)...") + +(os/setenv "JOLT_DIRECT_LINK" "1") +(def ctx (api/init {:compile? true})) +(api/eval-string ctx "(ns p3)") +(def pns (types/ctx-find-ns ctx "jolt.passes")) +(def reinfer (types/var-get (types/ns-find pns "reinfer-def"))) +# helper: analyze a defn, reinfer with seeded param types, count guards +(defn guards [src ptmap] + (def d (backend/analyze-form ctx (reader/parse-string src))) + (length (string/find-all ":jolt/type" (string/format "%p" (backend/emit-ir ctx (reinfer d ptmap)))))) + +# a reduce closure's element param gets the vector's element type +(def red "(defn f [coll] (reduce (fn [acc h] (+ acc (:r h))) 0 coll))") +(assert (= 0 (guards red @{"coll" {:vec :struct-map}})) "reduce element typed -> bare lookup in closure") +(assert (= 1 (guards red @{"coll" {:vec :any}})) "reduce over vector of unknown -> guard kept") +(assert (= 1 (guards red @{})) "untyped coll -> guard kept") + +# mapv over a vector-of-structs types the closure element too +(def mp "(defn g [coll] (mapv (fn [h] (:r h)) coll))") +(assert (= 0 (guards mp @{"coll" {:vec :struct-map}})) "mapv element typed -> bare lookup") +(assert (= 1 (guards mp @{"coll" {:vec :any}})) "mapv over unknown element -> guard") + +# element type is DERIVED, not just seeded: a vector literal of structs, reduced +(def derived "(defn h2 [] (reduce (fn [acc x] (+ acc (:r x))) 0 [{:r 1 :g 2} {:r 3 :g 4}]))") +(assert (= 0 (guards derived @{})) "vector literal of structs -> element struct -> bare lookup") + +# correctness: the specialized closures compute the same +(assert (= 4 (api/eval-string ctx "((fn [coll] (reduce (fn [acc h] (+ acc (:r h))) 0 coll)) [{:r 1} {:r 3}])")) "reduce value") +(assert (= 4 (api/eval-string ctx "(reduce (fn [acc x] (+ acc (:r x))) 0 [{:r 1 :g 2} {:r 3 :g 4}])")) "derived value") + +(print "Type inference Phase 3 passed!")