Merge pull request #118 from jolt-lang/refactor-phase3a-ir-walk

Refactor phase 3a: one map-ir-children combinator for IR rewrite walks (jolt-26dm)

jolt-core/jolt/ir.clj

@@ -59,3 +59,56 @@
 (defn throw-node [expr] {:op :throw :expr expr})
 
 (defn op [node] (:op node))
+
+;; ---------------------------------------------------------------------------
+;; Structural recursion over IR child nodes (jolt-26dm / phase 3a).
+;;
+;; A tree-rewriting pass recurses into each op's child NODE positions and
+;; rebuilds the node; this combinator does that one place, so the per-op child
+;; layout is single-sourced and adding an op is a one-site change here (was: an
+;; edit to every walk). `(map-ir-children f node)` returns node with f applied to
+;; each child IR node — re-applied per element for seq positions (:args/:items/
+;; :statements), per value for :map pairs, per init for :let/:loop bindings, and
+;; per arity :body for :fn. Non-node positions (binding NAMES, fn :params/:rest,
+;; the :op tag, :ns/:name/:val) are left intact. Leaf ops and any op with no
+;; child nodes pass through unchanged, so walks built on this are TOTAL over the
+;; op set (an unknown op recurses nowhere rather than being silently dropped).
+;;
+;; Uses cond/=/get only — same constructs as the passes that consume it, so it
+;; loads at the same compiler tier with no new macro dependency.
+(defn map-ir-children [f node]
+  (let [op (get node :op)]
+    (cond
+      (= op :if)     (assoc node :test (f (get node :test))
+                                 :then (f (get node :then))
+                                 :else (f (get node :else)))
+      (= op :do)     (assoc node :statements (mapv f (get node :statements))
+                                 :ret (f (get node :ret)))
+      (= op :throw)  (assoc node :expr (f (get node :expr)))
+      (= op :invoke) (assoc node :fn (f (get node :fn))
+                                 :args (mapv f (get node :args)))
+      (= op :vector) (assoc node :items (mapv f (get node :items)))
+      (= op :set)    (assoc node :items (mapv f (get node :items)))
+      (= op :map)    (assoc node :pairs (mapv (fn [pr] [(f (nth pr 0)) (f (nth pr 1))])
+                                              (get node :pairs)))
+      (= op :let)    (assoc node :bindings (mapv (fn [b] [(nth b 0) (f (nth b 1))])
+                                                 (get node :bindings))
+                                 :body (f (get node :body)))
+      (= op :loop)   (assoc node :bindings (mapv (fn [b] [(nth b 0) (f (nth b 1))])
+                                                 (get node :bindings))
+                                 :body (f (get node :body)))
+      (= op :recur)  (assoc node :args (mapv f (get node :args)))
+      (= op :fn)     (assoc node :arities (mapv (fn [a] (assoc a :body (f (get a :body))))
+                                                (get node :arities)))
+      (= op :def)    (assoc node :init (f (get node :init)))
+      ;; :catch-body / :finally are optional; recurse them only when PRESENT.
+      ;; Assoc'ing them nil-when-absent would turn the node into a phm (jolt's
+      ;; nil-valued-key representation) and force backend densification — so we
+      ;; preserve the node's shape and never introduce a nil key.
+      (= op :try)
+      (let [n (assoc node :body (f (get node :body)))
+            n (if (get node :catch-body) (assoc n :catch-body (f (get node :catch-body))) n)
+            n (if (get node :finally) (assoc n :finally (f (get node :finally))) n)]
+        n)
+      ;; :const :local :var :host :the-var :rt :quote — no child nodes
+      :else node)))

jolt-core/jolt/passes/fold.clj

@@ -3,7 +3,8 @@
   Bottom-up numeric folding + dead-branch removal, total over node :ops (unknown
   ops pass through with folded children). Portable Clojure: kernel-tier fns +
   seed primitives only — it loads with the compiler namespaces, before the later
-  core tiers.")
+  core tiers."
+  (:require [jolt.ir :refer [map-ir-children]]))
 
 ;; Folding computes with THE ACTUAL jolt fns, so a folded result matches what
 ;; the unfolded code would produce at runtime by construction. Conservative:
@@ -38,14 +39,15 @@
   (let [op (get node :op)]
     (cond
       (= op :invoke)
-      (let [f (const-fold (get node :fn))
+      ;; fold children first, then this call if the fn is foldable over consts
-            args (mapv const-fold (get node :args))
+      (let [n (map-ir-children const-fold node)
-            ff (fold-fn f)
+            ff (fold-fn (get n :fn))
+            args (get n :args)
             folded (when (and ff (pos? (count args)) (every? const-num? args))
                      (try
                        {:op :const :val (apply ff (mapv (fn [a] (get a :val)) args))}
                        (catch Exception e nil)))]
-        (or folded (assoc node :fn f :args args)))
+        (or folded n))
 
       (= op :if)
       (let [t (const-fold (get node :test))]
@@ -59,41 +61,9 @@
                  :then (const-fold (get node :then))
                  :else (const-fold (get node :else)))))
 
-      (= op :do)
+      ;; every other op: fold each child (let/loop bindings are [name init]
-      (assoc node
+      ;; pairs, handled by the combinator)
-             :statements (mapv const-fold (get node :statements))
+      :else (map-ir-children const-fold node))))
-             :ret (const-fold (get node :ret)))
-
-      ;; let/loop bindings are [name-string init-ir] PAIRS (see
-      ;; analyzer/analyze-bindings), not maps.
-      (= op :let)
-      (assoc node
-             :bindings (mapv (fn [b] [(nth b 0) (const-fold (nth b 1))])
-                             (get node :bindings))
-             :body (const-fold (get node :body)))
-
-      (= op :loop)
-      (assoc node
-             :bindings (mapv (fn [b] [(nth b 0) (const-fold (nth b 1))])
-                             (get node :bindings))
-             :body (const-fold (get node :body)))
-
-      (= op :recur)
-      (assoc node :args (mapv const-fold (get node :args)))
-
-      (= op :fn)
-      (assoc node
-             :arities (mapv (fn [a] (assoc a :body (const-fold (get a :body))))
-                            (get node :arities)))
-
-      (= op :def)    (assoc node :init (const-fold (get node :init)))
-      (= op :throw)  (assoc node :expr (const-fold (get node :expr)))
-      (= op :vector) (assoc node :items (mapv const-fold (get node :items)))
-      (= op :set)    (assoc node :items (mapv const-fold (get node :items)))
-      (= op :map)    (assoc node :pairs (mapv (fn [pr] (mapv const-fold pr)) (get node :pairs)))
-
-      ;; leaves and anything this pass doesn't know: unchanged
-      :else node)))
 
 ;; A const node whose value is a scalar literal (kw/str/num/bool). Shared by the
 ;; scalar-replace pass (jolt.passes.inline) and the collection-type inference

jolt-core/jolt/passes/inline.clj

@@ -4,6 +4,7 @@
   share the alpha-rename invariant (every spliced binder is made globally fresh)
   and the `dirty` fixpoint flag. Portable Clojure (compiler-tier)."
   (:require [jolt.host :refer [inline-ir]]
+            [jolt.ir :refer [map-ir-children]]
             [jolt.passes.fold :refer [scalar-const?]]))
 
 ;; ---------------------------------------------------------------------------
@@ -85,22 +86,9 @@
                           (assoc r :hint (get node :hint))
                           r)
                         node))
-      (= op :if) (assoc node
+      ;; :let alpha-renames each binder to a fresh name, threading the extended
-                        :test (subst (get node :test) env)
+      ;; env left-to-right — sequential scope the uniform combinator can't model,
-                        :then (subst (get node :then) env)
+      ;; so it stays explicit.
-                        :else (subst (get node :else) env))
-      (= op :do) (assoc node
-                        :statements (mapv (fn [s] (subst s env)) (get node :statements))
-                        :ret (subst (get node :ret) env))
-      (= op :throw) (assoc node :expr (subst (get node :expr) env))
-      (= op :invoke) (assoc node
-                            :fn (subst (get node :fn) env)
-                            :args (mapv (fn [a] (subst a env)) (get node :args)))
-      (= op :vector) (assoc node :items (mapv (fn [x] (subst x env)) (get node :items)))
-      (= op :set) (assoc node :items (mapv (fn [x] (subst x env)) (get node :items)))
-      (= op :map) (assoc node :pairs (mapv (fn [pr] [(subst (nth pr 0) env)
-                                                     (subst (nth pr 1) env)])
-                                           (get node :pairs)))
       (= op :let)
       (let [res (reduce (fn [acc b]
                           (let [e (nth acc 0)
@@ -112,8 +100,10 @@
                         [env []]
                         (get node :bindings))]
         (assoc node :bindings (nth res 1) :body (subst (get node :body) (nth res 0))))
-      ;; :const :var :host :the-var :quote — no locals to substitute
+      ;; every other op substitutes env uniformly into its children. Inline
-      :else node)))
+      ;; bodies only contain safe ops (see safe-op?), so loop/recur/fn/def/try
+      ;; never reach here; the combinator handles them harmlessly regardless.
+      :else (map-ir-children (fn [c] (subst c env)) node))))
 
 (defn- trivial-arg? [n]
   ;; safe to substitute directly (immutable, free to duplicate): a local read or
@@ -201,41 +191,10 @@
 (defn inline-node
   "Bottom-up: inline children first, then attempt to inline this node."
   [node ctx]
-  (let [op (get node :op)]
+  (if (= :invoke (get node :op))
-    (cond
+    ;; inline children first, then attempt to splice this call
-      (= op :invoke)
+    (try-inline (map-ir-children (fn [c] (inline-node c ctx)) node) ctx)
-      (try-inline (assoc node
+    (map-ir-children (fn [c] (inline-node c ctx)) node)))
-                         :fn (inline-node (get node :fn) ctx)
-                         :args (mapv (fn [a] (inline-node a ctx)) (get node :args)))
-                  ctx)
-      (= op :if) (assoc node
-                        :test (inline-node (get node :test) ctx)
-                        :then (inline-node (get node :then) ctx)
-                        :else (inline-node (get node :else) ctx))
-      (= op :do) (assoc node
-                        :statements (mapv (fn [s] (inline-node s ctx)) (get node :statements))
-                        :ret (inline-node (get node :ret) ctx))
-      (= op :let) (assoc node
-                         :bindings (mapv (fn [b] [(nth b 0) (inline-node (nth b 1) ctx)]) (get node :bindings))
-                         :body (inline-node (get node :body) ctx))
-      (= op :loop) (assoc node
-                          :bindings (mapv (fn [b] [(nth b 0) (inline-node (nth b 1) ctx)]) (get node :bindings))
-                          :body (inline-node (get node :body) ctx))
-      (= op :recur) (assoc node :args (mapv (fn [a] (inline-node a ctx)) (get node :args)))
-      (= op :fn) (assoc node :arities (mapv (fn [a] (assoc a :body (inline-node (get a :body) ctx)))
-                                            (get node :arities)))
-      (= op :def) (assoc node :init (inline-node (get node :init) ctx))
-      (= op :throw) (assoc node :expr (inline-node (get node :expr) ctx))
-      (= op :vector) (assoc node :items (mapv (fn [x] (inline-node x ctx)) (get node :items)))
-      (= op :set) (assoc node :items (mapv (fn [x] (inline-node x ctx)) (get node :items)))
-      (= op :map) (assoc node :pairs (mapv (fn [pr] [(inline-node (nth pr 0) ctx)
-                                                     (inline-node (nth pr 1) ctx)])
-                                           (get node :pairs)))
-      (= op :try) (assoc node
-                         :body (inline-node (get node :body) ctx)
-                         :catch-body (when (get node :catch-body) (inline-node (get node :catch-body) ctx))
-                         :finally (when (get node :finally) (inline-node (get node :finally) ctx)))
-      :else node)))
 
 ;; ---------------------------------------------------------------------------
 ;; flatten-lets: (let [a (let [b X] Y) ..] body) -> (let [b X a Y ..] body).
@@ -256,40 +215,11 @@
           binds))
 
 (defn flatten-lets [node]
-  (let [op (get node :op)]
+  (if (= :let (get node :op))
-    (cond
+    ;; flatten children first, then hoist any let-valued binding inits
-      (= op :let) (assoc node
+    (let [n (map-ir-children flatten-lets node)]
-                         :bindings (flatten-let-bindings
+      (assoc n :bindings (flatten-let-bindings (get n :bindings))))
-                                    (mapv (fn [b] [(nth b 0) (flatten-lets (nth b 1))]) (get node :bindings)))
+    (map-ir-children flatten-lets node)))
-                         :body (flatten-lets (get node :body)))
-      (= op :if) (assoc node
-                        :test (flatten-lets (get node :test))
-                        :then (flatten-lets (get node :then))
-                        :else (flatten-lets (get node :else)))
-      (= op :do) (assoc node
-                        :statements (mapv flatten-lets (get node :statements))
-                        :ret (flatten-lets (get node :ret)))
-      (= op :throw) (assoc node :expr (flatten-lets (get node :expr)))
-      (= op :invoke) (assoc node
-                            :fn (flatten-lets (get node :fn))
-                            :args (mapv flatten-lets (get node :args)))
-      (= op :vector) (assoc node :items (mapv flatten-lets (get node :items)))
-      (= op :set) (assoc node :items (mapv flatten-lets (get node :items)))
-      (= op :map) (assoc node :pairs (mapv (fn [pr] [(flatten-lets (nth pr 0))
-                                                     (flatten-lets (nth pr 1))])
-                                           (get node :pairs)))
-      (= op :loop) (assoc node
-                          :bindings (mapv (fn [b] [(nth b 0) (flatten-lets (nth b 1))]) (get node :bindings))
-                          :body (flatten-lets (get node :body)))
-      (= op :recur) (assoc node :args (mapv flatten-lets (get node :args)))
-      (= op :fn) (assoc node :arities (mapv (fn [a] (assoc a :body (flatten-lets (get a :body))))
-                                            (get node :arities)))
-      (= op :def) (assoc node :init (flatten-lets (get node :init)))
-      (= op :try) (assoc node
-                         :body (flatten-lets (get node :body))
-                         :catch-body (when (get node :catch-body) (flatten-lets (get node :catch-body)))
-                         :finally (when (get node :finally) (flatten-lets (get node :finally))))
-      :else node)))
 
 ;; ---------------------------------------------------------------------------
 ;; scalar-replace (AOT escape analysis). A map allocation whose ONLY use is
@@ -485,42 +415,13 @@
   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 mapnode]
-  (let [op (get node :op)
+  (let [k (lookup-key node nm)]
-        k (lookup-key node nm)]
+    (if k
-    (cond
+      (map-val mapnode k)
-      k (map-val mapnode k)
+      ;; the caller's escape check guarantees nm is never rebound below, so we
-      (= op :if) (assoc node
+      ;; recurse uniformly into every child — leaving any lookup of nm
-                        :test (subst-lookup (get node :test) nm mapnode)
+      ;; un-substituted would dangle.
-                        :then (subst-lookup (get node :then) nm mapnode)
+      (map-ir-children (fn [c] (subst-lookup c nm mapnode)) node))))
-                        :else (subst-lookup (get node :else) nm mapnode))
-      (= op :do) (assoc node
-                        :statements (mapv (fn [s] (subst-lookup s nm mapnode)) (get node :statements))
-                        :ret (subst-lookup (get node :ret) nm mapnode))
-      (= op :throw) (assoc node :expr (subst-lookup (get node :expr) nm mapnode))
-      (= op :invoke) (assoc node
-                            :fn (subst-lookup (get node :fn) nm mapnode)
-                            :args (mapv (fn [a] (subst-lookup a nm mapnode)) (get node :args)))
-      (= op :vector) (assoc node :items (mapv (fn [x] (subst-lookup x nm mapnode)) (get node :items)))
-      (= op :set) (assoc node :items (mapv (fn [x] (subst-lookup x nm mapnode)) (get node :items)))
-      (= op :map) (assoc node :pairs (mapv (fn [pr] [(subst-lookup (nth pr 0) nm mapnode)
-                                                     (subst-lookup (nth pr 1) nm mapnode)])
-                                           (get node :pairs)))
-      (= op :let) (assoc node
-                         :bindings (mapv (fn [b] [(nth b 0) (subst-lookup (nth b 1) nm mapnode)]) (get node :bindings))
-                         :body (subst-lookup (get node :body) nm mapnode))
-      ;; the caller's escape check guarantees nm is not rebound in these, so we
-      ;; recurse uniformly — leaving any lookup of nm un-substituted would dangle.
-      (= op :recur) (assoc node :args (mapv (fn [a] (subst-lookup a nm mapnode)) (get node :args)))
-      (= op :loop) (assoc node
-                          :bindings (mapv (fn [b] [(nth b 0) (subst-lookup (nth b 1) nm mapnode)]) (get node :bindings))
-                          :body (subst-lookup (get node :body) nm mapnode))
-      (= op :fn) (assoc node :arities (mapv (fn [a] (assoc a :body (subst-lookup (get a :body) nm mapnode)))
-                                            (get node :arities)))
-      (= op :try) (assoc node
-                         :body (subst-lookup (get node :body) nm mapnode)
-                         :catch-body (when (get node :catch-body) (subst-lookup (get node :catch-body) nm mapnode))
-                         :finally (when (get node :finally) (subst-lookup (get node :finally) nm mapnode)))
-      :else node)))
 
 (defn- fold-kw-literal
   "(a) (:k {:k a ..}) -> a (siblings pure)."
@@ -566,36 +467,8 @@
   [node]
   (let [op (get node :op)]
     (cond
-      (= op :invoke)
+      ;; (a) fold (:k {:k a ..}) at invokes, after scalar-replacing children
-      (fold-kw-literal (assoc node
+      (= op :invoke) (fold-kw-literal (map-ir-children scalar-replace node))
-                              :fn (scalar-replace (get node :fn))
+      ;; (b) drop a non-escaping const-key-map let binding, after children
-                              :args (mapv scalar-replace (get node :args))))
+      (= op :let) (elim-let-maps (map-ir-children scalar-replace node))
-      (= op :let)
+      :else (map-ir-children scalar-replace node))))
-      (elim-let-maps (assoc node
-                            :bindings (mapv (fn [b] [(nth b 0) (scalar-replace (nth b 1))]) (get node :bindings))
-                            :body (scalar-replace (get node :body))))
-      (= op :if) (assoc node
-                        :test (scalar-replace (get node :test))
-                        :then (scalar-replace (get node :then))
-                        :else (scalar-replace (get node :else)))
-      (= op :do) (assoc node
-                        :statements (mapv scalar-replace (get node :statements))
-                        :ret (scalar-replace (get node :ret)))
-      (= op :throw) (assoc node :expr (scalar-replace (get node :expr)))
-      (= op :vector) (assoc node :items (mapv scalar-replace (get node :items)))
-      (= op :set) (assoc node :items (mapv scalar-replace (get node :items)))
-      (= op :map) (assoc node :pairs (mapv (fn [pr] [(scalar-replace (nth pr 0))
-                                                     (scalar-replace (nth pr 1))])
-                                           (get node :pairs)))
-      (= op :loop) (assoc node
-                          :bindings (mapv (fn [b] [(nth b 0) (scalar-replace (nth b 1))]) (get node :bindings))
-                          :body (scalar-replace (get node :body)))
-      (= op :recur) (assoc node :args (mapv scalar-replace (get node :args)))
-      (= op :fn) (assoc node :arities (mapv (fn [a] (assoc a :body (scalar-replace (get a :body))))
-                                            (get node :arities)))
-      (= op :def) (assoc node :init (scalar-replace (get node :init)))
-      (= op :try) (assoc node
-                         :body (scalar-replace (get node :body))
-                         :catch-body (when (get node :catch-body) (scalar-replace (get node :catch-body)))
-                         :finally (when (get node :finally) (scalar-replace (get node :finally))))
-      :else node)))

test/integration/ir-passes-test.janet

@@ -39,4 +39,13 @@
 # broader matrix; this pins a couple end-to-end)
 (assert (= 6 (api/eval-string ctx "(+ 1 2 3)")) "folded eval")
 (assert (= :yes (api/eval-string ctx "(if (< 1 2) :yes :no)")) "folded if eval")
+
+# Folding reaches into every op's children via map-ir-children (phase 3a) — it is
+# now total, so a constant nested in a fn arity / loop / try body folds too
+# (const-fold previously passed :try through unfolded). Sound: folding preserves
+# runtime results regardless of position, so these eval end-to-end (3-mode
+# conformance covers the broader matrix).
+(assert (= 3 ((api/eval-string ctx "(fn [x] (+ 1 2))") 0)) "folded fn arity body eval")
+(assert (= 10 (api/eval-string ctx "(loop [i 0] (if (< i (* 2 5)) (recur (inc i)) i))")) "folded loop eval")
+(assert (= 5 (api/eval-string ctx "(try (+ 2 3) (catch Throwable e 0))")) "folded try eval")
 (print "IR passes passed!")