Whole-program param-type inference (closed world) (#226)

Re-derive each app fn's param types from its call sites under --opt, so a
record type flows across fn boundaries: a ctor's return reaches a callee
param, and a typed vector's element reaches a HOF closure's param. The back
end can then bare-index field reads and devirtualize protocol calls at those
sites (it reads the resulting :hint/:devirt annotations; consuming them is
separate work).

This rebuilds the inter-procedural driver the Janet host had — the API
(infer-body/reinfer-def) survived the rehost but nothing drove it, and the
record-shapes/protocol-methods registries were empty stubs.

- records.ss: populate record-shapes (ctor key -> fields/tags/type, resolving
  nested record field tags) and protocol-methods (method var -> [proto method])
  registries at deftype/defprotocol load time; jolt.host accessors materialize
  them.
- passes/types.clj: wp-infer! runs a closed-world fixpoint joining call-site
  arg types into callee params; reinfer-def re-seeds each def at emit. Self-
  recursive calls and fn-level recur are collected so a recursive fn's params
  are constrained by its recursion, not just external callers — else a param
  the recursion widens (e.g. binary-trees check-tree, whose untagged child can
  be nil) would be unsoundly typed non-nil. A fn used in value position keeps
  :any params (callers unknown). Megamorphic sites join to :any.
- build.ss: analyze all app forms and run the fixpoint before per-form emit.
- run-wp.ss: gate (cross-fn propagation, escape soundness, self-recursion).

make test / shakesmoke green, 0 new divergences, selfhost holds.

Co-authored-by: Yogthos <yogthos@gmail.com>

host/chez/build.ss

@@ -170,6 +170,38 @@
 ;; The loop itself is emit-image's ei-emit-ns* (optimize? #t, guard? #f).
 (define (bld-emit-ns ns-name src) (ei-emit-ns* ns-name src #t #f))
 
+;; --- whole-program inference pre-pass ---------------------------------------
+;; Analyze every app form (all namespaces, deps-first) to IR and run the
+;; closed-world param-type fixpoint, so each fn's param types pick up the record
+;; types its callers pass. The per-ns emit below then bare-indexes field reads and
+;; devirtualizes protocol calls at those sites (the back end reads the resulting
+;; :hint/:devirt annotations). Optimized builds only; registries come from the
+;; runtime tables populated as the app loaded.
+(define jolt-wp-infer!             (var-deref "jolt.passes.types" "wp-infer!"))
+(define jolt-wp-set-record-shapes! (var-deref "jolt.passes.types" "set-record-shapes!"))
+(define jolt-wp-set-proto-methods! (var-deref "jolt.passes.types" "set-protocol-methods!"))
+(define jolt-wp-host-record-shapes (var-deref "jolt.host" "record-shapes"))
+(define jolt-wp-host-proto-methods (var-deref "jolt.host" "protocol-methods"))
+
+(define (bld-wp-infer! ordered)
+  (jolt-wp-set-record-shapes! (jolt-wp-host-record-shapes #f))
+  (jolt-wp-set-proto-methods! (jolt-wp-host-proto-methods #f))
+  (let ((nodes '()))
+    (for-each
+      (lambda (nf)
+        (set-chez-ns! (car nf))
+        (let ((src (read-file-string (cdr nf))))
+          (parameterize ((rdr-source-file (cdr nf)))
+            (for-each
+              (lambda (f)
+                (ce-scan-requires! f (car nf))
+                (unless (or (ei-ns-form? f) (ce-macro-form? f))
+                  (guard (e (#t #f))
+                    (set! nodes (cons (jolt-ce-analyze (make-analyze-ctx (car nf)) f) nodes)))))
+              (ei-read-all src)))))
+      ordered)
+    (jolt-wp-infer! (apply jolt-vector (reverse nodes)))))
+
 ;; Strings emitted before each app ns's forms, replaying what the source loader
 ;; does per file: (1) set chez-current-ns so runtime ns-sensitive setup forms
 ;; (defmulti/defmethod resolve their target var through it) land in the right ns;
@@ -326,7 +358,9 @@
                 (set-optimize! (string=? mode "optimized"))
                 (when direct-link?
                   ((var-deref "jolt.backend-scheme" "set-direct-link!") #t)
-                  ((var-deref "jolt.backend-scheme" "direct-link-reset!"))))
+                  ((var-deref "jolt.backend-scheme" "direct-link-reset!")))
+                ;; whole-program param-type fixpoint before per-form emit
+                (when (string=? mode "optimized") (bld-wp-infer! ordered)))
               (lambda ()
                 (if tree-shake?
                     (dce-shake

host/chez/host-contract.ss

@@ -363,7 +363,10 @@
   (hc-sq-lower ctx inner (make-hashtable string-hash string=?)))
 (define (hc-record-type? ctx name) #f)
 (define (hc-record-ctor-key ctx name) jolt-nil)
-(define (hc-record-shapes ctx) (jolt-hash-map))
+;; record + protocol-method shapes for the inference, from the runtime registries
+;; (records.ss) populated as deftype/defprotocol forms load.
+(define (hc-record-shapes ctx) (chez-record-shapes-map))
+(define (hc-protocol-methods ctx) (chez-protocol-methods-map))
 ;; Optimization gate. Off for ordinary runs (open world, redefinition); `jolt
 ;; build` flips it on during app emission for release/optimized modes (closed
 ;; world), turning on the inference + flatten + scalar-replace passes.
@@ -432,6 +435,7 @@
   (def-var! "jolt.host" "record-type?" hc-record-type?)
   (def-var! "jolt.host" "record-ctor-key" hc-record-ctor-key)
   (def-var! "jolt.host" "record-shapes" hc-record-shapes)
+  (def-var! "jolt.host" "protocol-methods" hc-protocol-methods)
   (def-var! "jolt.host" "inline-enabled?" hc-inline-enabled?)
   (def-var! "jolt.host" "inline-ir" hc-inline-ir)
   (def-var! "jolt.host" "stash-inline!" hc-stash-inline!))

host/chez/records.ss

@@ -36,6 +36,71 @@
 ;; in the extension map is distinguished from a genuine miss).
 (define jrec-absent (list 'jrec-absent))
 
+;; --- whole-program inference registries -------------------------------------
+;; Populated at definition/load time (deftype/defrecord and defprotocol forms run
+;; before `jolt build` re-emits), read by the inference driver to seed record and
+;; protocol-method types across fn boundaries. A no-op for the runtime itself; the
+;; tables just accumulate. jolt.host/record-shapes and /protocol-methods (host-
+;; contract.ss) materialize them into the shape jolt.passes.types expects.
+
+;; ctor-key "ns/->Name" -> (vector field-kw-list field-tag-list type-tag).
+;; field-tag-list parallels the fields: "num", a record simple-name string, or #f.
+(define chez-record-shapes-tbl (make-hashtable string-hash string=?))
+;; method var-key "ns/method" -> (cons proto-name method-name).
+(define chez-protocol-methods-tbl (make-hashtable string-hash string=?))
+
+(define (register-record-shape! ctor-key field-kws field-tags type-tag)
+  (hashtable-set! chez-record-shapes-tbl ctor-key
+                  (vector field-kws field-tags type-tag)))
+
+;; simple name of a dotted/slashed string: the segment after the last . or /.
+(define (chez-shape-simple-name s)
+  (let loop ((i (- (string-length s) 1)))
+    (cond ((< i 0) s)
+          ((or (char=? (string-ref s i) #\.) (char=? (string-ref s i) #\/))
+           (substring s (+ i 1) (string-length s)))
+          (else (loop (- i 1))))))
+
+;; resolve a field's declared type tag to what jolt.passes.types wants: "num"
+;; passes through; a record name (simple "Vec3" or qualified "ns.Vec3") resolves
+;; to its ctor-key (so the field reads back as that record); anything else -> nil.
+(define (chez-resolve-field-tag tag by-name)
+  (cond ((or (not tag) (jolt-nil-t? tag)) jolt-nil)
+        ((string=? tag "num") "num")
+        (else (let ((ck (hashtable-ref by-name (chez-shape-simple-name tag) #f)))
+                (if ck ck jolt-nil)))))
+
+;; materialize chez-record-shapes-tbl into "ns/->Name" -> {:fields :tags :type},
+;; the shape record-type-from-entry consumes.
+(define (chez-record-shapes-map)
+  (let ((by-name (make-hashtable string-hash string=?))
+        (kw-fields (keyword #f "fields")) (kw-tags (keyword #f "tags")) (kw-type (keyword #f "type"))
+        (out (jolt-hash-map)))
+    ;; index simple record name (from the type tag "ns.Name") -> ctor-key for
+    ;; nested-field-tag resolution.
+    (let-values (((ks vs) (hashtable-entries chez-record-shapes-tbl)))
+      (vector-for-each
+        (lambda (k v) (hashtable-set! by-name (chez-shape-simple-name (vector-ref v 2)) k)) ks vs)
+      (vector-for-each
+        (lambda (k v)
+          (let* ((fields (vector-ref v 0)) (tags (vector-ref v 1)) (type-tag (vector-ref v 2))
+                 (rtags (map (lambda (t) (chez-resolve-field-tag t by-name)) tags)))
+            (set! out (jolt-assoc out k
+                                  (jolt-hash-map kw-fields (apply jolt-vector fields)
+                                                 kw-tags   (apply jolt-vector rtags)
+                                                 kw-type   type-tag)))))
+        ks vs))
+    out))
+
+;; materialize chez-protocol-methods-tbl into "ns/method" -> [proto method].
+(define (chez-protocol-methods-map)
+  (let ((out (jolt-hash-map)))
+    (let-values (((ks vs) (hashtable-entries chez-protocol-methods-tbl)))
+      (vector-for-each
+        (lambda (k v) (set! out (jolt-assoc out k (jolt-vector (car v) (cdr v)))))
+        ks vs))
+    out))
+
 ;; index of a declared field key, or #f (only an interned keyword can be one).
 (define (jrec-field-index r k) (hashtable-ref (jrdesc-index (jrec-desc r)) k #f))
 ;; a vector-copy that doesn't depend on the optional rnrs vector-copy being present.
@@ -351,9 +416,10 @@
 ;; ---- the native that handles the analyzer/overlay call ----------------------
 ;; make-deftype-ctor: (name-sym field-kws field-tags field-muts) -> ctor closure.
 ;; The tag is baked at definition time in the type's ns (chez-current-ns).
-(define (make-deftype-ctor name-sym field-kws . _ignored)
+(define (make-deftype-ctor name-sym field-kws . rest-args)
   (let* ((tag (string-append (chez-current-ns) "." (symbol-t-name name-sym)))
          (kws (seq->list field-kws))
+         (field-tags (if (pair? rest-args) (seq->list (car rest-args)) '()))
          (desc (make-jrdesc tag kws))
          (nf (length kws))
          (ctor (lambda args
@@ -368,6 +434,10 @@
     ;; even when the runtime current ns is the caller's, not the defining ns
     ;; (host-new checks class-ctors-tbl before the current-ns var fallback).
     (register-class-ctor! (symbol-t-name name-sym) ctor)
+    ;; record the shape for whole-program inference, keyed by the positional
+    ;; ctor var "ns/->Name" the analyzer resolves a (->Name …) call to.
+    (register-record-shape! (string-append (chez-current-ns) "/->" (symbol-t-name name-sym))
+                            kws field-tags tag)
     ctor))
 
 ;; make-protocol: a protocol value the overlay reads via (get p :name)/(get p :methods).
@@ -376,10 +446,18 @@
                  (keyword #f "name") (jolt-symbol jolt-nil name-str)
                  (keyword #f "methods") methods))
 
-;; register-protocol-methods!: intentional no-op. Chez dispatches a protocol method
-;; by the receiver's type tag at call time, so there is no method table to register;
-;; this binding exists only because defprotocol-emitted code calls it.
-(define (register-protocol-methods! proto-name method-names) jolt-nil)
+;; register-protocol-methods!: record each method's var-key -> [proto method] for
+;; the inference driver (devirtualization). Dispatch itself is by the receiver's
+;; type tag at call time, so this table is read only by `jolt build` inference.
+;; Called by defprotocol-emitted code in the protocol's ns.
+(define (register-protocol-methods! proto-name method-names)
+  (let ((ns (chez-current-ns)))
+    (for-each (lambda (mn)
+                (let ((m (if (symbol-t? mn) (symbol-t-name mn) mn)))
+                  (hashtable-set! chez-protocol-methods-tbl
+                                  (string-append ns "/" m) (cons proto-name m))))
+              (seq->list method-names)))
+  jolt-nil)
 
 ;; register-method: extend-type/extend register an impl. Host type names keep a
 ;; bare canonical tag; record names qualify to the current ns.

host/chez/run-wp.ss

@@ -0,0 +1,96 @@
+;; run-wp.ss — whole-program param-type fixpoint gate (jolt.passes.types/wp-infer!).
+;;
+;; run-infer.ss drives the per-form inference; this drives the inter-procedural
+;; driver: analyze a multi-def unit, run wp-infer!, and assert that a record type
+;; flows across fn boundaries — a callee's param picks up its caller's ctor return
+;; type, so a field read off it is marked for the bare-index back-end path.
+;;
+;;   chez --script host/chez/run-wp.ss
+(import (chezscheme))
+(load "host/chez/rt.ss")
+(set-chez-ns! "clojure.core")
+(load "host/chez/seed/prelude.ss")
+(load "host/chez/post-prelude.ss")
+(set-chez-ns! "user")
+(load "host/chez/host-contract.ss")
+(load "host/chez/seed/image.ss")
+(load "host/chez/compile-eval.ss")
+
+(define analyze            (var-deref "jolt.analyzer" "analyze"))
+(define run-inference      (var-deref "jolt.passes.types" "run-inference"))
+(define set-record-shapes! (var-deref "jolt.passes.types" "set-record-shapes!"))
+(define set-protocol-methods! (var-deref "jolt.passes.types" "set-protocol-methods!"))
+(define wp-infer!          (var-deref "jolt.passes.types" "wp-infer!"))
+(define param-seeds-for    (var-deref "jolt.passes.types" "param-seeds-for"))
+(define reinfer-def        (var-deref "jolt.passes.types" "reinfer-def"))
+(define pr-str             (var-deref "clojure.core" "pr-str"))
+
+(define (anode src) (analyze (make-analyze-ctx "user") (jolt-ce-read src)))
+(define (contains-sub? s sub)
+  (let ((n (string-length s)) (m (string-length sub)))
+    (let loop ((i 0))
+      (cond ((> (+ i m) n) #f)
+            ((string=? (substring s i (+ i m)) sub) #t)
+            (else (loop (+ i 1)))))))
+
+(define fails 0) (define total 0)
+(define (check label actual expected)
+  (set! total (+ total 1))
+  (unless (equal? actual expected)
+    (set! fails (+ fails 1))
+    (printf "  FAIL ~a: got ~s expected ~s\n" label actual expected)))
+
+;; Node record shape (left/right untagged), like binary-trees.
+(set-record-shapes!
+  (jolt-hash-map "user/->Node"
+                 (jolt-hash-map (keyword #f "fields") (jolt-vector (keyword #f "left") (keyword #f "right"))
+                                (keyword #f "tags")   (jolt-vector jolt-nil jolt-nil)
+                                (keyword #f "type")   "user.Node")))
+(set-protocol-methods! (jolt-hash-map))
+
+;; a 3-def unit: make-tree returns ->Node, run calls check-tree with a make-tree
+;; result, so check-tree's `node` param must be inferred as a Node.
+(define mt (anode "(def make-tree (fn [depth] (if (zero? depth) (->Node nil nil) (->Node (make-tree (dec depth)) (make-tree (dec depth))))))"))
+(define ct (anode "(def check-tree (fn [node] (:left node)))"))
+(define rn (anode "(def run (fn [d] (check-tree (make-tree d))))"))
+
+(wp-infer! (jolt-vector mt ct rn))
+
+;; check-tree's param `node` should be seeded with a struct carrying the Node type
+(define seed (param-seeds-for "user/check-tree"))
+(check "check-tree has a param seed" (jolt-truthy? seed) #t)
+(when (jolt-truthy? seed)
+  (check "node seeded as user.Node struct"
+         (contains-sub? (pr-str seed) "user.Node") #t))
+
+;; reinfer-def then must mark the (:left node) read site for the bare-index path
+(define marked (reinfer-def ct seed))
+(check "read site marked :hint :struct" (contains-sub? (pr-str marked) ":hint :struct") #t)
+
+;; a fn used only via value position (escape) must NOT be specialized — unknown
+;; callers make a concrete seed unsound.
+(define ev (anode "(def use-it (fn [f] (f 1)))"))
+(define ec (anode "(def caller (fn [] (use-it check-tree)))"))  ; check-tree escapes
+(wp-infer! (jolt-vector mt ct rn ev ec))
+(check "escaped fn keeps no param seed" (jolt-truthy? (param-seeds-for "user/check-tree")) #f)
+
+;; a self-recursive fn that recurses on a NILABLE field (an untagged record field
+;; is :any, so the child can be nil) must NOT be specialized — the recursion can
+;; pass nil, so typing the param as a non-nil record would be unsound.
+(define ctr (anode "(def walk (fn [node] (let [l (:left node)] (if (nil? l) 1 (walk l)))))"))
+(define rnr (anode "(def run2 (fn [d] (walk (make-tree d))))"))
+(wp-infer! (jolt-vector mt ctr rnr))
+(check "self-recursive nilable param not specialized"
+       (jolt-truthy? (param-seeds-for "user/walk")) #f)
+
+;; a self-recursive fn that recurses passing the SAME record type (make-tree always
+;; returns a Node) is still safe to specialize — the recursion preserves the type.
+(define mtt (anode "(def grow (fn [n acc] (if (zero? n) acc (grow (dec n) (->Node acc acc)))))"))
+(define gcl (anode "(def gcaller (fn [] (grow 5 (->Node nil nil))))"))
+(wp-infer! (jolt-vector mtt gcl))
+(check "self-recursive same-type param keeps its seed"
+       (jolt-truthy? (param-seeds-for "user/grow")) #t)
+
+(if (= fails 0)
+    (begin (printf "wp gate: ~a/~a passed\n" total total) (exit 0))
+    (begin (printf "wp gate: ~a/~a passed (~a failed)\n" (- total fails) total fails) (exit 1)))