;; lib/datalog/magic.sx — adornment analysis + sideways info passing.
;;
;; First step of the magic-sets transformation (Phase 6). Right now
;; the saturator does not consume these — they are introspection
;; helpers that future magic-set rewriting will build on top of.
;;
;; Definitions:
;;   - An *adornment* of an n-ary literal is an n-character string
;;     of "b" (bound — value already known at the call site) and
;;     "f" (free — to be derived).
;;   - SIPS (Sideways Information Passing Strategy) walks the body
;;     of an adorned rule left-to-right tracking which variables
;;     have been bound so far, computing each body literal's
;;     adornment in turn.
;;
;; Usage:
;;
;;   (dl-adorn-goal '(ancestor tom X))
;;     => "bf"
;;
;;   (dl-rule-sips
;;     {:head (ancestor X Z)
;;      :body ((parent X Y) (ancestor Y Z))}
;;     "bf")
;;   => ({:lit (parent X Y) :adornment "bf"}
;;       {:lit (ancestor Y Z) :adornment "bf"})

;; Per-arg adornment under the current bound-var name set.
(define
  dl-adorn-arg
  (fn
    (arg bound)
    (cond
      ((dl-var? arg)
        (if (dl-member-string? (symbol->string arg) bound) "b" "f"))
      (else "b"))))

;; Adornment for the args of a literal (after the relation name).
(define
  dl-adorn-args
  (fn
    (args bound)
    (cond
      ((= (len args) 0) "")
      (else
        (str
          (dl-adorn-arg (first args) bound)
          (dl-adorn-args (rest args) bound))))))

;; Adornment of a top-level goal under the empty bound-var set.
(define
  dl-adorn-goal
  (fn (goal) (dl-adorn-args (rest goal) (list))))

;; Adornment of a literal under an explicit bound set.
(define
  dl-adorn-lit
  (fn (lit bound) (dl-adorn-args (rest lit) bound)))

;; The set of variable names made bound by walking a positive
;; literal whose adornment is known. Free positions add their
;; vars to the bound set.
(define
  dl-vars-bound-by-lit
  (fn
    (lit bound)
    (let ((args (rest lit)) (out (list)))
      (do
        (for-each
          (fn (a)
            (when
              (and (dl-var? a)
                   (not (dl-member-string? (symbol->string a) bound))
                   (not (dl-member-string? (symbol->string a) out)))
              (append! out (symbol->string a))))
          args)
        out))))

;; Walk the rule body left-to-right tracking bound vars seeded by the
;; head adornment. Returns a list of {:lit :adornment} entries.
;;
;; Negation, comparison, and built-ins are passed through with their
;; adornment computed from the current bound set; they don't add new
;; bindings (except `is`, which binds its left arg if a var). Aggregates
;; are treated like is — the result var becomes bound.
(define
  dl-init-head-bound
  (fn
    (head adornment)
    (let ((args (rest head)) (out (list)))
      (do
        (define
          dl-ihb-loop
          (fn
            (i)
            (when
              (< i (len args))
              (do
                (let
                  ((c (slice adornment i (+ i 1)))
                    (a (nth args i)))
                  (when
                    (and (= c "b") (dl-var? a))
                    (let ((n (symbol->string a)))
                      (when
                        (not (dl-member-string? n out))
                        (append! out n)))))
                (dl-ihb-loop (+ i 1))))))
        (dl-ihb-loop 0)
        out))))

(define
  dl-rule-sips
  (fn
    (rule head-adornment)
    (let
      ((bound (dl-init-head-bound (get rule :head) head-adornment))
        (out (list)))
      (do
        (for-each
          (fn
            (lit)
            (cond
              ((and (dict? lit) (has-key? lit :neg))
                (let ((target (get lit :neg)))
                  (append!
                    out
                    {:lit lit :adornment (dl-adorn-lit target bound)})))
              ((dl-builtin? lit)
                (let ((adn (dl-adorn-lit lit bound)))
                  (do
                    (append! out {:lit lit :adornment adn})
                    ;; `is` binds its left arg (if var) once RHS is ground.
                    (when
                      (and (= (dl-rel-name lit) "is") (dl-var? (nth lit 1)))
                      (let ((n (symbol->string (nth lit 1))))
                        (when
                          (not (dl-member-string? n bound))
                          (append! bound n)))))))
              ((and (list? lit) (dl-aggregate? lit))
                (let ((adn (dl-adorn-lit lit bound)))
                  (do
                    (append! out {:lit lit :adornment adn})
                    ;; Result var (first arg) becomes bound.
                    (when (dl-var? (nth lit 1))
                      (let ((n (symbol->string (nth lit 1))))
                        (when
                          (not (dl-member-string? n bound))
                          (append! bound n)))))))
              ((and (list? lit) (> (len lit) 0))
                (let ((adn (dl-adorn-lit lit bound)))
                  (do
                    (append! out {:lit lit :adornment adn})
                    (for-each
                      (fn (n)
                        (when (not (dl-member-string? n bound))
                          (append! bound n)))
                      (dl-vars-bound-by-lit lit bound)))))))
          (get rule :body))
        out))))

;; ── Magic predicate naming + bound-args extraction ─────────────
;; These are building blocks for the magic-sets *transformation*
;; itself. The transformation (which generates rewritten rules
;; with magic_<rel>^<adornment> filters) is future work — for now
;; these helpers can be used to inspect what such a transformation
;; would produce.

;; "magic_p^bf" given relation "p" and adornment "bf".
(define
  dl-magic-rel-name
  (fn (rel adornment) (str "magic_" rel "^" adornment)))

;; A magic predicate literal:
;;   (magic_<rel>^<adornment> arg1 arg2 ...)
(define
  dl-magic-lit
  (fn
    (rel adornment bound-args)
    (cons (string->symbol (dl-magic-rel-name rel adornment)) bound-args)))

;; Extract bound args (those at "b" positions in `adornment`) from a
;; literal `(rel arg1 arg2 ... argN)`. Returns the list of arg values.
(define
  dl-bound-args
  (fn
    (lit adornment)
    (let ((args (rest lit)) (out (list)))
      (do
        (define
          dl-ba-loop
          (fn
            (i)
            (when
              (< i (len args))
              (do
                (when
                  (= (slice adornment i (+ i 1)) "b")
                  (append! out (nth args i)))
                (dl-ba-loop (+ i 1))))))
        (dl-ba-loop 0)
        out))))

;; ── Magic-sets rewriter ─────────────────────────────────────────
;;
;; Given the original rule list and a query (rel, adornment) pair,
;; generates the magic-rewritten program: a list of rules that
;; (a) gate each original rule with a `magic_<rel>^<adn>` filter and
;; (b) propagate the magic relation through SIPS so that only
;; query-relevant tuples are derived. Seed facts are returned
;; separately and must be added to the db at evaluation time.
;;
;; Output: {:rules <rewritten-rules> :seed <magic-seed-literal>}
;;
;; The rewriter only rewrites IDB rules; EDB facts pass through.
;; Built-in predicates and negation in body literals are kept in
;; place but do not generate propagation rules of their own.

(define
  dl-magic-pair-key
  (fn (rel adornment) (str rel "^" adornment)))

(define
  dl-magic-rewrite
  (fn
    (rules query-rel query-adornment query-args)
    (let
      ((seen (list))
        (queue (list))
        (out (list)))
      (do
        (define
          dl-mq-mark!
          (fn
            (rel adornment)
            (let ((k (dl-magic-pair-key rel adornment)))
              (when
                (not (dl-member-string? k seen))
                (do
                  (append! seen k)
                  (append! queue {:rel rel :adn adornment}))))))

        (define
          dl-mq-rewrite-rule!
          (fn
            (rule adn)
            (let
              ((head (get rule :head))
                (body (get rule :body))
                (sips (dl-rule-sips rule adn)))
              (let
                ((magic-filter
                   (dl-magic-lit
                     (dl-rel-name head)
                     adn
                     (dl-bound-args head adn))))
                (do
                  ;; Adorned rule: head :- magic-filter, body...
                  (let ((new-body (list)))
                    (do
                      (append! new-body magic-filter)
                      (for-each
                        (fn (lit) (append! new-body lit))
                        body)
                      (append! out {:head head :body new-body})))
                  ;; Propagation rules for each positive non-builtin
                  ;; body literal at position i.
                  (define
                    dl-mq-prop-loop
                    (fn
                      (i)
                      (when
                        (< i (len body))
                        (do
                          (let
                            ((lit (nth body i))
                              (sip-entry (nth sips i)))
                            (when
                              (and (list? lit)
                                   (> (len lit) 0)
                                   (not (and (dict? lit) (has-key? lit :neg)))
                                   (not (dl-builtin? lit))
                                   (not (dl-aggregate? lit)))
                              (let
                                ((lit-adn (get sip-entry :adornment))
                                  (lit-rel (dl-rel-name lit)))
                                (let
                                  ((prop-head
                                     (dl-magic-lit
                                       lit-rel
                                       lit-adn
                                       (dl-bound-args lit lit-adn))))
                                  (let ((prop-body (list)))
                                    (do
                                      (append! prop-body magic-filter)
                                      (define
                                        dl-mq-prefix-loop
                                        (fn
                                          (j)
                                          (when
                                            (< j i)
                                            (do
                                              (append!
                                                prop-body
                                                (nth body j))
                                              (dl-mq-prefix-loop (+ j 1))))))
                                      (dl-mq-prefix-loop 0)
                                      (append!
                                        out
                                        {:head prop-head :body prop-body})
                                      (dl-mq-mark! lit-rel lit-adn)))))))
                          (dl-mq-prop-loop (+ i 1))))))
                  (dl-mq-prop-loop 0))))))

        (dl-mq-mark! query-rel query-adornment)

        (define
          dl-mq-process
          (fn
            ()
            (when
              (> (len queue) 0)
              (let ((item (first queue)))
                (do
                  (set! queue (rest queue))
                  (let
                    ((rel (get item :rel)) (adn (get item :adn)))
                    (for-each
                      (fn
                        (rule)
                        (when
                          (= (dl-rel-name (get rule :head)) rel)
                          (dl-mq-rewrite-rule! rule adn)))
                      rules))
                  (dl-mq-process))))))
        (dl-mq-process)

        {:rules out
         :seed
         (dl-magic-lit
           query-rel
           query-adornment
           query-args)}))))

;; ── Top-level magic-sets driver ─────────────────────────────────
;;
;; (dl-magic-query db query-goal) — run `query-goal` under magic-sets
;; evaluation. Builds a fresh internal db with:
;;   - the caller's EDB facts (relations not headed by any rule),
;;   - the magic seed fact, and
;;   - the rewritten rules.
;; Saturates and queries, returning the substitution list. The
;; caller's db is untouched.
;;
;; Useful primarily as a perf alternative for queries that only
;; need a small slice of a recursive relation. Equivalent to
;; dl-query for any single fully-stratifiable program.

(define
  dl-magic-rule-heads
  (fn
    (rules)
    (let ((seen (list)))
      (do
        (for-each
          (fn
            (r)
            (let ((h (dl-rel-name (get r :head))))
              (when
                (and (not (nil? h)) (not (dl-member-string? h seen)))
                (append! seen h))))
          rules)
        seen))))

(define
  dl-magic-query
  (fn
    (db query-goal)
    ;; Magic-sets only applies to positive non-builtin / non-aggregate
    ;; literals against rule-defined relations. For other goal shapes
    ;; (built-ins, aggregates, EDB-only relations) the seed is either
    ;; non-ground or unused; fall back to dl-query.
    (cond
      ((or (dl-builtin? query-goal)
           (dl-aggregate? query-goal)
           (and (dict? query-goal) (has-key? query-goal :neg)))
        (dl-query db query-goal))
      (else
        (let
          ((query-rel (dl-rel-name query-goal))
            (query-adn (dl-adorn-goal query-goal)))
          (let
            ((query-args (dl-bound-args query-goal query-adn))
              (rules (dl-rules db)))
            (let
              ((rewritten (dl-magic-rewrite rules query-rel query-adn query-args))
                (mdb (dl-make-db))
                (rule-heads (dl-magic-rule-heads rules)))
          (do
            ;; Copy ALL existing facts. EDB-only relations bring their
            ;; tuples; mixed EDB+IDB relations bring both their EDB
            ;; portion and any pre-saturated IDB tuples (which the
            ;; rewritten rules would re-derive anyway). Skipping facts
            ;; for rule-headed relations would leave the magic run
            ;; without the EDB portion of mixed relations.
            (for-each
              (fn
                (rel)
                (for-each
                  (fn (t) (dl-add-fact! mdb t))
                  (dl-rel-tuples db rel)))
              (keys (get db :facts)))
            ;; Seed + rewritten rules.
            (dl-add-fact! mdb (get rewritten :seed))
            (for-each (fn (r) (dl-add-rule! mdb r)) (get rewritten :rules))
            (dl-query mdb query-goal)))))))))