;; lib/datalog/eval.sx — fixpoint evaluator (naive + semi-naive).
;;
;; Two saturators are exposed:
;;   `dl-saturate-naive!` — re-joins each rule against the full DB every
;;     iteration. Reference implementation; useful for differential tests.
;;   `dl-saturate!`       — semi-naive default. Tracks per-relation delta
;;     sets and substitutes one positive body literal per rule with the
;;     delta of its relation, joining the rest against the previous-
;;     iteration DB. Same fixpoint, dramatically less work on recursive
;;     rules.
;;
;; Body literal kinds:
;;   positive (rel arg ... arg)  → match against EDB+IDB tuples
;;   built-in (< X Y), (is X e)  → constraint via dl-eval-builtin
;;   negation {:neg lit}         → Phase 7

(define
  dl-match-positive
  (fn
    (lit db subst)
    (let
      ((rel (dl-rel-name lit)) (results (list)))
      (cond
        ((nil? rel) (error (str "dl-match-positive: bad literal " lit)))
        (else
          (let
            ;; If the first argument walks to a non-variable (constant
            ;; or already-bound var), use the first-arg index for
            ;; this relation. Otherwise scan the full tuple list.
            ((tuples
               (cond
                 ((>= (len lit) 2)
                   (let ((walked (dl-walk (nth lit 1) subst)))
                     (cond
                       ((dl-var? walked) (dl-rel-tuples db rel))
                       (else (dl-index-lookup db rel walked)))))
                 (else (dl-rel-tuples db rel)))))
            (do
              (for-each
                (fn
                  (tuple)
                  (let
                    ((s (dl-unify lit tuple subst)))
                    (when (not (nil? s)) (append! results s))))
                tuples)
              results)))))))

;; Match a positive literal against the delta set for its relation only.
(define
  dl-match-positive-delta
  (fn
    (lit delta subst)
    (let
      ((rel (dl-rel-name lit)) (results (list)))
      (let
        ((tuples (if (has-key? delta rel) (get delta rel) (list))))
        (do
          (for-each
            (fn
              (tuple)
              (let
                ((s (dl-unify lit tuple subst)))
                (when (not (nil? s)) (append! results s))))
            tuples)
          results)))))

;; Naive matcher (for dl-saturate-naive! and dl-query post-saturation).
(define
  dl-match-negation
  (fn
    (inner db subst)
    (let
      ((walked (dl-apply-subst inner subst))
        (matches (dl-match-positive inner db subst)))
      (cond
        ((= (len matches) 0) (list subst))
        (else (list))))))

(define
  dl-match-lit
  (fn
    (lit db subst)
    (cond
      ((and (dict? lit) (has-key? lit :neg))
        (dl-match-negation (get lit :neg) db subst))
      ((dl-aggregate? lit) (dl-eval-aggregate lit db subst))
      ((dl-builtin? lit)
        (let
          ((s (dl-eval-builtin lit subst)))
          (if (nil? s) (list) (list s))))
      ((and (list? lit) (> (len lit) 0))
        (dl-match-positive lit db subst))
      (else (error (str "datalog: unknown body-literal shape: " lit))))))

(define
  dl-find-bindings
  (fn (lits db subst) (dl-fb-aux lits db subst 0 (len lits))))

(define
  dl-fb-aux
  (fn
    (lits db subst i n)
    (cond
      ((nil? subst) (list))
      ((>= i n) (list subst))
      (else
        (let
          ((options (dl-match-lit (nth lits i) db subst))
            (results (list)))
          (do
            (for-each
              (fn
                (s)
                (for-each
                  (fn (s2) (append! results s2))
                  (dl-fb-aux lits db s (+ i 1) n)))
              options)
            results))))))

;; Naive: apply each rule against full DB until no new tuples.
(define
  dl-apply-rule!
  (fn
    (db rule)
    (let
      ((head (get rule :head)) (body (get rule :body)) (new? false))
      (do
        (for-each
          (fn
            (s)
            (let
              ((derived (dl-apply-subst head s)))
              (when (dl-add-derived! db derived) (set! new? true))))
          (dl-find-bindings body db (dl-empty-subst)))
        new?))))

;; Returns true iff one more saturation step would derive no new
;; tuples (i.e. the db is at fixpoint). Useful in tests that want
;; to assert "no work left" after a saturation call. Works under
;; either saturator since both compute the same fixpoint.
(define
  dl-saturated?
  (fn
    (db)
    (let ((any-new false))
      (do
        (for-each
          (fn
            (rule)
            (when (not any-new)
              (for-each
                (fn
                  (s)
                  (let ((derived (dl-apply-subst (get rule :head) s)))
                    (when
                      (and (not any-new)
                           (not (dl-tuple-member?
                                  derived
                                  (dl-rel-tuples
                                    db (dl-rel-name derived)))))
                      (set! any-new true))))
                (dl-find-bindings (get rule :body) db (dl-empty-subst)))))
          (dl-rules db))
        (not any-new)))))

(define
  dl-saturate-naive!
  (fn
    (db)
    (let
      ((changed true))
      (do
        (define
          dl-snloop
          (fn
            ()
            (when
              changed
              (do
                (set! changed false)
                (for-each
                  (fn (r) (when (dl-apply-rule! db r) (set! changed true)))
                  (dl-rules db))
                (dl-snloop)))))
        (dl-snloop)
        db))))

;; ── Semi-naive ───────────────────────────────────────────────────

;; Take a snapshot dict {rel -> tuples} of every relation currently in
;; the DB. Used as initial delta for the first iteration.
(define
  dl-snapshot-facts
  (fn
    (db)
    (let
      ((facts (get db :facts)) (out {}))
      (do
        (for-each
          (fn (k) (dict-set! out k (dl-copy-list (get facts k))))
          (keys facts))
        out))))

(define
  dl-copy-list
  (fn
    (xs)
    (let
      ((out (list)))
      (do (for-each (fn (x) (append! out x)) xs) out))))

;; Does any relation in `delta` have ≥1 tuple?
(define
  dl-delta-empty?
  (fn
    (delta)
    (let
      ((ks (keys delta)) (any-non-empty false))
      (do
        (for-each
          (fn
            (k)
            (when
              (> (len (get delta k)) 0)
              (set! any-non-empty true)))
          ks)
        (not any-non-empty)))))

;; Find substitutions such that `lits` are all satisfied AND `delta-idx`
;; is matched against the per-relation delta only. The other positive
;; literals match against the snapshot DB (db.facts read at iteration
;; start). Built-ins and negations behave as in `dl-match-lit`.
(define
  dl-find-bindings-semi
  (fn
    (lits db delta delta-idx subst)
    (dl-fbs-aux lits db delta delta-idx 0 subst)))

(define
  dl-fbs-aux
  (fn
    (lits db delta delta-idx i subst)
    (cond
      ((nil? subst) (list))
      ((>= i (len lits)) (list subst))
      (else
        (let
          ((lit (nth lits i))
            (options
              (cond
                ((and (dict? lit) (has-key? lit :neg))
                  (dl-match-negation (get lit :neg) db subst))
                ((dl-aggregate? lit) (dl-eval-aggregate lit db subst))
                ((dl-builtin? lit)
                  (let
                    ((s (dl-eval-builtin lit subst)))
                    (if (nil? s) (list) (list s))))
                ((and (list? lit) (> (len lit) 0))
                  (if
                    (= i delta-idx)
                    (dl-match-positive-delta lit delta subst)
                    (dl-match-positive lit db subst)))
                (else (error (str "datalog: unknown body-lit: " lit)))))
            (results (list)))
          (do
            (for-each
              (fn
                (s)
                (for-each
                  (fn (s2) (append! results s2))
                  (dl-fbs-aux lits db delta delta-idx (+ i 1) s)))
              options)
            results))))))

;; Collect candidate head tuples from a rule using delta. Walks every
;; positive body position and unions the resulting heads. For rules
;; with no positive body literal, falls back to a naive single-pass
;; (so static facts like `(p X) :- (= X 5).` derive on iteration 1).
(define
  dl-collect-rule-candidates
  (fn
    (rule db delta)
    (let
      ((head (get rule :head))
        (body (get rule :body))
        (out (list))
        (saw-pos false))
      (do
        (define
          dl-cri
          (fn
            (i)
            (when
              (< i (len body))
              (do
                (let
                  ((lit (nth body i)))
                  (when
                    (dl-positive-lit? lit)
                    (do
                      (set! saw-pos true)
                      (for-each
                        (fn (s) (append! out (dl-apply-subst head s)))
                        (dl-find-bindings-semi
                          body
                          db
                          delta
                          i
                          (dl-empty-subst))))))
                (dl-cri (+ i 1))))))
        (dl-cri 0)
        (when
          (not saw-pos)
          (for-each
            (fn (s) (append! out (dl-apply-subst head s)))
            (dl-find-bindings body db (dl-empty-subst))))
        out))))

;; Add a list of candidate tuples to db; collect newly-added ones into
;; the new-delta dict (keyed by relation name).
(define
  dl-commit-candidates!
  (fn
    (db candidates new-delta)
    (for-each
      (fn
        (lit)
        (when
          (dl-add-derived! db lit)
          (let
            ((rel (dl-rel-name lit)))
            (do
              (when
                (not (has-key? new-delta rel))
                (dict-set! new-delta rel (list)))
              (append! (get new-delta rel) lit)))))
      candidates)))

(define
  dl-saturate-rules!
  (fn
    (db rules)
    (let
      ((delta (dl-snapshot-facts db)))
      (do
        (define
          dl-sr-step
          (fn
            ()
            (let
              ((pending (list)) (new-delta {}))
              (do
                (for-each
                  (fn
                    (rule)
                    (for-each
                      (fn (cand) (append! pending cand))
                      (dl-collect-rule-candidates rule db delta)))
                  rules)
                (dl-commit-candidates! db pending new-delta)
                (cond
                  ((dl-delta-empty? new-delta) nil)
                  (else (do (set! delta new-delta) (dl-sr-step))))))))
        (dl-sr-step)
        db))))

;; Stratified driver: rejects non-stratifiable programs at saturation
;; time, then iterates strata in increasing order, running semi-naive on
;; the rules whose head sits in that stratum.
(define
  dl-saturate!
  (fn
    (db)
    (let
      ((err (dl-check-stratifiable db)))
      (cond
        ((not (nil? err)) (error (str "dl-saturate!: " err)))
        (else
          (let
            ((strata (dl-compute-strata db)))
            (let
              ((grouped (dl-group-rules-by-stratum db strata)))
              (let
                ((groups (get grouped :groups))
                  (max-s (get grouped :max)))
                (do
                  (define
                    dl-strat-loop
                    (fn
                      (s)
                      (when
                        (<= s max-s)
                        (let
                          ((sk (str s)))
                          (do
                            (when
                              (has-key? groups sk)
                              (dl-saturate-rules! db (get groups sk)))
                            (dl-strat-loop (+ s 1)))))))
                  (dl-strat-loop 0)
                  db)))))))))

;; ── Querying ─────────────────────────────────────────────────────

;; Coerce a query argument to a list of body literals. A single literal
;; like `(p X)` (positive — head is a symbol) or `{:neg ...}` becomes
;; `((p X))`. A list of literals like `((p X) (q X))` is returned as-is.
(define
  dl-query-coerce
  (fn
    (goal)
    (cond
      ((and (dict? goal) (has-key? goal :neg)) (list goal))
      ((and (list? goal) (> (len goal) 0) (symbol? (first goal)))
        (list goal))
      ((list? goal) goal)
      (else (error (str "dl-query: unrecognised goal shape: " goal))))))

(define
  dl-query
  (fn
    (db goal)
    (do
      (dl-saturate! db)
      ;; Rename anonymous '_' vars in each goal literal so multiple
      ;; occurrences do not unify together. Keep the user-facing var
      ;; list (taken before renaming) so projected results retain user
      ;; names.
      (let
        ((goals (dl-query-coerce goal))
          ;; Start the renamer past any `_anon<N>` symbols the user
          ;; may have written in the query — avoids collision.
          (renamer
            (dl-make-anon-renamer (dl-max-anon-num-list goal 0 0))))
        (let
          ((user-vars (dl-query-user-vars goals))
            (renamed (map (fn (g) (dl-rename-anon-lit g renamer)) goals)))
          (let
            ((substs (dl-find-bindings renamed db (dl-empty-subst)))
              (results (list)))
            (do
              (for-each
                (fn
                  (s)
                  (let
                    ((proj (dl-project-subst s user-vars)))
                    (when
                      (not (dl-tuple-member? proj results))
                      (append! results proj))))
                substs)
              results)))))))

(define
  dl-query-user-vars
  (fn
    (goals)
    (let ((seen (list)))
      (do
        (for-each
          (fn
            (g)
            (cond
              ((and (dict? g) (has-key? g :neg))
                (for-each
                  (fn
                    (v)
                    (when
                      (and (not (= v "_")) (not (dl-member-string? v seen)))
                      (append! seen v)))
                  (dl-vars-of (get g :neg))))
              ((dl-aggregate? g)
                ;; Only the result var (first arg of the aggregate
                ;; literal) is user-facing. The aggregated var and
                ;; any vars in the inner goal are internal.
                (let ((r (nth g 1)))
                  (when
                    (dl-var? r)
                    (let ((rn (symbol->string r)))
                      (when
                        (and (not (= rn "_"))
                             (not (dl-member-string? rn seen)))
                        (append! seen rn))))))
              (else
                (for-each
                  (fn
                    (v)
                    (when
                      (and (not (= v "_")) (not (dl-member-string? v seen)))
                      (append! seen v)))
                  (dl-vars-of g)))))
          goals)
        seen))))

(define
  dl-project-subst
  (fn
    (subst names)
    (let
      ((out {}))
      (do
        (for-each
          (fn
            (n)
            (let
              ((sym (string->symbol n)))
              (let
                ((v (dl-walk sym subst)))
                (dict-set! out n (dl-apply-subst v subst)))))
          names)
        out))))

(define dl-relation (fn (db name) (dl-rel-tuples db name)))
