(ns ^:typed.clojure com.tailrecursion.blacktide.compiler.passes.loop-facts
"Recover source loop representations using Clojure's monotone boxing analysis.
Reanalyze expanded bindings in order without repeating user macros. Native
syntax retains its initializers; lowered loops materialize the conversions."
(:require [com.tailrecursion.nanopass :as np]
[com.tailrecursion.blacktide.compiler.languages :as l]
[com.tailrecursion.blacktide.compiler.jvm :as jvm]
[com.tailrecursion.blacktide.compiler.errors :as errors]
[com.tailrecursion.blacktide.compiler.passes.source-jvm :as facts]
[typed.clojure :as-alias t]))
(np/defpass children {:from [l/Bound Expr] :to [l/Bound Exprs] :body true}
[expr]
(let [node (:node expr)]
(case (:op node)
:if [(:test node) (:then node) (:else node)]
:case (into [(:selector node)] (:arms node))
(:dispatch :throw) [(:value node)]
:do (:forms node)
:collection (cond-> (:items node) (:metadata node) (conj (:metadata node)))
(:host :recur) (:args node)
:invoke (into [(:function node)] (:args node))
(:let :loop) (conj (mapv :init (:bindings node)) (:body node))
(:setq :psetq) (mapv :value (:pairs node))
[])))
(np/defpass primitive {:from [l/JVM Fact] :to (t/Option java.lang.Class) :body true}
[fact]
(let [type (:type fact)]
(when (and (:primitive? fact) (class? type) (.isPrimitive ^Class type) (not= Void/TYPE type)) type)))
(np/defpass inferred {:from [l/Bound Expr] :to [l/Bound Expr] :args [l/Bound-Node] :body true}
[node expr]
(assoc expr :node node :jvm (facts/infer {:hint (:hint (:jvm expr)) :var-meta (:var-meta (:jvm expr))} node)))
(np/defpass converted {:from [l/Bound Expr] :to [l/Bound Expr] :args [java.lang.Class] :body true}
[^Class type expr]
(let [node {:op :host :operation {:op :static :class clojure.lang.RT
:member (symbol (str (.getName type) "Cast")) :tags [type]}
:args [expr]}]
(assoc expr :node node :jvm (facts/infer {:hint nil :var-meta {}} node))))
(np/defpass boxed {:from [l/Bound Expr] :to [l/Bound Expr] :body true}
[expr]
(let [node {:op :host :operation {:op :static :class clojure.lang.RT :member 'box :tags nil} :args [expr]}]
;; Reanalysis can select RT.box(Object), rather than RT.box(long/double).
(assoc expr :node node :jvm (facts/infer {:hint nil :var-meta {}} node))))
(np/defpass initializer {:from [l/Bound Expr] :to [l/Bound Expr] :args [t/Bool] :body true}
[box? init]
(let [type (primitive (:jvm init))]
(if box? (boxed init)
(if-let [promoted (cond (= type Integer/TYPE) Long/TYPE (= type Float/TYPE) Double/TYPE)]
(converted promoted init) init))))
(np/defpass recurs {:from [l/Bound Expr] :to [l/Bound Exprs] :args [t/Sym] :body true}
[target expr]
(if (and (= :recur (:op (:node expr))) (= target (:id (:target (:node expr)))))
[expr] (into [] (mapcat #(recurs target %)) (children expr))))
(np/defpass compatible? {:from [l/JVM Fact] :to t/Bool :args [java.lang.Class] :body true}
[type fact]
(let [actual (primitive fact)]
(or (= type actual) (and (= type Long/TYPE) (contains? #{Integer/TYPE Short/TYPE Byte/TYPE Character/TYPE} actual))
(and (= type Double/TYPE) (= actual Float/TYPE)))))
(np/defpasses
(bindings
{:from [l/Bound Bindings] :to (t/HVec [(t/Map t/Sym l/Bound-Local) l/Bound-Bindings])
:args [(t/Map t/Sym l/Bound-Local) t/Bool (t/Set t/Sym)] :body true}
[env loop? boxed-ids originals]
(reduce (fn [^{::t/- (t/HVec [(t/Map t/Sym l/Bound-Local) l/Bound-Bindings])} state
^{::t/- l/Bound-Binding} binding]
(let [[env out] state init (expression env (:init binding)) local (:local binding)
effective (if loop? (initializer (contains? boxed-ids (:id local)) init) init)
value (facts/binding-facts (:name local) (:mutable? local) (:jvm effective))
local (assoc local :jvm value)]
[(assoc env (:id local) local) (conj out {:local local :init init})])) [env []] originals))
(expression
{:from [l/Bound Expr] :to [l/Bound Expr] :args [(t/Map t/Sym l/Bound-Local)] :body true}
[env expr]
(let [node (:node expr)
node (case (:op node)
:local (assoc node :local (get env (:id (:local node)) (:local node)))
:if (assoc node :test (expression env (:test node))
:then (expression env (:then node)) :else (expression env (:else node)))
:case (assoc node :selector (expression env (:selector node))
:arms (mapv #(expression env %) (:arms node)))
(:dispatch :throw) (assoc node :value (expression env (:value node)))
:do (assoc node :forms (mapv #(expression env %) (:forms node)))
:collection (assoc node :items (mapv #(expression env %) (:items node))
:metadata (some->> (:metadata node) (expression env)))
:host (assoc node :args (mapv #(expression env %) (:args node)))
:recur (assoc node :args (mapv #(expression env %) (:args node))
:target (assoc (:target node) :bindings
(mapv #(get env (:id %) %) (:bindings (:target node)))))
:invoke (assoc node :function (expression env (:function node))
:args (mapv #(expression env %) (:args node)))
:let (let [[inner values] (bindings env false #{} (:bindings node))]
(assoc node :bindings values :body (expression inner (:body node))))
:loop (:node (first (solve env expr)))
(:setq :psetq) (assoc node :pairs (mapv #(assoc % :value (expression env (:value %))) (:pairs node)))
node)]
(inferred node expr)))
(solve
{:from [l/Bound Expr] :to (t/HVec [l/Bound-Expr (t/Set t/Sym)])
:args [(t/Map t/Sym l/Bound-Local)] :body true}
[env expr]
(let [node (:node expr)]
(if-not (= :loop (:op node)) (throw (IllegalStateException. "Expected source loop"))
((fn ^{::t/- [(t/Set t/Sym) :-> (t/HVec [l/Bound-Expr (t/Set t/Sym)])]} fixed [boxed-ids]
(let [[env values] (bindings env true boxed-ids (:bindings node))
body (expression env (:body node))
mismatches (reduce
(fn [^{::t/- (t/Set t/Sym)} out ^{::t/- l/Bound-Expr} edge]
(into out (keep identity)
(map (fn [^{::t/- l/Bound-Binding} binding ^{::t/- l/Bound-Expr} arg]
(let [local (:local binding) type (primitive (:jvm local))]
(when (and type (contains? #{Long/TYPE Double/TYPE} type)
(not (compatible? type (:jvm arg)))) (:id local))))
values (:args (:node edge))))) boxed-ids (recurs (:id node) body))]
(if (not= boxed-ids mismatches) (recur mismatches)
[(inferred (assoc node :bindings values :body body) expr) boxed-ids]))) #{})))))
;; RecurExpr's parser admits byte/short/char to long, but its emitter rejects
;; them. Preserve that validation before replacing recur with graph edges.
(np/defpass recur-error
{:from [l/Bound Expr] :to (t/Option (t/HVec [l/Bound-Local l/Bound-Expr])) :body true}
[expr]
(let [node (:node expr)]
(when (= :loop (:op node))
(first
(keep (fn [^{::t/- l/Bound-Expr} edge]
(first
(keep identity
(map (fn [^{::t/- l/Bound-Binding} binding ^{::t/- l/Bound-Expr} arg]
(let [local (:local binding) type (primitive (:jvm local)) actual (primitive (:jvm arg))]
(when (and type (not (or (= type actual)
(and (= type Long/TYPE) (= actual Integer/TYPE))
(and (= type Double/TYPE) (= actual Float/TYPE)))))
[local arg]))) (:bindings node) (:args (:node edge))))))
(recurs (:id node) (:body node)))))))
(np/defpass recurrence {:from [l/Bound Expr] :to [l/Bound Expr] :body true}
[expr]
(let [node (:node expr)]
(case (:op node)
:recur (assoc expr :node
(assoc node :args
(mapv (fn [local arg]
(if-let [type (primitive (:jvm local))]
(if (= type (primitive (:jvm arg))) arg (converted type arg)) arg))
(:bindings (:target node)) (:args node))))
:if (assoc expr :node (assoc node :then (recurrence (:then node)) :else (recurrence (:else node))))
:do (assoc expr :node (assoc node :forms (mapv recurrence (:forms node))))
:case (assoc expr :node (assoc node :arms (mapv recurrence (:arms node))))
:let (assoc expr :node (assoc node :body (recurrence (:body node))))
expr)))
(np/defpass materialize {:from [l/Bound Expr] :to [l/Bound Expr] :args [(t/Set t/Sym)] :body true}
[boxed-ids expr]
(let [node (:node expr)]
(if-not (= :loop (:op node)) (throw (IllegalStateException. "Expected source loop"))
(assoc expr :node
(assoc node :body (recurrence (:body node)) :bindings
(mapv (fn [binding]
(let [local (:local binding) init (:init binding)
effective (initializer (contains? boxed-ids (:id local)) init)
type (primitive (:jvm init))
;; Keep primitive provenance available to storage analysis.
init (if (and type (contains? boxed-ids (:id local))) (converted type init) effective)
hint (jvm/tag-class (:hint (:jvm local)))
init (if hint (inferred {:op :host :operation {:op :adapt :class hint} :args [init]} init) init)]
(assoc binding :init init))) (:bindings node)))))))
(np/defpass lower {:from [l/Bound Expr] :to [l/Bound Expr] :body true}
[expr]
(let [[expr boxed-ids] (solve {} expr)]
(when-let [[local arg] (recur-error expr)]
(errors/fail :invalid-recur-type "recur argument does not match its primitive loop binding"
{:binding (:name local) :expected (:type (:jvm local)) :actual (:type (:jvm arg))} (:site arg)))
(materialize boxed-ids expr)))