Scalar-replace short-lived record allocations (jolt-15jq) (#132)

scalar-replace already folds non-escaping const-key map literals
((:k {:k a ..}) -> a, and drops a let-bound map that doesn't escape).
Extend the same fold to record constructors: a (->Rec a b c) is a
positional struct whose declared field order lives in the record-shapes
registry, so a field read on a non-escaping ctor folds to the matching
positional arg and the allocation disappears.

Direct form (:field (->Rec ..)) and the let-bound form both handled,
threaded through run-passes via a per-unit shape registry (new
jolt.host/record-shapes accessor). Soundness: ctor args must be pure
(duplicated/discarded like map vals), arg count must equal the field
count, and only declared-field reads fold — a record answers the virtual
:jolt/deftype key with its type tag and any other key with nil, neither
of which is a positional arg, so those keep the allocation. pure? now
treats a record ctor of pure args as pure, so nested records (a Ray
holding a Vec3) fold bottom-up.

Allocation-bound microbench (non-escaping record built + field-read in a
hot loop): 69.6s -> 2.4s, landing on the no-record arithmetic baseline.
The ray tracer is unchanged — its vec3 results escape (returned/stored
each op), so they genuinely allocate; that's a separate problem.

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

jolt-core/jolt/passes.clj

@@ -13,9 +13,9 @@
   :refer, so jolt.passes stays the only namespace the back end imports.
 
   Portable Clojure: kernel-tier fns + seed primitives only."
-  (:require [jolt.host :refer [inline-enabled?]]
+  (:require [jolt.host :refer [inline-enabled? record-shapes]]
             [jolt.passes.fold :refer [const-fold]]
-            [jolt.passes.inline :refer [inline-node flatten-lets scalar-replace dirty]]
+            [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
                                        set-rtenv! set-vtypes! join-types
@@ -33,7 +33,8 @@
   type is proven. Otherwise (core + bootstrap) just const-fold, as before."
   [node ctx]
   (if (inline-enabled? ctx)
-    (let [opt (loop [i 0 n (const-fold node)]
+    (let [_ (set-rec-shapes! (record-shapes ctx))   ;; record ctor fold (jolt-15jq)
+          opt (loop [i 0 n (const-fold node)]
                 (reset! dirty false)
                 (let [n2 (const-fold (scalar-replace (flatten-lets (inline-node n ctx))))]
                   (if (and @dirty (< i 8))

jolt-core/jolt/passes/inline.clj

@@ -14,6 +14,16 @@
 (def dirty (atom false))   ;; read/reset by the run-passes fixpoint (jolt.passes)
 (defn- mark! [] (reset! dirty true))
 
+;; Record-ctor shape registry ("ns/->Name" -> {:fields (:k ..) :type tag}), fed
+;; per unit by run-passes (set-rec-shapes!) before the fixpoint so scalar-replace
+;; can recognize a (->Rec ..) call and map its positional args to declared fields
+;; — the record analogue of the inline keys a map literal already carries in the
+;; IR (jolt-15jq).
+(def ^:private rec-shapes (atom {}))
+(defn set-rec-shapes!
+  "Install the record-ctor shape registry the record fold consults (jolt-15jq)."
+  [m] (reset! rec-shapes (or m {})))
+
 (def ^:private fresh-counter (atom 0))
 (defn- fresh [base]
   (let [n @fresh-counter]
@@ -244,10 +254,16 @@
       (= op :host) (contains? pure-fns (get f :name))
       :else false)))
 
+;; forward ref: a record ctor (allocating an immutable struct from its args) is
+;; side-effect-free, so pure? treats (->Rec pure-args..) as pure — which lets a
+;; nested record (a Ray holding a Vec3) fold bottom-up (jolt-15jq).
+(declare ctor-shape)
+
 (defn- pure?
   "Conservative: true only for expressions with no side effects that are safe to
-  duplicate or discard. A var/host ref is a pure read; an invoke is pure only
+  duplicate or discard. A var/host ref is a pure read; an invoke is pure for a
-  for a known-pure fn (arithmetic, comparison, keyword lookup, get)."
+  known-pure fn (arithmetic, comparison, keyword lookup, get) or a record
+  constructor (an immutable struct alloc) whose args are themselves pure."
   [node]
   (let [op (get node :op)]
     (cond
@@ -263,7 +279,8 @@
       (= op :vector) (every? pure? (get node :items))
       (= op :set) (every? pure? (get node :items))
       (= op :map) (every? (fn [pr] (and (pure? (nth pr 0)) (pure? (nth pr 1)))) (get node :pairs))
-      (= op :invoke) (and (pure-fn? (get node :fn)) (every? pure? (get node :args)))
+      (= op :invoke) (and (or (pure-fn? (get node :fn)) (ctor-shape node))
+                          (every? pure? (get node :args)))
       :else false)))
 
 (defn- const-key-map? [node]
@@ -410,46 +427,136 @@
       (= op :def) (local-escapes? (get node :init) nm)
       :else true)))
 
-(defn- subst-lookup
+;; --- record constructors as foldable struct sources (jolt-15jq) -------------
-  "Replace every (:k nm)/(get nm :k) in node with the map value at k. The caller
+;; A record ctor (->Rec a b ..) is a positional struct: the registry maps its
-  guarantees (via local-escapes?) that nm is never rebound here and appears only
+;; ctor key ("ns/->Name", exactly how the IR names the call head) to the DECLARED
-  as a lookup subject, so no shadowing logic is needed."
+;; field order. A field read on a non-escaping ctor folds to the matching arg,
-  [node nm mapnode]
+;; just as (:k {:k a ..}) folds to a. Two soundness differences from maps:
+;;   - the ctor's args are duplicated/discarded, so they must be pure (like map
+;;     vals), and the arg count must equal the field count (a positional call);
+;;   - a record answers the virtual :jolt/deftype key with its type tag and any
+;;     other non-field key with nil — neither is a positional arg, so we only
+;;     fold DECLARED-field reads and keep the allocation otherwise.
+
+(defn- ctor-shape
+  "If node is a record-constructor :invoke (its :fn a :var whose ns/name is a
+  registered ctor key, with arg count matching the declared field count), return
+  that record's shape entry; else nil."
+  [node]
+  (if (= :invoke (get node :op))
+    (let [f (get node :fn)]
+      (if (= :var (get f :op))
+        (let [rs (get @rec-shapes (str (get f :ns) "/" (get f :name)))]
+          (if (and rs (= (count (get rs :fields)) (count (get node :args))))
+            rs
+            nil))
+        nil))
+    nil))
+
+(defn- ctor-all-args-pure? [node] (every? pure? (get node :args)))
+
+(defn- field-index
+  "Index of scalar key k in the declared field tuple fields, or nil."
+  [fields k]
+  (let [n (count fields)]
+    (loop [i 0]
+      (if (< i n)
+        (if (= (nth fields i) k) i (recur (inc i)))
+        nil))))
+
+(defn- ctor-val
+  "The positional arg IR at declared field k of record ctor node (shape rs). Only
+  called for a key known to be a field, so the index is always present."
+  [ctor rs k]
+  (nth (get ctor :args) (field-index (get rs :fields) k)))
+
+(defn- collect-keys!
+  "Accumulate (into atom acc) every constant-keyword lookup key applied to local
+  nm in node. The caller has proven (via local-escapes?) that nm appears only as
+  a lookup subject and is never rebound, so a uniform recursion suffices: at a
+  lookup of nm we record the key and stop (its subject is nm itself); elsewhere
+  we recurse into children."
+  [node nm acc]
   (let [k (lookup-key node nm)]
     (if k
-      (map-val mapnode k)
+      (swap! acc conj k)
+      (map-ir-children (fn [c] (collect-keys! c nm acc) c) node))))
+
+(defn- lookups-all-fields?
+  "True if every lookup of nm across nodes uses a declared field in fields — the
+  record-only guard that keeps a :jolt/deftype/unknown-key read (not a positional
+  arg) from being folded to the wrong value."
+  [nodes nm fields]
+  (every? (fn [node]
+            (let [acc (atom #{})]
+              (collect-keys! node nm acc)
+              (every? (fn [k] (field-index fields k)) @acc)))
+          nodes))
+
+(defn- src-val
+  "Field value at k from a foldable struct source — a const-key map (absent key
+  -> nil, struct-map semantics) or a record ctor (k is always a declared field
+  here, guaranteed by lookups-all-fields?)."
+  [src k]
+  (if (= :map (get src :op))
+    (map-val src k)
+    (ctor-val src (ctor-shape src) k)))
+
+(defn- subst-lookup
+  "Replace every (:k nm)/(get nm :k) in node with the source value at k. The
+  caller guarantees (via local-escapes?) that nm is never rebound here and
+  appears only as a lookup subject, so no shadowing logic is needed."
+  [node nm src]
+  (let [k (lookup-key node nm)]
+    (if k
+      (src-val src k)
       ;; the caller's escape check guarantees nm is never rebound below, so we
       ;; recurse uniformly into every child — leaving any lookup of nm
       ;; un-substituted would dangle.
-      (map-ir-children (fn [c] (subst-lookup c nm mapnode)) node))))
+      (map-ir-children (fn [c] (subst-lookup c nm src)) node))))
 
 (defn- fold-kw-literal
-  "(a) (:k {:k a ..}) -> a (siblings pure)."
+  "(a) (:k <source>) -> the value at k. <source> is a const-key pure map
+  ((:k {:k a ..}) -> a) or a pure record ctor ((:k (->Rec a ..)) -> the arg for
+  field k). Siblings are duplicated/discarded, so all must be pure; a record
+  lookup folds only for a declared field."
   [node]
   (let [f (get node :fn) args (get node :args)]
     (if (and (= :const (get f :op)) (keyword? (get f :val)) (= 1 (count args)))
-      (let [m (nth args 0)]
+      (let [m (nth args 0) k (get f :val)]
         (if (and (= :map (get m :op)) (const-key-map? m) (all-vals-pure? m))
-          (do (mark!) (map-val m (get f :val)))
+          (do (mark!) (map-val m k))
-          node))
+          (let [rs (ctor-shape m)]
+            (if (and rs (ctor-all-args-pure? m) (field-index (get rs :fields) k))
+              (do (mark!) (ctor-val m rs k))
+              node))))
       node)))
 
-(defn- elim-let-maps
+(defn- elim-let-structs
-  "(b) Drop the first non-escaping let binding whose init is a pure const-key map
+  "(b) Drop the first non-escaping let binding whose init is a foldable struct
-  literal, substituting its field lookups into the remaining bindings and body.
+  source — a pure const-key map literal or a pure record ctor — substituting its
-  Fixpoint re-runs us for the rest, so one elimination per call keeps it simple."
+  field reads into the remaining bindings and body. Fixpoint re-runs us for the
+  rest, so one elimination per call keeps it simple. For a record every lookup
+  of the binding must hit a declared field, else we keep the allocation."
   [node]
   (let [binds (get node :bindings) n (count binds) body (get node :body)]
     (loop [i 0]
       (if (< i n)
-        (let [b (nth binds i) nm (nth b 0) init (nth b 1)]
+        (let [b (nth binds i) nm (nth b 0) init (nth b 1)
-          (if (and (= :map (get init :op)) (const-key-map? init) (all-vals-pure? init)
+              ismap (and (= :map (get init :op)) (const-key-map? init) (all-vals-pure? init))
+              rs (when (not ismap) (ctor-shape init))
+              isrec (and rs (ctor-all-args-pure? init))]
+          (if (and (or ismap isrec)
                    (not (any-binding-named? (subvec binds (inc i) n) nm))
                    (not (loop [j (inc i)]
                           (if (< j n)
                             (if (local-escapes? (nth (nth binds j) 1) nm) true (recur (inc j)))
                             false)))
-                   (not (local-escapes? body nm)))
+                   (not (local-escapes? body nm))
+                   (or ismap
+                       (lookups-all-fields?
+                         (conj (mapv (fn [bb] (nth bb 1)) (subvec binds (inc i) n)) body)
+                         nm (get rs :fields))))
             (let [head (subvec binds 0 i)
                   tail (mapv (fn [bb] [(nth bb 0) (subst-lookup (nth bb 1) nm init)])
                              (subvec binds (inc i) n))
@@ -467,8 +574,8 @@
   [node]
   (let [op (get node :op)]
     (cond
-      ;; (a) fold (:k {:k a ..}) at invokes, after scalar-replacing children
+      ;; (a) fold (:k <map|ctor>) at invokes, after scalar-replacing children
       (= op :invoke) (fold-kw-literal (map-ir-children scalar-replace node))
-      ;; (b) drop a non-escaping const-key-map let binding, after children
+      ;; (b) drop a non-escaping foldable-struct let binding, after children
-      (= op :let) (elim-let-maps (map-ir-children scalar-replace node))
+      (= op :let) (elim-let-structs (map-ir-children scalar-replace node))
       :else (map-ir-children scalar-replace node))))

src/jolt/host_iface.janet

@@ -239,6 +239,13 @@
   # not just one whose ->Name happens to be interned in the current ns).
   (if (record-hint-ctor-key ctx name) true false))
 
+# The record-ctor shape registry ("ns/->Name" -> {:fields (:k ..) :type tag ..}),
+# or an empty table. Lets scalar-replace (jolt-15jq) recognize a (->Rec ..) call
+# and map its positional args to the declared field order — the record analogue
+# of the inline map-literal keys the map fold already sees in the IR.
+(defn h-record-shapes [ctx]
+  (or (get (ctx :env) :record-shapes) @{}))
+
 (def- exports
   {"form-sym?" h-sym? "form-sym-name" h-sym-name "form-sym-ns" h-sym-ns
    "ref-put!" h-ref-put!
@@ -255,7 +262,7 @@
    "host-intern!" h-intern!
    "inline-enabled?" h-inline-enabled? "inline-ir" h-inline-ir
    "record-type?" h-record-type? "form-position" h-form-position
-   "record-ctor-key" h-record-ctor-key})
+   "record-ctor-key" h-record-ctor-key "record-shapes" h-record-shapes})
 
 (defn install! [ctx]
   (def ns (ctx-find-ns ctx "jolt.host"))

test/integration/scalar-replace-record-test.janet

@@ -0,0 +1,61 @@
+# Scalar-replacement of short-lived RECORD allocations (jolt-15jq). The pass
+# already folds the const-key MAP-literal form ((:k {:k a ..}) -> a and drops a
+# non-escaping let-bound map); this extends it to record CONSTRUCTORS. A record
+# ctor (->Rec a b ..) is a positional struct whose declared field order lives in
+# the record-shapes registry, so a field read on a non-escaping ctor result folds
+# to the corresponding positional arg and the allocation disappears.
+#
+# Probe: count occurrences of the ctor var "->V3" in the analyzed IR. Folded =>
+# the ctor is gone (0). Mirrors type-infer-test's guard-counting harness.
+(import ../../src/jolt/api :as api)
+(import ../../src/jolt/backend :as backend)
+(import ../../src/jolt/reader :as reader)
+
+(print "Scalar-replace of records (jolt-15jq)...")
+
+(os/setenv "JOLT_DIRECT_LINK" "1")
+(def ctx (api/init {:compile? true}))
+(api/eval-string ctx "(ns srr)")
+(api/eval-string ctx "(defrecord V3 [r g b])")
+
+(defn ctors [src]
+  (length (string/find-all "->V3"
+                           (string/format "%p" (backend/analyze-form ctx (reader/parse-string src))))))
+(defn ev [src] (api/eval-string ctx src))
+
+# --- direct form: (:field (->V3 a b c)) -> the positional arg ----------------
+(assert (= 0 (ctors "(fn [] (:r (->V3 1 2 3)))")) "direct record lookup :r -> arg, ctor gone")
+(assert (= 0 (ctors "(fn [] (:g (->V3 1 2 3)))")) "direct record lookup :g -> arg, ctor gone")
+(assert (= 0 (ctors "(fn [] (:b (->V3 1 2 3)))")) "direct record lookup :b -> arg, ctor gone")
+# pure non-constant args fold too (each discarded sibling is pure)
+(assert (= 0 (ctors "(fn [a b] (:r (->V3 (+ a 1) (* b 2) 7)))")) "direct fold with pure arith args")
+
+# --- let form: non-escaping let-bound record, field reads -> args ------------
+(assert (= 0 (ctors "(fn [a b c] (let [v (->V3 a b c)] (+ (:r v) (:g v) (:b v))))")) "let-bound record, all field reads folded")
+(assert (= 0 (ctors "(fn [a b c] (let [v (->V3 a b c)] (:r v)))")) "let-bound record, single field read folded (siblings discarded, pure)")
+
+# --- sound fallbacks: keep the allocation ------------------------------------
+# escaping record (passed to a sink) must NOT be folded
+(assert (>= (ctors "(fn [sink a b c] (let [v (->V3 a b c)] (sink v) (:r v)))") 1) "escaping record keeps the allocation")
+# returning the record escapes it
+(assert (>= (ctors "(fn [a b c] (->V3 a b c))") 1) "returned record keeps the allocation")
+# a non-field key read (:jolt/deftype is a virtual key) -> not folded, keep alloc
+(assert (>= (ctors "(fn [a b c] (let [v (->V3 a b c)] (:jolt/deftype v)))") 1) ":jolt/deftype lookup keeps the allocation")
+
+# --- correctness: folded path evaluates identically -------------------------
+(assert (= 1 (ev "((fn [] (:r (->V3 1 2 3))))")) "direct :r value")
+(assert (= 2 (ev "((fn [] (:g (->V3 1 2 3))))")) "direct :g value")
+(assert (= 3 (ev "((fn [] (:b (->V3 1 2 3))))")) "direct :b value")
+(assert (= 6 (ev "((fn [a b c] (let [v (->V3 a b c)] (+ (:r v) (:g v) (:b v)))) 1 2 3)")) "let-bound sum value")
+(assert (= 10 (ev "((fn [a b] (:r (->V3 (+ a 1) (* b 2) 7))) 9 5)")) "arith-arg direct value")
+# correctness of the kept-allocation fallbacks
+(assert (= 1 (ev "((fn [sink a b c] (let [v (->V3 a b c)] (sink v) (:r v))) (fn [_] nil) 1 2 3)")) "escaping record reads correctly")
+(assert (= "srr.V3" (ev "((fn [a b c] (let [v (->V3 a b c)] (:jolt/deftype v))) 1 2 3)")) ":jolt/deftype reads the type tag")
+
+# --- nested records fold compositionally (bottom-up) -------------------------
+(api/eval-string ctx "(defrecord Ray [orig dir])")
+# (:r (:orig (->Ray (->V3 a b c) d))): inner ctors both fold away
+(assert (= 0 (ctors "(fn [a b c d] (:r (:orig (->Ray (->V3 a b c) d))))")) "nested record reads fold both ctors")
+(assert (= 7 (ev "((fn [a b c d] (:r (:orig (->Ray (->V3 a b c) d)))) 7 8 9 0)")) "nested record fold value")
+
+(print "Scalar-replace of records passed!")