; lib/artdag/optimize.sx — Phase 5: result-preserving DAG rewrites.
;   DCE    — drop nodes not reachable upstream from the requested outputs.
;   CSE    — free from content addressing: structurally identical subexpressions
;            already collapse to one node at build time (artdag/cse == build).
;   Fusion — collapse a maximal 1-to-1 chain of fusible unary ops into a single
;            "artdag/pipeline" node that replays the stages; output-equivalent.
;   optimize — fuse then DCE in one pass.
; Depends on dag.sx and analyze.sx.

; ---- dict helper ----

(define
  artdag/-dict-filter
  (fn
    (d keep?)
    (reduce
      (fn (acc k) (if (keep? k (get d k)) (assoc acc k (get d k)) acc))
      {}
      (keys d))))

(define
  artdag/-union
  (fn
    (a b)
    (reduce (fn (acc x) (if (artdag/member? x acc) acc (cons x acc))) a b)))

; ---- dead-node elimination ----
; keep only the outputs and their transitive dependencies; ids are preserved.
(define
  artdag/dce
  (fn
    (dag outputs)
    (let
      ((db (artdag/analyze dag)))
      (let
        ((live (reduce (fn (acc out) (artdag/-union (artdag/-union acc (list out)) (artdag/ancestors-of db out))) (list) outputs)))
        {:names (artdag/-dict-filter (artdag/dag-names dag) (fn (k v) (artdag/member? v live))) :order (filter (fn (id) (artdag/member? id live)) (artdag/dag-order dag)) :ok true :nodes (artdag/-dict-filter (artdag/dag-nodes dag) (fn (k v) (artdag/member? k live)))}))))

; ---- common-subexpression elimination ----
; structural sharing is inherent to content addressing: build already maps
; structurally identical specs to a single node/id.
(define artdag/cse artdag/build)

; ---- adjacent-op fusion (entry-level rewrite) ----

(define artdag/pipeline-op "artdag/pipeline")

(define
  artdag/-name->entry
  (fn
    (entries)
    (reduce
      (fn (m e) (assoc m (artdag/entry-name e) e))
      {}
      entries)))

; name -> list of dependent names
(define
  artdag/-deps-map
  (fn
    (entries)
    (reduce
      (fn
        (m e)
        (reduce
          (fn
            (mm i)
            (assoc
              mm
              i
              (cons
                (artdag/entry-name e)
                (if (has-key? mm i) (get mm i) (list)))))
          m
          (artdag/entry-inputs e)))
      {}
      entries)))

(define artdag/-stage (fn (e) {:op (artdag/entry-op e) :params (artdag/entry-params e)}))

; the single predecessor that `name` may absorb, or nil. Requires: name is a
; fusible unary op; its one input is a locally-defined fusible node whose ONLY
; dependent is name (so fusing cannot break sharing).
(define
  artdag/-absorbs
  (fn
    (n->e deps fusible? name)
    (let
      ((e (get n->e name)))
      (let
        ((ins (artdag/entry-inputs e)))
        (if
          (= (len ins) 1)
          (let
            ((x (first ins)))
            (if
              (and
                (has-key? n->e x)
                (fusible? (artdag/entry-op e))
                (fusible? (artdag/entry-op (get n->e x)))
                (= (get deps x) (list name)))
              x
              nil))
          nil)))))

(define
  artdag/-absorbed-set
  (fn
    (n->e deps fusible? names)
    (reduce
      (fn
        (acc y)
        (let
          ((p (artdag/-absorbs n->e deps fusible? y)))
          (if (nil? p) acc (cons p acc))))
      (list)
      names)))

; walk predecessors from a tail, building stages head->tail.
(define
  artdag/-fuse-chain
  (fn
    (n->e deps fusible? cur stages)
    (let
      ((p (artdag/-absorbs n->e deps fusible? cur)))
      (if
        (nil? p)
        {:stages (cons (artdag/-stage (get n->e cur)) stages) :head cur}
        (artdag/-fuse-chain
          n->e
          deps
          fusible?
          p
          (cons (artdag/-stage (get n->e cur)) stages))))))

(define
  artdag/fuse-entries
  (fn
    (entries fusible?)
    (let
      ((n->e (artdag/-name->entry entries))
        (deps (artdag/-deps-map entries))
        (names (map artdag/entry-name entries)))
      (let
        ((absorbed (artdag/-absorbed-set n->e deps fusible? names)))
        (map
          (fn
            (name)
            (let
              ((c (artdag/-fuse-chain n->e deps fusible? name (list))))
              (if
                (> (len (get c :stages)) 1)
                (list
                  name
                  artdag/pipeline-op
                  (artdag/entry-inputs (get n->e (get c :head)))
                  {:stages (get c :stages)})
                (get n->e name))))
          (filter (fn (name) (not (artdag/member? name absorbed))) names))))))

(define
  artdag/fuse
  (fn
    (entries fusible?)
    (artdag/build (artdag/fuse-entries entries fusible?))))

; runner that replays a fused pipeline over its single input, delegating each
; stage to a base runner; non-pipeline ops fall through unchanged.
(define
  artdag/pipeline-run
  (fn
    (base-runner)
    (fn
      (params inputs)
      (reduce
        (fn
          (val stage)
          (base-runner (get stage :op) (get stage :params) (list val)))
        (first inputs)
        (get params :stages)))))

(define
  artdag/fusing-runner
  (fn
    (base-runner)
    (fn
      (op params inputs)
      (if
        (= op artdag/pipeline-op)
        ((artdag/pipeline-run base-runner) params inputs)
        (base-runner op params inputs)))))

; ---- full optimization pass ----
; fuse the entry list, then drop everything not feeding the requested output
; names. Output names survive fusion (sinks are never absorbed).
(define
  artdag/optimize
  (fn
    (entries outputs fusible?)
    (let
      ((fused (artdag/fuse entries fusible?)))
      (artdag/dce fused (map (fn (nm) (artdag/dag-id fused nm)) outputs)))))