(ns ^:typed.clojure com.tailrecursion.blacktide.compiler.passes.compact
"Discard unused storage without discarding evaluation. Substitute compatible
copies and single-use, total primitive computations only before their inputs
are rebound. Native bodies remain opaque, including their free local reads."
(:require [clojure.set :as set]
[com.tailrecursion.nanopass :as np]
[com.tailrecursion.blacktide.compiler.languages :as l]
[com.tailrecursion.blacktide.compiler.passes.native :as native]
[com.tailrecursion.blacktide.compiler.jvm :as jvm]
[typed.clojure :as-alias t]))
(np/defpasses
(atom-reads {:from [l/Compact Atom] :to (t/Vec t/Sym)}
[node]
:local [(:id node)] :literal [] :inline (rhs-reads (:value node)))
(rhs-reads {:from [l/Compact Rhs] :to (t/Vec t/Sym)}
[node]
:value (atom-reads (:value node))
:global [] :external [] :var []
:native (vec (native/reads (:value node)))
:dispatch (atom-reads (:value (:value node)))
:convert (atom-reads (:value (:value node)))
:construct (into [] (mapcat atom-reads) (:items node))
:metadata (into (atom-reads (:value node)) (atom-reads (:metadata node)))
:host (into [] (mapcat #(atom-reads (:value %))) (:args node))
:invoke (into (atom-reads (:function node)) (mapcat atom-reads) (:args node))))
(np/defpass reads {:from [l/Compact Expr] :to (t/Map t/Sym t/Int)}
[node]
:bind (merge-with + (frequencies (rhs-reads (:value (:binding node))))
(dissoc (reads (:body node)) (:target (:binding node))))
:effect (merge-with + (frequencies (rhs-reads (:value (:binding node)))) (reads (:body node)))
:join (merge-with + (reads (:value node)) (apply dissoc (reads (:body node)) (:targets node)))
:branch (merge-with + (frequencies (atom-reads (:test node))) (reads (:yes node)) (reads (:no node)))
:switch (reduce (fn [^{::t/- (t/Map t/Sym t/Int)} refs ^{::t/- l/Compact-Expr} arm]
(merge-with + refs (reads arm))) (frequencies (atom-reads (:test node))) (:arms node))
:loop (reduce (fn [^{::t/- (t/Map t/Sym t/Int)} refs ^{::t/- l/Compact-Binding} binding]
(merge-with + (frequencies (atom-reads (:value binding))) (dissoc refs (:target binding))))
(reads (:body node)) (rseq (:bindings node)))
:recur (frequencies (into [] (mapcat atom-reads) (:args node)))
:return (frequencies (into [] (mapcat atom-reads) (:values node))) :throw (frequencies (atom-reads (:value node))))
(np/defpass stable?
{:from [l/Compact Expr] :to t/Bool :args [t/Sym (t/Set t/Sym)]}
[id dependencies node]
:bind (if (contains? dependencies (:target (:binding node)))
(not (contains? (reads (:body node)) id)) (stable? id dependencies (:body node)))
:effect (stable? id dependencies (:body node))
:join (and (stable? id dependencies (:value node))
(if (some dependencies (:targets node))
(not (contains? (reads (:body node)) id)) (stable? id dependencies (:body node))))
:branch (and (stable? id dependencies (:yes node)) (stable? id dependencies (:no node)))
:switch (every? #(stable? id dependencies %) (:arms node))
;; Loop binders also rebind on the backedge. Do not move an outer snapshot
;; into a loop that changes one of its inputs, even if the first entry agrees.
:loop (and (empty? (set/intersection dependencies (set (map :target (:bindings node)))))
(stable? id dependencies (:body node)))
:recur true :return true :throw true)
(np/defpass atom-type
{:from [l/Compact Atom] :to (t/Option java.lang.Class) :args [l/Compact-Types]}
[types node]
:local (get types (:id node))
:literal (cond (instance? Long (:value node)) Long/TYPE (instance? Double (:value node)) Double/TYPE :else nil)
:inline (:type node))
(np/defpass movable?
{:from [l/Compact Rhs] :to t/Bool :args [l/Compact-Types] :body true}
[types rhs]
(and (= :host (:op rhs))
(let [operation (:operation rhs)]
(and (= :static (:op operation))
(or (and (identical? clojure.lang.Numbers (:class operation))
(contains? '#{lt lte gt gte equiv unchecked_add unchecked_minus unchecked_multiply
unchecked_inc unchecked_dec unchecked_negate} (:member operation)))
(and (identical? clojure.lang.Util (:class operation)) (= 'equiv (:member operation))))
(every? (fn [^{::t/- l/Compact-Operand} operand]
(let [type (atom-type types (:value operand))]
(and (contains? #{Long/TYPE Double/TYPE} type) (= type (:type (:jvm operand)))
;; An inlined operand can still throw. Moving a pure
;; comparison must not move its checked conversion.
(or (not= :inline (:op (:value operand)))
(movable? types (:value (:value operand)))))))
(:args rhs))
(let [type (jvm/host-type (mapv #(:type (:jvm %)) (:args rhs)) operation)]
(and (class? type) (.isPrimitive ^Class type) (not= Void/TYPE type)))))))
(np/defpass replacement
{:from [l/Compact Operation] :to (t/Option l/Compact-Atom) :args [l/Compact-Types t/Int] :body true}
[types uses binding]
(let [rhs (:value binding) ^Class type (get types (:target binding))]
(case (:op rhs)
:value (let [value (:value rhs)]
;; A bound reference must not become several independently boxed reads.
(when (and (= type (atom-type types value))
(or type (= :local (:op value)))
;; Reading a primitive local boxes afresh; emitting its
;; literal at an Object use instead loads a pooled box.
(or (not= :literal (:op value))
(and (instance? Long (:value value)) (<= -128 (:value value) 127)))
(or (not= :inline (:op value)) (= uses 1))) value))
:convert (let [conversion (:conversion rhs) operand (:value rhs)]
(when (and type (:parameter? conversion) (= type (:type conversion))
(= type (atom-type types (:value operand))) (= type (:type (:jvm operand))))
(:value operand)))
:host (let [operation (:operation rhs) operand (first (:args rhs))
cast-name (when (class? type) (symbol (str (.getName ^Class type) "Cast")))]
;; RT's same-width primitive casts are identities only after both
;; storage and source operand types agree. Keep Object results:
;; those casts may establish a fresh box at a source-local read.
(if (and type (= :static (:op operation)) (= clojure.lang.RT (:class operation))
(= 1 (count (:args rhs))) operand
(= type (atom-type types (:value operand)) (:type (:jvm operand)))
(= cast-name (:member operation))
(= [type] (:tags operation)))
(:value operand)
(when (and type (= uses 1) (movable? types rhs)) {:op :inline :value rhs :type type})))
nil)))
(np/defpasses
(substitute-atom
{:from [l/Compact Atom] :to [l/Compact Atom] :args [t/Sym l/Compact-Atom]
:processors {Rhs substitute-rhs}}
[from replacement node]
:local (if (= from (:id node)) replacement node))
(substitute-operand
{:from [l/Compact Operand] :to [l/Compact Operand] :args [t/Sym l/Compact-Atom]
:processors {Atom substitute-atom}} [from replacement node])
(substitute-rhs
{:from [l/Compact Rhs] :to [l/Compact Rhs] :args [t/Sym l/Compact-Atom]
:processors {Atom substitute-atom Operand substitute-operand}} [from replacement node]))
(np/defpass substitute-operation
{:from [l/Compact Operation] :to [l/Compact Operation] :args [t/Sym l/Compact-Atom]
:processors {Rhs substitute-rhs}} [from replacement node])
(np/defpass substitute-binding
{:from [l/Compact Binding] :to [l/Compact Binding] :args [t/Sym l/Compact-Atom]
:processors {Atom substitute-atom}} [from replacement node])
(np/defpass substitute
{:from [l/Compact Expr] :to [l/Compact Expr] :args [t/Sym l/Compact-Atom]
:processors {Atom substitute-atom Rhs substitute-rhs Operation substitute-operation Binding substitute-binding}}
[from replacement node]
:bind (assoc node :binding (substitute-operation from replacement (:binding node))
:body (if (= from (:target (:binding node))) (:body node)
(substitute from replacement (:body node))))
:join (assoc node :value (substitute from replacement (:value node))
:body (if (some #{from} (:targets node)) (:body node)
(substitute from replacement (:body node))))
;; Loop initializers see earlier bindings, but the body sees all loop binders.
:loop ((fn ^{::t/- [t/Int (t/Vec l/Compact-Binding) t/Bool :-> l/Compact-Expr]} bindings [index pairs shadowed?]
(if (< index (count (:bindings node)))
(let [binding (nth (:bindings node) index)]
(recur (inc index) (conj pairs (if shadowed? binding (substitute-binding from replacement binding)))
(or shadowed? (= from (:target binding)))))
(assoc node :bindings pairs :body (if shadowed? (:body node)
(substitute from replacement (:body node)))))) 0 [] false))
(np/defpass discardable? {:from [l/Compact Rhs] :to t/Bool :body true}
[rhs]
;; Keep native validation, including constants that may fail AOT embedding.
(and (= :value (:op rhs))
(let [value (:value rhs)]
(or (= :local (:op value)) (and (= :literal (:op value)) (nil? (:value value)))))))
(np/defpasses
(simplify-bindings
{:from [l/Compact Expr] :to (t/HVec [l/Compact-Expr (t/Map t/Sym t/Int)])
:args [l/Compact-Types] :body true}
[types node]
;; Process sequential bindings inside out without growing the JVM stack.
;; Each result carries free-read counts for the next enclosing binding.
((fn ^{::t/- [(t/Vec l/Compact-Expr) l/Compact-Expr
:-> (t/HVec [l/Compact-Expr (t/Map t/Sym t/Int)])]} collect [parents node]
(case (:op node)
(:bind :effect) (recur (conj parents node) (:body node))
(reduce
(fn [^{::t/- (t/HVec [l/Compact-Expr (t/Map t/Sym t/Int)])} result
^{::t/- l/Compact-Expr} parent]
(let [[body refs] result]
(case (:op parent)
:bind
(let [binding (:binding parent) id (:target binding)
uses (get refs id 0) value (replacement types uses binding)
keep-binding (fn [] [(assoc parent :body body)
(merge-with + (frequencies (rhs-reads (:value binding))) (dissoc refs id))])]
(cond
(and (= :native (:op (:value binding))) (not (contains? types id))
(= body {:op :return :values [{:op :local :id id}]}))
[{:op :return :values [{:op :inline :value (:value binding) :type nil}]}
(frequencies (rhs-reads (:value binding)))]
(zero? uses) (if (discardable? (:value binding)) [body refs]
[{:op :effect :binding binding :body body}
(merge-with + (frequencies (rhs-reads (:value binding))) refs)])
(and value (stable? id (set (atom-reads value)) body))
(let [out (substitute id value body) refs (reads out)]
;; Substitution deliberately leaves native code alone.
(if (contains? refs id) (keep-binding) [out refs]))
:else (keep-binding)))
:effect
(let [rhs (:value (:binding parent))]
(if (discardable? rhs) [body refs]
[(assoc parent :body body) (merge-with + (frequencies (rhs-reads rhs)) refs)])))))
(simplify types node) (rseq parents)))) [] node))
(simplify
{:from [l/Compact Expr] :to (t/HVec [l/Compact-Expr (t/Map t/Sym t/Int)]) :args [l/Compact-Types]}
[types node]
:bind (simplify-bindings types node)
:effect (simplify-bindings types node)
:join (let [[value value-refs] (simplify types (:value node))
[body body-refs] (simplify types (:body node)) targets (:targets node)]
(if (and (= 1 (count targets)) (not (contains? types (first targets)))
(= body {:op :return :values [{:op :local :id (first targets)}]}))
[value value-refs]
[(assoc node :value value :body body) (merge-with + value-refs (apply dissoc body-refs targets))]))
:branch (let [[yes yes-refs] (simplify types (:yes node)) [no no-refs] (simplify types (:no node))]
[(assoc node :yes yes :no no)
(merge-with + (frequencies (atom-reads (:test node))) yes-refs no-refs)])
:switch (let [arms (mapv #(simplify types %) (:arms node))]
[(assoc node :arms (mapv first arms))
(reduce (fn [^{::t/- (t/Map t/Sym t/Int)} refs
^{::t/- (t/HVec [l/Compact-Expr (t/Map t/Sym t/Int)])} arm]
(merge-with + refs (second arm))) (frequencies (atom-reads (:test node))) arms)])
:loop (let [[body refs] (simplify types (:body node))]
[(assoc node :body body)
(reduce (fn [^{::t/- (t/Map t/Sym t/Int)} refs ^{::t/- l/Compact-Binding} binding]
(merge-with + (frequencies (atom-reads (:value binding))) (dissoc refs (:target binding))))
refs (rseq (:bindings node)))])
:recur [node (reads node)] :return [node (reads node)] :throw [node (reads node)]))
(np/defpass program {:from [l/Compact Program] :to [l/Compact Program] :body true}
[program]
{:types (:types program) :body (first (simplify (:types program) (:body program)))})