(ns ^:typed.clojure com.tailrecursion.blacktide.compiler.passes.structure
"Recover lexical control from reducible loop regions. Nested loops yield
their live results; tail loops keep their continuation without a packet.
Recovery bounds code duplication and preserves dead-code validation."
(:require [com.tailrecursion.nanopass :as np]
[com.tailrecursion.blacktide.compiler.languages :as l]
[com.tailrecursion.blacktide.compiler.passes.graph :as graph]
[typed.clojure :as-alias t]))
(np/defpass block
{:from [l/Machine Label] :to [l/Numeric Block] :args [l/Numeric-Program] :body true}
[machine id]
(or (get (:blocks machine) id) (throw (IllegalStateException. "Missing checked block"))))
(np/defpass harmless-atom? {:from [l/Machine Atom] :to t/Bool}
[atom]
:local true
:literal (nil? (:value atom)))
(np/defpass harmless-block? {:from [l/Numeric Block] :to t/Bool :body true}
[block]
;; Lowering leaves nil result copies after transfers. No host/native form or
;; arbitrary quoted constant may disappear before Clojure validates it.
(and (every? (fn [op] (let [rhs (:value op)]
(and (= :value (:op rhs)) (harmless-atom? (:value rhs))))) (:ops block))
(let [term (:term block)]
(case (:op term)
:jump true
(:return :throw) (harmless-atom? (:value term))
(:branch :switch) (harmless-atom? (:test term))))))
(np/defpass wrap
{:from [l/Numeric Block] :to [l/Joined Expr] :args [l/Joined-Expr] :body true}
[end block]
(reduce (fn [^{::t/- l/Joined-Expr} body ^{::t/- l/Numeric-Operation} op]
{:op :bind :binding op :body body}) end (reverse (:ops block))))
(np/defpass within-stop?
{:from [l/Planned Join] :to t/Bool :args [l/Joined-Context] :body true}
[ctx join]
(or (nil? (:stop ctx)) (not (contains? (:blocks join) (first (:stop ctx))))))
(np/defpasses
(arms
{:from [l/Machine Term] :to [l/Joined Arms] :args [l/Joined-Context t/Int] :body true}
[ctx budget term]
(if (= :switch (:op term))
((fn ^{::t/- [t/Int t/Int (t/Vec l/Joined-Expr) :-> l/Joined-Arms]} build [index budget out]
(if (= index (count (:targets term))) {:op :built :budget budget :arms out}
(let [result (edge ctx budget (nth (:targets term) index))]
(if (= :built (:op result))
(recur (inc index) (:budget result) (conj out (:expr result))) result)))) 0 budget [])
{:op :exhausted :reason :budget}))
(terminator
{:from [l/Machine Term] :to [l/Joined Build] :args [l/Joined-Context t/Int]}
[ctx budget term]
:jump (edge ctx budget (:target term))
:return (if (:stop ctx) {:op :exhausted :reason :cycles}
{:op :built :budget budget :expr {:op :return :values [(:value term)]}})
:throw {:op :built :budget budget :expr term}
:switch (let [result (arms ctx budget term)]
(if (= :built (:op result))
{:op :built :budget (:budget result) :expr {:op :switch :test (:test term) :arms (:arms result)}} result))
:branch (let [yes (edge ctx budget (:yes term))]
(if (= :exhausted (:op yes)) yes
(let [no (edge ctx (:budget yes) (:no term))]
(if (= :exhausted (:op no)) no
{:op :built :budget (:budget no)
:expr {:op :branch :test (:test term) :yes (:expr yes) :no (:expr no)}})))))
(joined
{:from [l/Planned Join] :to [l/Joined Build] :args [l/Joined-Context t/Int l/Machine-Term] :body true}
[ctx budget term join]
(let [inner (assoc ctx :stop [(:target join) (:results join)])
result (terminator inner budget term)]
(if (= :built (:op result))
(let [tail (edge ctx (:budget result) (:target join))]
(if (= :built (:op tail))
(assoc tail :expr {:op :join :targets (:results join) :value (:expr result) :body (:expr tail)}) tail))
result)))
(body
{:from [l/Machine Label] :to [l/Joined Build] :args [l/Joined-Context t/Int] :body true}
[ctx budget id]
(let [block (block (:machine ctx) id)
budget (- budget (inc (count (:ops block))))]
(if (neg? budget) {:op :exhausted :reason :budget}
(let [join (get (:joins ctx) id)
result (if (and join (within-stop? ctx join))
(joined ctx budget (:term block) join)
(terminator ctx budget (:term block)))]
(if (= :built (:op result)) (assoc result :expr (wrap (:expr result) block)) result)))))
(loop-form
{:from [l/Machine Label] :to [l/Joined Build] :args [l/Joined-Context t/Int] :body true}
[ctx budget id]
(let [slots (vec (sort (get (:live (:machine ctx)) id #{})))
inner (assoc ctx :header id :slots slots :active (conj (:active ctx) id))
result (body inner (dec budget) id)]
(if (= :built (:op result))
(assoc result :expr {:op :loop :slots slots :body (:expr result)}) result)))
(edge
{:from [l/Machine Label] :to [l/Joined Build] :args [l/Joined-Context t/Int] :body true}
[ctx budget id]
(let [stop (:stop ctx)]
(cond
(and stop (= id (first stop)))
{:op :built :budget budget
:expr {:op :return :values (mapv (fn [id] {:op :local :id id}) (second stop))}}
(= id (:header ctx))
(if (pos? budget) {:op :built :budget (dec budget) :expr {:op :recur :slots (:slots ctx)}}
{:op :exhausted :reason :budget})
(contains? (:active ctx) id) {:op :exhausted :reason :cycles}
(contains? (:loops ctx) id)
(if (contains? (get (:loops ctx) (:header ctx) #{}) id)
(if-let [join (get (:exits ctx) id)]
(if (within-stop? ctx join)
(let [result (loop-form (assoc ctx :stop [(:target join) (:results join)]) budget id)]
(if (= :built (:op result))
(let [tail (edge ctx (:budget result) (:target join))]
(if (= :built (:op tail))
(assoc tail :expr {:op :join :targets (:results join) :value (:expr result) :body (:expr tail)}) tail)) result))
{:op :exhausted :reason :cycles})
{:op :exhausted :reason :cycles})
(loop-form ctx budget id))
:else (body ctx budget id)))))
(np/defpass recover {:from [l/Planned Program] :to [l/Joined Decision] :body true}
[plan]
(let [machine (:machine plan)
reachable (graph/reachable (graph/collect-facts (:blocks machine)) (:entry machine))
dead (apply dissoc (:blocks machine) reachable)
loops (:loops plan)]
(cond
(nil? loops) {:op :fallback :reason :cycles}
(not (every? harmless-block? (vals dead))) {:op :fallback :reason :validation}
:else
(let [ctx {:machine machine :header nil :slots [] :joins (:joins plan) :stop nil
:loops loops :exits (:exits plan) :active #{}}
budget (* 4 (reduce + 0 (map #(inc (count (:ops %))) (vals (select-keys (:blocks machine) reachable)))))
result (edge ctx budget (:entry machine))]
(if (= :built (:op result))
{:op :structured :body (:expr result)}
{:op :fallback :reason (:reason result)})))))