Merge pull request #64 from jolt-lang/ir-passes

compiler: IR pass pipeline + constant folding (jolt-2om)

jolt-core/jolt/passes.clj

@@ -0,0 +1,99 @@
+(ns jolt.passes
+  "IR optimization passes (nanopass-lite, jolt-2om). Each pass is a pure
+  IR -> IR rewrite, total over node :ops (unknown ops pass through with
+  folded children, so adding a node kind can't silently break a pass), run
+  in a fixed order by run-passes between the analyzer and the back end.
+  Portable Clojure: same constraint as jolt.analyzer — kernel-tier fns +
+  seed primitives only (it loads with the compiler namespaces).")
+
+;; Folding computes with THE ACTUAL jolt fns, so a folded result matches what
+;; the unfolded code would produce at runtime by construction. Conservative:
+;; numbers only, the op table only names pure numeric fns, and any throw
+;; during folding (e.g. (mod x 0)) leaves the node alone for runtime.
+(def ^:private foldable
+  ;; SEED fns only: this ns loads with the compiler, BEFORE the later core
+  ;; tiers — a name from 20-coll (min/max/abs) wouldn't resolve yet.
+  {"+" + "-" - "*" * "/" /
+   "<" < ">" > "<=" <= ">=" >= "=" =
+   "inc" inc "dec" dec
+   "mod" mod "rem" rem "quot" quot
+   "bit-and" bit-and "bit-or" bit-or "bit-xor" bit-xor})
+
+(defn- const? [n] (= :const (get n :op)))
+(defn- const-num? [n] (and (const? n) (number? (get n :val))))
+
+(defn- fold-fn [fnode]
+  (let [op (get fnode :op)]
+    (when (or (and (= op :var) (= "clojure.core" (get fnode :ns)))
+              (= op :host))
+      (get foldable (get fnode :name)))))
+
+(defn const-fold
+  "Bottom-up constant folding: a call of a foldable numeric fn whose args are
+  all constant numbers becomes a constant; an if with a constant test becomes
+  the taken branch."
+  [node]
+  (let [op (get node :op)]
+    (cond
+      (= op :invoke)
+      (let [f (const-fold (get node :fn))
+            args (mapv const-fold (get node :args))
+            ff (fold-fn f)
+            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)))
+
+      (= op :if)
+      (let [t (const-fold (get node :test))]
+        (if (const? t)
+          ;; jolt truthiness = Clojure's: nil/false take else
+          (if (or (nil? (get t :val)) (= false (get t :val)))
+            (const-fold (get node :else))
+            (const-fold (get node :then)))
+          (assoc node
+                 :test t
+                 :then (const-fold (get node :then))
+                 :else (const-fold (get node :else)))))
+
+      (= op :do)
+      (assoc node
+             :statements (mapv const-fold (get node :statements))
+             :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)))
+
+(defn run-passes
+  "All passes, in order. The back end applies this to every analyzed form."
+  [node]
+  (const-fold node))

src/jolt/backend.janet

@@ -385,7 +385,8 @@
 # kernel tier must already be loaded (see api/load-core-overlay!).
 (defn- build-compiler! [ctx]
   (compile-load ctx "jolt.ir")
-  (compile-load ctx "jolt.analyzer"))
+  (compile-load ctx "jolt.analyzer")
+  (compile-load ctx "jolt.passes"))
 
 (defn- ensure-analyzer [ctx]
   # Don't build until the kernel tier is loaded (see api/load-core-overlay! and
@@ -435,7 +436,20 @@
   (def r (protect ((var-get av) ctx form)))
   (put (ctx :env) :compile-ns nil)
   (ctx-set-current-ns ctx saved-ns)
-  (if (r 0) (r 1) (error (r 1))))
+  (unless (r 0) (error (r 1)))
+  # IR passes (jolt.passes/run-passes — nanopass-lite, jolt-2om): pure IR->IR
+  # rewrites (constant folding, ...) between the analyzer and the back end.
+  # Resolved lazily; absent during the pre-passes bootstrap window.
+  (def pv (unless (= "1" (os/getenv "JOLT_NO_IR_PASSES"))
+            (ns-find (ctx-find-ns ctx "jolt.passes") "run-passes")))
+  (if pv
+    (let [pr (protect ((var-get pv) (r 1)))]
+      # the pass runs interpreted; a throw inside it unwinds past the
+      # interpreter's ns restores — put the compile ns back either way, or
+      # the REST of this compilation resolves in jolt.passes
+      (ctx-set-current-ns ctx saved-ns)
+      (if (pr 0) (pr 1) (r 1)))
+    (r 1)))
 
 # The analyzer's deliberate punt signal — (uncompilable why) throws the string
 # "jolt/uncompilable: <why>". Anything else escaping the compile step is an

test/integration/ir-passes-test.janet

@@ -0,0 +1,36 @@
+# IR pass pipeline (jolt-2om, nanopass-lite): jolt.passes/run-passes applies
+# pure IR->IR rewrites between the analyzer and the back end. The first pass
+# is constant folding — it computes with the ACTUAL jolt fns, so folded
+# results match runtime semantics by construction.
+(import ../../src/jolt/api :as api)
+(import ../../src/jolt/backend :as backend)
+(import ../../src/jolt/reader :as reader)
+
+(print "IR passes (constant folding)...")
+(def ctx (api/init-cached {:compile? true}))
+(defn ir [src] (backend/analyze-form ctx (reader/parse-string src)))
+
+(defn check-const [src want]
+  (def n (ir src))
+  (assert (= :const (n :op)) (string src " folds to a constant"))
+  (assert (= want (n :val)) (string src " folds to " (string/format "%q" want))))
+
+(check-const "(+ 1 2 3)" 6)
+(check-const "(* 2 (+ 3 4))" 14)
+(check-const "(quot 7 2)" 3)
+(check-const "(mod -7 3)" 2)
+(check-const "(if (< 1 2) :yes :no)" :yes)
+# dead-branch elimination: the untaken branch never evaluates
+(check-const "(if false (this-would-not-resolve) 2)" 2)
+
+# non-constants stay calls; folding must be conservative
+(assert (= :invoke ((ir "(+ x 2)") :op)) "free var stays a call")
+(assert (= :invoke ((ir "(mod x 0)") :op)) "non-const args stay calls")
+# a fold that would THROW is left for runtime
+(assert (= :invoke ((ir "(mod 5 0)") :op)) "throwing fold left to runtime")
+
+# and the folded code evaluates identically (3-mode conformance covers the
+# 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")
+(print "IR passes passed!")