;; lib/datalog/db.sx — Datalog database (EDB + IDB + rules) + safety.
;;
;; A db is a mutable dict:
;;   {:facts {<rel-name-string> -> (literal ...)}
;;    :rules ({:head literal :body (literal ...)} ...)}
;;
;; Facts are stored as full literals `(rel arg ... arg)` so they unify
;; directly against rule body literals. Each relation's tuple list is
;; deduplicated on insert.
;;
;; Phase 3 makes no distinction between EDB (asserted) and IDB (derived);
;; a future phase may track provenance for retraction.

(define dl-make-db (fn () {:facts {} :rules (list)}))

;; The relation name (as string) of a positive literal, or of the
;; underlying literal of a negative one. nil for malformed input.
(define
  dl-rel-name
  (fn
    (lit)
    (cond
      ((and (dict? lit) (has-key? lit :neg)) (dl-rel-name (get lit :neg)))
      ((and (list? lit) (> (len lit) 0) (symbol? (first lit)))
        (symbol->string (first lit)))
      (else nil))))

(define dl-builtin-rels (list "<" "<=" ">" ">=" "=" "!=" "is"))

(define
  dl-member-string?
  (fn
    (s xs)
    (cond
      ((= (len xs) 0) false)
      ((= (first xs) s) true)
      (else (dl-member-string? s (rest xs))))))

(define
  dl-builtin?
  (fn
    (lit)
    (and
      (list? lit)
      (> (len lit) 0)
      (let
        ((rel (dl-rel-name lit)))
        (cond
          ((nil? rel) false)
          (else (dl-member-string? rel dl-builtin-rels)))))))

(define
  dl-positive-lit?
  (fn
    (lit)
    (cond
      ((and (dict? lit) (has-key? lit :neg)) false)
      ((dl-builtin? lit) false)
      ((and (list? lit) (> (len lit) 0)) true)
      (else false))))

;; Membership using dl-deep tuple equality (handles var/constant symbols
;; and numbers consistently).
(define
  dl-tuple-equal?
  (fn
    (a b)
    (cond
      ((and (list? a) (list? b))
        (and (= (len a) (len b)) (dl-tuple-equal-list? a b 0)))
      ((and (number? a) (number? b)) (= a b))
      (else (equal? a b)))))

(define
  dl-tuple-equal-list?
  (fn
    (a b i)
    (cond
      ((>= i (len a)) true)
      ((not (dl-tuple-equal? (nth a i) (nth b i))) false)
      (else (dl-tuple-equal-list? a b (+ i 1))))))

(define
  dl-tuple-member?
  (fn
    (lit lits)
    (cond
      ((= (len lits) 0) false)
      ((dl-tuple-equal? lit (first lits)) true)
      (else (dl-tuple-member? lit (rest lits))))))

;; Ensure :facts has a list for the given relation key, then return it.
(define
  dl-ensure-rel!
  (fn
    (db rel-key)
    (let
      ((facts (get db :facts)))
      (do
        (when
          (not (has-key? facts rel-key))
          (dict-set! facts rel-key (list)))
        (get facts rel-key)))))

;; All tuples currently known for a relation (EDB ∪ IDB). Returns empty
;; list when relation hasn't been seen.
(define
  dl-rel-tuples
  (fn
    (db rel-key)
    (let
      ((facts (get db :facts)))
      (if (has-key? facts rel-key) (get facts rel-key) (list)))))

;; Add a ground literal. Returns true iff the literal was new.
(define
  dl-add-fact!
  (fn
    (db lit)
    (cond
      ((not (and (list? lit) (> (len lit) 0)))
        (error (str "dl-add-fact!: expected literal list, got " lit)))
      ((not (dl-ground? lit (dl-empty-subst)))
        (error (str "dl-add-fact!: expected ground literal, got " lit)))
      (else
        (let
          ((rel-key (dl-rel-name lit)))
          (let
            ((tuples (dl-ensure-rel! db rel-key)))
            (cond
              ((dl-tuple-member? lit tuples) false)
              (else (do (append! tuples lit) true)))))))))

;; Collect variables appearing in the positive body literals of `body`.
(define
  dl-positive-body-vars
  (fn
    (body)
    (let
      ((vars (list)))
      (do
        (for-each
          (fn
            (lit)
            (when
              (dl-positive-lit? lit)
              (for-each
                (fn
                  (v)
                  (when (not (dl-member-string? v vars)) (append! vars v)))
                (dl-vars-of lit))))
          body)
        vars))))

;; Collect variables in any literal (positive, negated, or built-in).
(define
  dl-all-body-vars
  (fn
    (body)
    (let
      ((vars (list)))
      (do
        (for-each
          (fn
            (lit)
            (let
              ((target (if (and (dict? lit) (has-key? lit :neg)) (get lit :neg) lit)))
              (for-each
                (fn
                  (v)
                  (when (not (dl-member-string? v vars)) (append! vars v)))
                (dl-vars-of target))))
          body)
        vars))))

;; Return the list of head variables NOT covered by some positive body
;; literal. A safe rule has an empty list. The check ignores '_' since
;; that is treated as a fresh anonymous variable per occurrence.
(define
  dl-rule-unsafe-head-vars
  (fn
    (rule)
    (let
      ((head (get rule :head))
        (body (get rule :body))
        (head-vars (dl-vars-of head))
        (body-vars (dl-positive-body-vars body))
        (out (list)))
      (do
        (for-each
          (fn
            (v)
            (when
              (and (not (dl-member-string? v body-vars)) (not (= v "_")))
              (append! out v)))
          head-vars)
        out))))

;; Add a rule. Rejects unsafe rules with a pointer at the offending var.
(define
  dl-add-rule!
  (fn
    (db rule)
    (cond
      ((not (dict? rule))
        (error (str "dl-add-rule!: expected rule dict, got " rule)))
      ((not (has-key? rule :head))
        (error (str "dl-add-rule!: rule missing :head, got " rule)))
      (else
        (let
          ((unsafe (dl-rule-unsafe-head-vars rule)))
          (cond
            ((> (len unsafe) 0)
              (error
                (str
                  "dl-add-rule!: unsafe rule — head variable(s) "
                  unsafe
                  " do not appear in any positive body literal of "
                  rule)))
            (else
              (let
                ((rules (get db :rules)))
                (do (append! rules rule) true)))))))))

;; Load a parsed clause (fact or rule) into the db. Queries are ignored.
(define
  dl-add-clause!
  (fn
    (db clause)
    (cond
      ((has-key? clause :query) false)
      ((and (has-key? clause :body) (= (len (get clause :body)) 0))
        (dl-add-fact! db (get clause :head)))
      (else (dl-add-rule! db clause)))))

;; Parse source text and load every clause into the db. Returns the db.
(define
  dl-load-program!
  (fn
    (db source)
    (let
      ((clauses (dl-parse source)))
      (do (for-each (fn (c) (dl-add-clause! db c)) clauses) db))))

;; Convenience: build a db from source in one step.
(define
  dl-program
  (fn (source) (let ((db (dl-make-db))) (dl-load-program! db source))))

(define dl-rules (fn (db) (get db :rules)))

;; Total number of stored ground tuples across all relations.
(define
  dl-fact-count
  (fn
    (db)
    (let
      ((facts (get db :facts)) (total 0))
      (do
        (for-each
          (fn (k) (set! total (+ total (len (get facts k)))))
          (keys facts))
        total))))