Infer monomorphic protocol-method return types (#234)

A protocol method whose impls all return the same record type has a monomorphic
return. collect-pm-rets! scans the unit's (register-(inline-)method ..) forms,
infers each impl fn's return type, and joins them per method; call-ret-type then
types a (method recv ..) call as that record, so a field read off the result
bare-indexes — e.g. (:ray (scatter m ..)) reads off a Ray. A disagreeing impl
joins to :any and keeps the generic path.

run-protoret.ss: a method with all-record impls bare-indexes + unboxes the field
read; a mixed-return method (one impl returns a number) stays generic. make test /
shakesmoke green, selfhost holds, 0 new divergences.

Foundation for auto-typing record values that flow through protocol dispatch. Does
not yet move the ray tracer: its scatter returns ScatterResult-or-nil (Metal
absorbs some rays), and the nil widens the join to :any — typing a nullable return
soundly needs flow-sensitive narrowing (a guarded (some? x) proves non-nil), filed
separately.

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

Makefile

@@ -4,7 +4,7 @@
 # build step. `make test` is the full gate. `make remint` rebuilds the seed after a
 # source change.
 
-.PHONY: test ci values corpus unit smoke buildsmoke selfhost sci certify ffi transient infer wp devirt fieldread numwp fieldnum directlink numeric inline shakesmoke remint
+.PHONY: test ci values corpus unit smoke buildsmoke selfhost sci certify ffi transient infer wp devirt fieldread numwp fieldnum protoret directlink numeric inline shakesmoke remint
 
 # Full gate (dev machine). Includes the self-host byte-fixpoint, which only holds
 # on the same Chez that minted the seed.
@@ -15,7 +15,7 @@ test: selfhost ci
 # lockfile) — it RUNS correctly on any Chez, but `selfhost` rebuilds it and a
 # different Chez version may emit byte-different (gensym/order) output, so the
 # byte-fixpoint is a dev-machine check, not a CI one (jolt-8479).
-ci: values corpus unit smoke buildsmoke sci ffi transient infer wp devirt fieldread numwp fieldnum directlink numeric inline certify
+ci: values corpus unit smoke buildsmoke sci ffi transient infer wp devirt fieldread numwp fieldnum protoret directlink numeric inline certify
 	@echo "OK: CI gates passed"
 
 # Self-host fixpoint: bootstrap.ss rebuild == checked-in seed.
@@ -92,6 +92,12 @@ numwp:
 fieldnum:
 	@chez --script host/chez/run-fieldnum.ss
 
+# Protocol-method return inference: a method whose impls all return the same record
+# type has a monomorphic return, so a (method recv ..) call types as that record and
+# a field read off the result bare-indexes; a disagreeing impl keeps the generic path.
+protoret:
+	@chez --script host/chez/run-protoret.ss
+
 # Direct-linking emission: a closed-world build binds top-level app defs to jv$
 # Scheme bindings and routes app->app calls/refs to them, skipping var-deref +
 # jolt-invoke; ^:dynamic/^:redef and nested defs opt out.

host/chez/run-protoret.ss

@@ -0,0 +1,70 @@
+;; run-protoret.ss — protocol-method return-type inference gate.
+;;
+;; A protocol method whose impls all return the same record type has a monomorphic
+;; return: collect-pm-rets! joins the impl return types, and call-ret-type then types
+;; a (method recv ..) call as that record — so a field read off the result bare-
+;; indexes. This is the ray tracer's (:ray (scatter material ..)): scatter's impls
+;; all return a ScatterResult, so the bounced ray types without a hint.
+;;
+;;   chez --script host/chez/run-protoret.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 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 run-passes (var-deref "jolt.passes" "run-passes"))
+(define emit (var-deref "jolt.backend-scheme" "emit"))
+(define (anode src) (analyze (make-analyze-ctx "user") (jolt-ce-read src)))
+(define (evals src) (jolt-compile-eval (string-append "(do " src ")") "user"))
+(define (sub? s t)(let((n(string-length s))(m(string-length t)))(let loop((i 0))(cond((>(+ i m)n)#f)((string=?(substring s i(+ i m))t)#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)))
+
+(evals "(defrecord R [^double k])")
+(evals "(defprotocol P (m [x]))")
+(evals "(defrecord A [v] P (m [x] (->R 1.0)))")
+(evals "(defrecord B [v] P (m [x] (->R 2.0)))")
+(evals "(defprotocol Q (q [x]))")
+(evals "(defrecord C [v] Q (q [x] (->R 3.0)))")
+(evals "(defrecord D [v] Q (q [x] 7)))")     ; one impl returns a number, not R
+(set-record-shapes! (chez-record-shapes-map))
+(set-protocol-methods! (chez-protocol-methods-map))
+(set-optimize! #t)
+
+;; analyze the impl-registering forms + a consumer; the fixpoint collects the
+;; impl return types. (the analyzed defrecord nodes carry register-inline-method.)
+(define na (anode "(defrecord A [v] P (m [x] (->R 1.0)))"))
+(define nb (anode "(defrecord B [v] P (m [x] (->R 2.0)))"))
+(define nc (anode "(defrecord C [v] Q (q [x] (->R 3.0)))"))
+(define nd (anode "(defrecord D [v] Q (q [x] 7))"))
+(define f  (anode "(def f (fn [a] (* (:k (m a)) 2.0)))"))
+(define g  (anode "(def g (fn [a] (:k (q a))))"))
+(wp-infer! (jolt-vector na nb nc nd f g))
+
+;; m's impls all return R -> (:k (m a)) reads off an R -> bare-index + unbox.
+(define fe (emit (run-passes f (make-analyze-ctx "user"))))
+(check "monomorphic protocol return bare-indexes the field read" (sub? fe "jrec-field-at") #t)
+(check "monomorphic protocol return unboxes the ^double field" (sub? fe "fl") #t)
+
+;; q's impls return R and a number -> joined to non-record -> stays generic (sound).
+(define ge (emit (run-passes g (make-analyze-ctx "user"))))
+(check "mixed-return protocol keeps generic jolt-get" (sub? ge "jolt-get") #t)
+(check "mixed-return protocol does not bare-index" (sub? ge "jrec-field-at") #f)
+
+(if (= fails 0)
+    (begin (printf "protoret gate: ~a/~a passed\n" total total) (exit 0))
+    (begin (printf "protoret gate: ~a/~a passed (~a failed)\n" (- total fails) total fails) (exit 1)))

jolt-core/jolt/passes/types.clj

@@ -5,7 +5,8 @@
   checker. Also the inter-procedural driver API the back end calls to
   propagate param types across a unit / the whole program. Weakly coupled to the
   IR-rewriting passes — shares the const-shape predicates (jolt.passes.fold)."
-  (:require [jolt.passes.fold :refer [scalar-const? kw-callee? get-callee?]]
+  (:require [jolt.ir :refer [reduce-ir-children]]
+            [jolt.passes.fold :refer [scalar-const? kw-callee? get-callee?]]
             [jolt.passes.types.check :refer
              [not-callable? type-name check-invoke register-user-fn!]]
             [jolt.passes.types.lattice :refer
@@ -45,6 +46,11 @@
 (def ^:private last-diags-box (atom []))
 ;; Whether run-inference also checks, and strictly. Set by set-check-mode!.
 (def ^:private check-mode-box (atom {:on false :strict false}))
+;; "Proto/method" -> the join of its impls' return types, so a protocol-method call
+;; types as that record when every impl returns the same one (monomorphic return —
+;; e.g. all Scatter impls return a ScatterResult). Set by collect-pm-rets! before
+;; the fixpoint, read by call-ret-type. A disagreeing impl widens it to :any.
+(def ^:private pm-rets-box (atom {}))
 
 ;; build a per-run env: a snapshot of the installed config plus this run's flags
 ;; and fresh accumulator/guard cells. escapes/user-sigs reference the sweep-level
@@ -104,11 +110,18 @@
                       ;; declared hints so nested records stay typed
                       (record-type-from-entry rs type-depth shapes)
                       (let [r (get (get env :rtenv) (var-key fnode))]
-                        (if r r (let [nm (and (= "clojure.core" (get fnode :ns)) (get fnode :name))]
-                                  (cond (nil? nm) :any
-                                        (contains? num-ret-fns nm) :num
-                                        (contains? vector-ret-fns nm) (mk-vec :any)
-                                        :else :any))))))
+                        (if r r
+                          ;; a protocol-method call types as its impls' joined return
+                          ;; (monomorphic): so (:ray (scatter m ..)) reads off a Ray.
+                          (let [pm (get (get env :protocol-methods) (var-key fnode))
+                                pmr (when pm (get @pm-rets-box (str (nth pm 0) "/" (nth pm 1))))]
+                            (if (and pmr (not= pmr :any))
+                              pmr
+                              (let [nm (and (= "clojure.core" (get fnode :ns)) (get fnode :name))]
+                                (cond (nil? nm) :any
+                                      (contains? num-ret-fns nm) :num
+                                      (contains? vector-ret-fns nm) (mk-vec :any)
+                                      :else :any))))))))
       (= op :host) (let [nm (get fnode :name)]
                      (cond (contains? num-ret-fns nm) :num
                            (contains? vector-ret-fns nm) (mk-vec :any)
@@ -668,6 +681,36 @@
          r (infer body tenv env)]
      [(nth r 0) (nth r 1) @(get env :calls)])))
 
+;; --- protocol-method return types -------------------------------------------
+;; An impl is emitted as (register-(inline-)method TAG "Proto" "method" (fn ...)).
+;; Its fn body's return type is one impl's contribution to the method's return; the
+;; join over every impl is the method's return type (monomorphic when all agree).
+(defn- impl-reg-ret [node]
+  (when (= :invoke (get node :op))
+    (let [f (get node :fn) args (get node :args)]
+      (when (and (= :var (get f :op))
+                 (or (= "register-inline-method" (get f :name))
+                     (= "register-method" (get f :name)))
+                 (= 4 (count args)))
+        (let [proto (get (nth args 1) :val)
+              method (get (nth args 2) :val)
+              fnn (nth args 3)]
+          (when (and (string? proto) (string? method)
+                     (= :fn (get fnn :op)) (seq (get fnn :arities)))
+            [(str proto "/" method)
+             (nth (infer-body (get (first (get fnn :arities)) :body) {}) 0)]))))))
+
+(defn- walk-pm-rets [node acc]
+  (let [kr (impl-reg-ret node)
+        acc (if kr (update acc (nth kr 0) (fn [t] (if t (join t (nth kr 1)) (nth kr 1)))) acc)]
+    (reduce-ir-children (fn [a c] (walk-pm-rets c a)) acc node)))
+
+(defn collect-pm-rets!
+  "Scan the unit's nodes for protocol-method impl registrations and stash each
+  method's joined impl-return type (record-shapes must already be installed)."
+  [nodes]
+  (reset! pm-rets-box (reduce (fn [acc n] (walk-pm-rets n acc)) {} nodes)))
+
 (defn reinfer-def
   "Re-run inference on a stashed :def's fn arity bodies with param types seeded
   (ptmap: param-name -> type), returning the def with annotated bodies. The back
@@ -792,6 +835,7 @@
   record-shapes / protocol-methods must already be installed. Idempotent — resets
   the seed box; called once per build before per-form emit."
   [nodes]
+  (collect-pm-rets! nodes)
   (let [spec (wp-specializable nodes)
         ks (keys spec)]
     (loop [iter 0 ptypes (wp-empty-ptypes spec ks) rets {}]