(ns ^:typed.clojure com.tailrecursion.blacktide.compiler.passes.joins
"Plan shared branch continuations and bounded loop exits. Postdominance
nominates follows; a least fixed point proves that every exit path reaches
its follow or throws. Joins carry only changed locals live at that follow."
(:require [clojure.set :as set]
[com.tailrecursion.nanopass :as np]
[com.tailrecursion.blacktide.compiler.languages :as l]
[com.tailrecursion.blacktide.compiler.passes.graph :as graph]
[com.tailrecursion.blacktide.compiler.passes.loops :as loops]
[typed.clojure :as-alias t]))
(np/defpass postdominators {:from [l/CFG Facts] :to [l/Planned Postdominators] :body true}
[facts]
(let [labels (set (keys facts))]
((fn ^{::t/- [l/Planned-Postdominators :-> l/Planned-Postdominators]} solve [known]
(let [next (reduce-kv
(fn [^{::t/- l/Planned-Postdominators} out ^{::t/- l/Machine-Label} id
^{::t/- l/CFG-BlockFacts} block]
(let [successors (:successors block)
common (if (seq successors)
(reduce set/intersection labels (map #(get known % #{}) successors)) #{})]
(assoc out id (conj common id)))) {} facts)]
(if (= known next) known (recur next)))) (zipmap labels (repeat labels)))))
(np/defpass regions
{:from [l/Numeric Program] :to [l/Planned Regions]
:args [(t/Option l/Machine-Label) (t/Set l/Machine-Label)] :body true}
[stop forbidden machine]
;; Prove blocks only after every successor is proved. Each block contributes
;; once; shared paths do not exhaust a traversal budget. Cycles never qualify.
((fn ^{::t/- [l/Planned-Regions :-> l/Planned-Regions]} solve [known]
(let [next (reduce-kv
(fn [^{::t/- l/Planned-Regions} out ^{::t/- l/Machine-Label} label
^{::t/- l/Numeric-Block} block]
(let [successors (graph/successors (:term block))]
(if (and (not (contains? forbidden label)) (not (contains? out label))
(or (= :throw (:op (:term block)))
(and (seq successors) (every? #(contains? out %) successors))))
(let [own (set (map :target (:ops block)))
children (map #(get out % {:blocks #{} :changed #{}}) successors)]
(assoc out label
{:blocks (into #{label} (mapcat :blocks) children)
:changed (into own (mapcat :changed) children)})) out))) known (:blocks machine))]
(if (= known next) known (recur next)))) (if stop {stop {:blocks #{} :changed #{}}} {})))
(np/defpass candidate
{:from [l/Machine Label] :to (t/Option l/Planned-Join)
:args [l/Numeric-Program l/Planned-Postdominators] :body true}
[machine known id]
(when-let [block (get (:blocks machine) id)]
(let [term (:term block)
exits (sort-by (fn [^{::t/- l/Machine-Label} target] (- (count (get known target #{}))))
(disj (get known id #{}) id))]
(when (contains? #{:branch :switch} (:op term))
(when-let [target (first exits)]
;; Only the nearest postdominator can form a nested lexical join.
;; Ignore the branch's preceding definitions when proving its result.
(let [branch-machine (update machine :blocks assoc id (assoc block :ops []))]
(when-let [proof (get (regions target #{} branch-machine) id)]
(let [results (vec (sort (set/intersection (:changed proof) (get (:live machine) target #{}))))]
;; Acyclic sharing must not turn a cheap branch into a tuple
;; allocation. Multiple values are reserved for loop exits.
(when (<= (count results) 1) {:target target :results results :blocks (:blocks proof)})))))))))
(np/defpass loop-exits
{:from [l/Planned LoopBodies] :to [l/Planned Joins] :args [l/Numeric-Program l/CFG-Facts] :body true}
[machine facts bodies]
(let [throws (set (keys (regions nil #{} machine)))
;; Postdominance in normal flow only nominates candidates. The region
;; proof below still checks every real edge, including throwing tails.
normal (postdominators (update-vals (apply dissoc facts throws)
#(update % :successors (fn [edges] (filterv (complement throws) edges)))))]
(reduce-kv
(fn [^{::t/- l/Planned-Joins} out ^{::t/- l/Machine-Label} header
^{::t/- (t/Set l/Machine-Label)} members]
(let [exits (set/difference (into #{} (mapcat #(:successors (graph/fact-at facts %))) members) members)
ordinary (set/difference exits throws)
common (when (seq ordinary) (reduce set/intersection (map #(get normal % #{}) ordinary)))
follow (first (sort-by (fn [^{::t/- l/Machine-Label} id] (- (count (get normal id #{}))))
(set/difference (or common #{}) members)))]
(if follow
(let [proof (regions follow members machine)]
(if (every? #(contains? proof %) exits)
(let [tails (map #(get proof % {:blocks #{} :changed #{}}) exits)
changed (into (into #{} (mapcat #(:defs (graph/fact-at facts %))) members)
(mapcat :changed) tails)]
(assoc out header {:target follow :blocks (into members (mapcat :blocks) tails)
:results (vec (sort (set/intersection changed
(get (:live machine) follow #{}))))})) out)) out)))
{} bodies)))
(np/defpass plan {:from [l/Numeric Program] :to [l/Planned Program] :body true}
[machine]
(let [facts (graph/collect-facts (:blocks machine))
known (postdominators facts)
bodies (loops/analyze (:entry machine) facts)]
{:machine machine
:loops bodies :exits (if bodies (loop-exits machine facts bodies) {})
:joins (reduce (fn [^{::t/- l/Planned-Joins} out ^{::t/- l/Machine-Label} id]
(if-let [join (candidate machine known id)] (assoc out id join) out))
{} (keys (:blocks machine)))}))