; Phase 5 — optimization: DCE, CSE (content-id sharing), adjacent-op fusion.

(define opt-BASE (artdag/op-table-runner {:in (fn (params inputs) (get params :v)) :sq (fn (params inputs) (* (first inputs) (first inputs))) :add (fn (params inputs) (+ (nth inputs 0) (nth inputs 1))) :inc (fn (params inputs) (+ 1 (first inputs)))}))
(define opt-RUN (artdag/fusing-runner opt-BASE))
(define opt-inc? (fn (op) (= op "inc")))
(define opt-incsq? (fn (op) (or (= op "inc") (= op "sq"))))

; linear chain a(in) -> b -> c -> d, all inc
(define
  opt-chain
  (list
    (list "a" "in" (list) {:v 5})
    (list "b" "inc" (list "a") {})
    (list "c" "inc" (list "b") {})
    (list "d" "inc" (list "c") {})))

; ---- DCE ----

(define
  dce-entries
  (list
    (list "a" "in" (list) {:v 5})
    (list "b" "inc" (list "a") {})
    (list "c" "inc" (list "b") {})
    (list "x" "sq" (list "a") {})))
(define dce-G (artdag/build dce-entries))

(artdag-test
  "dce: removes dead node"
  (artdag/node-count (artdag/dce dce-G (list (artdag/dag-id dce-G "c"))))
  3)

(artdag-test
  "dce: keeps live closure intact"
  (artdag/node-count (artdag/dce dce-G (list (artdag/dag-id dce-G "x"))))
  2)

(artdag-test
  "dce: preserves surviving node ids"
  (artdag/member?
    (artdag/dag-id dce-G "c")
    (keys
      (artdag/dag-nodes (artdag/dce dce-G (list (artdag/dag-id dce-G "c"))))))
  true)

(artdag-test
  "dce: output result unchanged after elimination"
  (let
    ((cache (persist/open)))
    (artdag/result-of
      (artdag/run
        (artdag/dce dce-G (list (artdag/dag-id dce-G "c")))
        opt-RUN
        cache)
      (artdag/dag-id dce-G "c")))
  7)

(artdag-test
  "dce: nothing dead is a no-op on count"
  (artdag/node-count
    (artdag/dce
      dce-G
      (list (artdag/dag-id dce-G "c") (artdag/dag-id dce-G "x"))))
  4)

; ---- CSE (free from content addressing) ----

(define
  cse-entries
  (list
    (list "a" "in" (list) {:v 3})
    (list "s1" "sq" (list "a") {})
    (list "s2" "sq" (list "a") {})
    (list "d" "add" (list "s1" "s2") {} true)))
(define cse-C (artdag/cse cse-entries))

(artdag-test
  "cse: identical subexpressions collapse to one node"
  (artdag/node-count cse-C)
  3)

(artdag-test
  "cse: shared node computes once"
  (let
    ((cache (persist/open)))
    (artdag/recompute-count (artdag/run cse-C opt-RUN cache)))
  3)

(artdag-test
  "cse: s1 and s2 are the same id"
  (equal? (artdag/dag-id cse-C "s1") (artdag/dag-id cse-C "s2"))
  true)

(artdag-test
  "cse: result is correct"
  (let
    ((cache (persist/open)))
    (artdag/result-of
      (artdag/run cse-C opt-RUN cache)
      (artdag/dag-id cse-C "d")))
  18)

; ---- fusion ----

(artdag-test
  "fusion: collapses a unary chain"
  (artdag/node-count (artdag/fuse opt-chain opt-inc?))
  2)

(artdag-test
  "fusion: unfused has all nodes"
  (artdag/node-count (artdag/build opt-chain))
  4)

(artdag-test
  "fusion: output-equivalent to unfused"
  (let
    ((c1 (persist/open)) (c2 (persist/open)))
    (=
      (artdag/result-of
        (artdag/run (artdag/build opt-chain) opt-RUN c1)
        (artdag/dag-id (artdag/build opt-chain) "d"))
      (artdag/result-of
        (artdag/run (artdag/fuse opt-chain opt-inc?) opt-RUN c2)
        (artdag/dag-id (artdag/fuse opt-chain opt-inc?) "d"))))
  true)

(artdag-test
  "fusion: leaf is never fused"
  (artdag/node-op
    (artdag/dag-node-by-name (artdag/fuse opt-chain opt-inc?) "a"))
  "in")

(artdag-test
  "fusion: tail becomes a pipeline node"
  (artdag/node-op
    (artdag/dag-node-by-name (artdag/fuse opt-chain opt-inc?) "d"))
  "artdag/pipeline")

(artdag-test
  "fusion: mixed fusible set fuses across op kinds"
  (artdag/node-count
    (artdag/fuse
      (list
        (list "a" "in" (list) {:v 2})
        (list "b" "inc" (list "a") {})
        (list "c" "sq" (list "b") {})
        (list "d" "inc" (list "c") {}))
      opt-incsq?))
  2)

(artdag-test
  "fusion: mixed chain replays correctly"
  (let
    ((cache (persist/open)))
    (let
      ((f (artdag/fuse (list (list "a" "in" (list) {:v 2}) (list "b" "inc" (list "a") {}) (list "c" "sq" (list "b") {}) (list "d" "inc" (list "c") {})) opt-incsq?)))
      (artdag/result-of (artdag/run f opt-RUN cache) (artdag/dag-id f "d"))))
  10)

(artdag-test
  "fusion: fanout node is not fused"
  (artdag/node-count
    (artdag/fuse
      (list
        (list "a" "in" (list) {:v 1})
        (list "b" "inc" (list "a") {})
        (list "c" "inc" (list "b") {})
        (list "e" "sq" (list "b") {}))
      opt-inc?))
  4)

(artdag-test
  "fusion: empty fusible set leaves dag unchanged"
  (artdag/node-count (artdag/fuse opt-chain (fn (op) false)))
  4)

; ---- full optimization pass (fuse + dce) ----

(define
  optp-entries
  (list
    (list "a" "in" (list) {:v 5})
    (list "b" "inc" (list "a") {})
    (list "c" "inc" (list "b") {})
    (list "x" "sq" (list "a") {})))

(artdag-test
  "optimize: fuses chain and drops dead node"
  (artdag/node-count (artdag/optimize optp-entries (list "c") opt-inc?))
  2)

(artdag-test
  "optimize: leaves dead node when it is an output"
  (artdag/node-count (artdag/optimize optp-entries (list "c" "x") opt-inc?))
  3)

(artdag-test
  "optimize: result equals the unoptimized dag"
  (let
    ((c1 (persist/open)) (c2 (persist/open)))
    (let
      ((o (artdag/optimize optp-entries (list "c") opt-inc?)))
      (=
        (artdag/result-of (artdag/run o opt-RUN c1) (artdag/dag-id o "c"))
        (artdag/result-of
          (artdag/run (artdag/build optp-entries) opt-RUN c2)
          (artdag/dag-id (artdag/build optp-entries) "c")))))
  true)

(artdag-test
  "optimize: no fusible ops still drops dead nodes"
  (artdag/node-count
    (artdag/optimize optp-entries (list "c") (fn (op) false)))
  3)