;; lib/datalog/tests/magic.sx — adornment + SIPS analysis tests.

(define dl-mt-pass 0)
(define dl-mt-fail 0)
(define dl-mt-failures (list))

(define
  dl-mt-deep=?
  (fn
    (a b)
    (cond
      ((and (list? a) (list? b))
        (and (= (len a) (len b)) (dl-mt-deq-l? a b 0)))
      ((and (dict? a) (dict? b))
        (let ((ka (keys a)) (kb (keys b)))
          (and (= (len ka) (len kb)) (dl-mt-deq-d? a b ka 0))))
      ((and (number? a) (number? b)) (= a b))
      (else (equal? a b)))))

(define
  dl-mt-deq-l?
  (fn
    (a b i)
    (cond
      ((>= i (len a)) true)
      ((not (dl-mt-deep=? (nth a i) (nth b i))) false)
      (else (dl-mt-deq-l? a b (+ i 1))))))

(define
  dl-mt-deq-d?
  (fn
    (a b ka i)
    (cond
      ((>= i (len ka)) true)
      ((let ((k (nth ka i)))
         (not (dl-mt-deep=? (get a k) (get b k))))
        false)
      (else (dl-mt-deq-d? a b ka (+ i 1))))))

(define
  dl-mt-test!
  (fn
    (name got expected)
    (if
      (dl-mt-deep=? got expected)
      (set! dl-mt-pass (+ dl-mt-pass 1))
      (do
        (set! dl-mt-fail (+ dl-mt-fail 1))
        (append!
          dl-mt-failures
          (str
            name
            "\n    expected: " expected
            "\n    got:      " got))))))

(define
  dl-mt-run-all!
  (fn
    ()
    (do
      ;; Goal adornment.
      (dl-mt-test! "adorn 0-ary"
        (dl-adorn-goal (list (quote ready)))
        "")
      (dl-mt-test! "adorn all bound"
        (dl-adorn-goal (list (quote p) 1 2 3))
        "bbb")
      (dl-mt-test! "adorn all free"
        (dl-adorn-goal (list (quote p) (quote X) (quote Y)))
        "ff")
      (dl-mt-test! "adorn mixed"
        (dl-adorn-goal (list (quote ancestor) (quote tom) (quote X)))
        "bf")
      (dl-mt-test! "adorn const var const"
        (dl-adorn-goal (list (quote p) (quote a) (quote X) (quote b)))
        "bfb")

      ;; dl-adorn-lit with explicit bound set.
      (dl-mt-test! "adorn lit with bound"
        (dl-adorn-lit (list (quote p) (quote X) (quote Y)) (list "X"))
        "bf")

      ;; Rule SIPS — chain ancestor.
      (dl-mt-test! "sips chain ancestor bf"
        (dl-rule-sips
          {:head (list (quote ancestor) (quote X) (quote Z))
           :body (list (list (quote parent) (quote X) (quote Y))
                       (list (quote ancestor) (quote Y) (quote Z)))}
          "bf")
        (list
          {:lit (list (quote parent) (quote X) (quote Y)) :adornment "bf"}
          {:lit (list (quote ancestor) (quote Y) (quote Z)) :adornment "bf"}))

      ;; SIPS — head fully bound.
      (dl-mt-test! "sips head bb"
        (dl-rule-sips
          {:head (list (quote q) (quote X) (quote Y))
           :body (list (list (quote p) (quote X) (quote Z))
                       (list (quote r) (quote Z) (quote Y)))}
          "bb")
        (list
          {:lit (list (quote p) (quote X) (quote Z)) :adornment "bf"}
          {:lit (list (quote r) (quote Z) (quote Y)) :adornment "bb"}))

      ;; SIPS — comparison; vars must be bound by prior body lit.
      (dl-mt-test! "sips with comparison"
        (dl-rule-sips
          {:head (list (quote q) (quote X))
           :body (list (list (quote p) (quote X))
                       (list (string->symbol "<") (quote X) 5))}
          "f")
        (list
          {:lit (list (quote p) (quote X)) :adornment "f"}
          {:lit (list (string->symbol "<") (quote X) 5) :adornment "bb"}))

      ;; SIPS — `is` binds its left arg.
      (dl-mt-test! "sips with is"
        (dl-rule-sips
          {:head (list (quote q) (quote X) (quote Y))
           :body (list (list (quote p) (quote X))
                       (list (quote is) (quote Y) (list (string->symbol "+") (quote X) 1)))}
          "ff")
        (list
          {:lit (list (quote p) (quote X)) :adornment "f"}
          {:lit (list (quote is) (quote Y)
                      (list (string->symbol "+") (quote X) 1))
           :adornment "fb"}))

      ;; Magic predicate naming.
      (dl-mt-test! "magic-rel-name"
        (dl-magic-rel-name "ancestor" "bf")
        "magic_ancestor^bf")

      ;; Bound-args extraction.
      (dl-mt-test! "bound-args bf"
        (dl-bound-args (list (quote ancestor) (quote tom) (quote X)) "bf")
        (list (quote tom)))

      (dl-mt-test! "bound-args mixed"
        (dl-bound-args (list (quote p) 1 (quote Y) 3) "bfb")
        (list 1 3))

      (dl-mt-test! "bound-args all-free"
        (dl-bound-args (list (quote p) (quote X) (quote Y)) "ff")
        (list))

      ;; Magic literal construction.
      (dl-mt-test! "magic-lit"
        (dl-magic-lit "ancestor" "bf" (list (quote tom)))
        (list (string->symbol "magic_ancestor^bf") (quote tom)))

      ;; Magic-sets rewriter: structural sanity.
      (dl-mt-test! "rewrite ancestor produces seed"
        (let
          ((rules
             (list
               {:head (list (quote ancestor) (quote X) (quote Y))
                :body (list (list (quote parent) (quote X) (quote Y)))}
               {:head (list (quote ancestor) (quote X) (quote Z))
                :body
                (list (list (quote parent) (quote X) (quote Y))
                      (list (quote ancestor) (quote Y) (quote Z)))})))
          (get
            (dl-magic-rewrite rules "ancestor" "bf" (list (quote a)))
            :seed))
        (list (string->symbol "magic_ancestor^bf") (quote a)))

      ;; Equivalence: rewritten program derives same ancestor tuples.
      ;; In a chain a→b→c→d, magic-rewritten run still derives all
      ;; ancestor pairs reachable from any node a/b/c/d propagated via
      ;; magic_ancestor^bf — i.e. the full closure (6 tuples). Magic
      ;; saves work only when the EDB has irrelevant nodes outside
      ;; the seed's transitive cone.
      (dl-mt-test! "magic-rewritten ancestor count"
        (let
          ((rules
             (list
               {:head (list (quote ancestor) (quote X) (quote Y))
                :body (list (list (quote parent) (quote X) (quote Y)))}
               {:head (list (quote ancestor) (quote X) (quote Z))
                :body
                (list (list (quote parent) (quote X) (quote Y))
                      (list (quote ancestor) (quote Y) (quote Z)))}))
            (edb (list
                   (list (quote parent) (quote a) (quote b))
                   (list (quote parent) (quote b) (quote c))
                   (list (quote parent) (quote c) (quote d)))))
          (let
            ((rewritten (dl-magic-rewrite rules "ancestor" "bf" (list (quote a))))
              (db (dl-make-db)))
            (do
              (for-each (fn (f) (dl-add-fact! db f)) edb)
              (dl-add-fact! db (get rewritten :seed))
              (for-each (fn (r) (dl-add-rule! db r)) (get rewritten :rules))
              (dl-saturate! db)
              (len (dl-relation db "ancestor")))))
        6)

      ;; dl-magic-query: end-to-end driver, doesn't mutate caller's db.
      ;; Magic over a rule with negated body literal — propagation
      ;; rules generated only for positive lits; negated lits pass
      ;; through unchanged.
      (dl-mt-test! "magic over rule with negation"
        (let
          ((db (dl-program
                 "u(a). u(b). u(c). banned(b).
                  active(X) :- u(X), not(banned(X)).")))
          (let
            ((semi (dl-query db (list (quote active) (quote X))))
              (magic (dl-magic-query db (list (quote active) (quote X)))))
            (= (len semi) (len magic))))
        true)

      ;; All-bound query (existence check) generates an "bb"
      ;; adornment chain. Verifies the rewriter walks multiple
      ;; (rel, adn) pairs through the worklist.
      (dl-mt-test! "magic existence check via bb"
        (let
          ((db (dl-program
                 "parent(a, b). parent(b, c). parent(c, d).
                  ancestor(X, Y) :- parent(X, Y).
                  ancestor(X, Z) :- parent(X, Y), ancestor(Y, Z).")))
          (let
            ((found (dl-magic-query
                      db (list (quote ancestor) (quote a) (quote c))))
              (missing (dl-magic-query
                         db (list (quote ancestor) (quote a) (quote z)))))
            (and (= (len found) 1) (= (len missing) 0))))
        true)

      ;; Magic equivalence on the federation demo.
      (dl-mt-test! "magic ≡ semi on foaf demo"
        (let
          ((db (dl-program-data
                 (quote ((follows alice bob)
                         (follows bob carol)
                         (follows alice dave)))
                 dl-demo-federation-rules)))
          (let
            ((semi (dl-query db (quote (foaf alice X))))
              (magic (dl-magic-query db (quote (foaf alice X)))))
            (= (len semi) (len magic))))
        true)

      ;; Shape validation: dl-magic-query rejects non-list / non-
      ;; dict goal shapes cleanly rather than crashing in `rest`.
      (dl-mt-test! "magic rejects string goal"
        (let ((threw false))
          (do
            (guard (e (#t (set! threw true)))
              (dl-magic-query (dl-make-db) "foo"))
            threw))
        true)

      (dl-mt-test! "magic rejects bare dict goal"
        (let ((threw false))
          (do
            (guard (e (#t (set! threw true)))
              (dl-magic-query (dl-make-db) {:foo "bar"}))
            threw))
        true)

      ;; 3-stratum program under magic — distinct rule heads at
      ;; strata 0/1/2 must all rewrite via the worklist.
      (dl-mt-test! "magic 3-stratum program"
        (let
          ((db (dl-program
                 "a(1). a(2). a(3). b(2).
                  c(X) :- a(X), not(b(X)).
                  d(X) :- c(X), not(banned(X)).
                  banned(3).")))
          (let
            ((semi (dl-query db (list (quote d) (quote X))))
              (magic (dl-magic-query db (list (quote d) (quote X)))))
            (= (len semi) (len magic))))
        true)

      ;; Aggregate -> derived -> threshold chain via magic.
      (dl-mt-test! "magic aggregate-derived chain"
        (let
          ((db (dl-program
                 "src(1). src(2). src(3).
                  cnt(N) :- count(N, X, src(X)).
                  active(N) :- cnt(N), >=(N, 2).")))
          (let
            ((semi (dl-query db (list (quote active) (quote N))))
              (magic (dl-magic-query db (list (quote active) (quote N)))))
            (= (len semi) (len magic))))
        true)

      ;; Multi-relation rewrite chain: query r4 → propagate to r3,
      ;; r2, r1, a. The worklist must process all of them; an
      ;; earlier bug stopped after only the head pair.
      (dl-mt-test! "magic chain through 4 rule levels"
        (let
          ((db (dl-program
                 "a(1). a(2). r1(X) :- a(X). r2(X) :- r1(X).
                  r3(X) :- r2(X). r4(X) :- r3(X).")))
          (= 2 (len (dl-magic-query db (list (quote r4) (quote X))))))
        true)

      ;; Shortest-path demo via magic — exercises the rewriter
      ;; against rules that mix recursive positive lits with an
      ;; aggregate body literal.
      (dl-mt-test! "magic on shortest-path demo"
        (let
          ((db (dl-program-data
                 (quote ((edge a b 5) (edge b c 3) (edge a c 10)))
                 dl-demo-shortest-path-rules)))
          (let
            ((semi (dl-query db (quote (shortest a c W))))
              (magic (dl-magic-query db (quote (shortest a c W)))))
            (and (= (len semi) (len magic))
                 (= (len semi) 1))))
        true)

      ;; Same relation called with different adornment patterns
      ;; in different rules. The worklist must enqueue and process
      ;; each (rel, adornment) pair.
      (dl-mt-test! "magic with multi-adornment same relation"
        (let
          ((db (dl-program
                 "parent(p1, alice). parent(p2, bob).
                  parent(g, p1). parent(g, p2).
                  sibling(P1, P2) :- parent(G, P1), parent(G, P2),
                                     !=(P1, P2).
                  cousin(X, Y) :- parent(P1, X), parent(P2, Y),
                                  sibling(P1, P2).")))
          (let
            ((semi (dl-query db (list (quote cousin) (quote alice) (quote Y))))
              (magic (dl-magic-query db (list (quote cousin) (quote alice) (quote Y)))))
            (= (len semi) (len magic))))
        true)

      ;; Magic over a rule whose body contains an aggregate.
      ;; The rewriter passes aggregate body lits through unchanged
      ;; (no propagation generated for them), so semi-naive's count
      ;; logic still fires correctly under the rewritten program.
      (dl-mt-test! "magic over rule with aggregate body"
        (let
          ((db (dl-program
                 "post(p1). post(p2). post(p3).
                  liked(u1, p1). liked(u2, p1). liked(u3, p1).
                  liked(u1, p2).
                  rich(P) :- post(P), count(N, U, liked(U, P)),
                             >=(N, 2).")))
          (let
            ((semi (dl-query db (list (quote rich) (quote P))))
              (magic (dl-magic-query db (list (quote rich) (quote P)))))
            (= (len semi) (len magic))))
        true)

      ;; Mixed EDB + IDB: a relation can be both EDB-seeded and
      ;; rule-derived. dl-magic-query must include the EDB portion
      ;; even though the relation has rules.
      (dl-mt-test! "magic mixed EDB+IDB"
        (len
          (dl-magic-query
            (dl-program
              "link(a, b). link(c, d). link(e, c).
               via(a, e).
               link(X, Y) :- via(X, M), link(M, Y).")
            (list (quote link) (quote a) (quote X))))
        2)

      ;; dl-magic-query falls back to dl-query for built-in,
      ;; aggregate, and negation goals (the magic seed would
      ;; otherwise be non-ground).
      (dl-mt-test! "magic-query falls back on aggregate"
        (let
          ((r (dl-magic-query
                (dl-program "p(1). p(2). p(3).")
                (list (quote count) (quote N) (quote X)
                      (list (quote p) (quote X))))))
          (and (= (len r) 1) (= (get (first r) "N") 3)))
        true)

      (dl-mt-test! "magic-query equivalent to dl-query"
        (let
          ((db (dl-program
                 "parent(a, b). parent(b, c). parent(c, d).
                  ancestor(X, Y) :- parent(X, Y).
                  ancestor(X, Z) :- parent(X, Y), ancestor(Y, Z).")))
          (let
            ((semi (dl-query db (list (quote ancestor) (quote a) (quote X))))
              (magic (dl-magic-query
                       db (list (quote ancestor) (quote a) (quote X)))))
            (= (len semi) (len magic))))
        true)

      ;; The magic rewriter passes aggregate body lits through
      ;; unchanged, so an aggregate over an IDB relation would see an
      ;; empty inner-goal in the magic db unless the IDB is already
      ;; materialised. dl-magic-query now pre-saturates the source db
      ;; to guarantee equivalence with dl-query for every stratified
      ;; program. Previously this returned `({:N 0})` because `active`
      ;; (IDB, derived through negation) was never derived in the
      ;; magic db.
      (dl-mt-test! "magic over aggregate-of-IDB matches vanilla"
        (let
          ((src
             "u(a). u(b). u(c). u(d). banned(b). banned(d).
              active(X) :- u(X), not(banned(X)).
              n(N) :- count(N, X, active(X))."))
          (let
            ((vanilla (dl-eval src "?- n(N)."))
              (magic (dl-eval-magic src "?- n(N).")))
            (and (= (len vanilla) 1)
                 (= (len magic) 1)
                 (= (get (first vanilla) "N")
                    (get (first magic) "N")))))
        true)

      ;; magic-query doesn't mutate caller db.
      (dl-mt-test! "magic-query preserves caller db"
        (let
          ((db (dl-program
                 "parent(a, b). parent(b, c).
                  ancestor(X, Y) :- parent(X, Y).
                  ancestor(X, Z) :- parent(X, Y), ancestor(Y, Z).")))
          (let
            ((rules-before (len (dl-rules db))))
            (do
              (dl-magic-query db (list (quote ancestor) (quote a) (quote X)))
              (= rules-before (len (dl-rules db))))))
        true)

      ;; Magic-sets benefit: query touches only one cluster of a
      ;; multi-component graph. Semi-naive derives the full closure
      ;; over both clusters; magic only the seeded one.
      ;; Magic-vs-semi work shape: chain of 12. Semi-naive
      ;; derives the full closure (78 = 12·13/2). A magic query
      ;; rooted at node 0 returns the 12 descendants only —
      ;; demonstrating that magic limits derivation to the
      ;; query's transitive cone.
      (dl-mt-test! "magic vs semi work-shape on chain-12"
        (let
          ((source (str
                     "parent(0, 1). parent(1, 2). parent(2, 3). "
                     "parent(3, 4). parent(4, 5). parent(5, 6). "
                     "parent(6, 7). parent(7, 8). parent(8, 9). "
                     "parent(9, 10). parent(10, 11). parent(11, 12). "
                     "ancestor(X, Y) :- parent(X, Y). "
                     "ancestor(X, Z) :- parent(X, Y), ancestor(Y, Z).")))
          (let
            ((db1 (dl-make-db)) (db2 (dl-make-db)))
            (do
              (dl-load-program! db1 source)
              (dl-saturate! db1)
              (dl-load-program! db2 source)
              (let
                ((semi-count (len (dl-relation db1 "ancestor")))
                  (magic-count
                    (len (dl-magic-query
                           db2 (list (quote ancestor) 0 (quote X))))))
                ;; Magic returns only descendants of 0 (12 of them).
                (and (= semi-count 78) (= magic-count 12))))))
        true)

      ;; Magic + arithmetic: rules with `is` clauses pass through
      ;; the rewriter unchanged (built-ins aren't propagated).
      (dl-mt-test! "magic preserves arithmetic"
        (let
          ((source "n(1). n(2). n(3).
                    doubled(X, Y) :- n(X), is(Y, *(X, 2))."))
          (let
            ((semi (dl-eval source "?- doubled(2, Y)."))
              (magic (dl-eval-magic source "?- doubled(2, Y).")))
            (= (len semi) (len magic))))
        true)

      (dl-mt-test! "magic skips irrelevant clusters"
        (let
          ;; Two disjoint chains. Query is rooted in cluster 1.
          ((db (dl-program
                 "parent(a, b). parent(b, c).
                  parent(x, y). parent(y, z).
                  ancestor(X, Y) :- parent(X, Y).
                  ancestor(X, Z) :- parent(X, Y), ancestor(Y, Z).")))
          (do
            (dl-saturate! db)
            (let
              ((semi-count (len (dl-relation db "ancestor")))
                (magic-results
                  (dl-magic-query
                    db (list (quote ancestor) (quote a) (quote X)))))
              ;; Semi-naive derives 6 (3 in each cluster). Magic
              ;; gives 3 query results (a's reachable: b, c).
              (and (= semi-count 6) (= (len magic-results) 2)))))
        true)

      (dl-mt-test! "magic-rewritten finds same answers"
        (let
          ((rules
             (list
               {:head (list (quote ancestor) (quote X) (quote Y))
                :body (list (list (quote parent) (quote X) (quote Y)))}
               {:head (list (quote ancestor) (quote X) (quote Z))
                :body
                (list (list (quote parent) (quote X) (quote Y))
                      (list (quote ancestor) (quote Y) (quote Z)))}))
            (edb (list
                   (list (quote parent) (quote a) (quote b))
                   (list (quote parent) (quote b) (quote c)))))
          (let
            ((rewritten (dl-magic-rewrite rules "ancestor" "bf" (list (quote a))))
              (db (dl-make-db)))
            (do
              (for-each (fn (f) (dl-add-fact! db f)) edb)
              (dl-add-fact! db (get rewritten :seed))
              (for-each (fn (r) (dl-add-rule! db r)) (get rewritten :rules))
              (dl-saturate! db)
              (len (dl-query db (list (quote ancestor) (quote a) (quote X)))))))
        2))))

(define
  dl-magic-tests-run!
  (fn
    ()
    (do
      (set! dl-mt-pass 0)
      (set! dl-mt-fail 0)
      (set! dl-mt-failures (list))
      (dl-mt-run-all!)
      {:passed dl-mt-pass
       :failed dl-mt-fail
       :total (+ dl-mt-pass dl-mt-fail)
       :failures dl-mt-failures})))
