;; lib/datalog/strata.sx — dependency graph, SCC analysis, stratum assignment.
;;
;; A program is stratifiable iff no cycle in its dependency graph passes
;; through a negative edge. The stratum of relation R is the depth at which
;; R can first be evaluated:
;;
;;   stratum(R) = max over edges (R → S) of:
;;       stratum(S)         if the edge is positive
;;       stratum(S) + 1     if the edge is negative
;;
;; All relations in the same SCC share a stratum (and the SCC must have only
;; positive internal edges, else the program is non-stratifiable).

;; Build dep graph: dict {head-rel-name -> ({:rel str :neg bool} ...)}.
(define
  dl-build-dep-graph
  (fn
    (db)
    (let ((g {}))
      (do
        (for-each
          (fn
            (rule)
            (let
              ((head-rel (dl-rel-name (get rule :head))))
              (when
                (not (nil? head-rel))
                (do
                  (when
                    (not (has-key? g head-rel))
                    (dict-set! g head-rel (list)))
                  (let ((existing (get g head-rel)))
                    (for-each
                      (fn
                        (lit)
                        (cond
                          ((dl-aggregate? lit)
                            (let
                              ((edge (dl-aggregate-dep-edge lit)))
                              (when
                                (not (nil? edge))
                                (append! existing edge))))
                          (else
                            (let
                              ((target
                                 (cond
                                   ((and (dict? lit) (has-key? lit :neg))
                                     (dl-rel-name (get lit :neg)))
                                   ((dl-builtin? lit) nil)
                                   ((and (list? lit) (> (len lit) 0))
                                     (dl-rel-name lit))
                                   (else nil)))
                                (neg?
                                  (and (dict? lit) (has-key? lit :neg))))
                              (when
                                (not (nil? target))
                                (append!
                                  existing
                                  {:rel target :neg neg?}))))))
                      (get rule :body)))))))
          (dl-rules db))
        g))))

;; All relations referenced — heads of rules + EDB names + body relations.
(define
  dl-all-relations
  (fn
    (db)
    (let ((seen (list)))
      (do
        (for-each
          (fn
            (k)
            (when (not (dl-member-string? k seen)) (append! seen k)))
          (keys (get db :facts)))
        (for-each
          (fn
            (rule)
            (do
              (let ((h (dl-rel-name (get rule :head))))
                (when
                  (and (not (nil? h)) (not (dl-member-string? h seen)))
                  (append! seen h)))
              (for-each
                (fn
                  (lit)
                  (let
                    ((t
                       (cond
                         ((dl-aggregate? lit)
                           (let ((edge (dl-aggregate-dep-edge lit)))
                             (if (nil? edge) nil (get edge :rel))))
                         ((and (dict? lit) (has-key? lit :neg))
                           (dl-rel-name (get lit :neg)))
                         ((dl-builtin? lit) nil)
                         ((and (list? lit) (> (len lit) 0))
                           (dl-rel-name lit))
                         (else nil))))
                    (when
                      (and (not (nil? t)) (not (dl-member-string? t seen)))
                      (append! seen t))))
                (get rule :body))))
          (dl-rules db))
        seen))))

;; reach: dict {from: dict {to: edge-info}} where edge-info is
;;   {:any bool :neg bool}
;; meaning "any path from `from` to `to`" and "exists a negative-passing
;; path from `from` to `to`".
;;
;; Floyd-Warshall over the dep graph. The 'neg' flag propagates through
;; concatenation: if any edge along the path is negative, the path's
;; flag is true.
(define
  dl-build-reach
  (fn
    (graph nodes)
    (let ((reach {}))
      (do
        (for-each
          (fn (n) (dict-set! reach n {}))
          nodes)
        (for-each
          (fn
            (head)
            (when
              (has-key? graph head)
              (for-each
                (fn
                  (edge)
                  (let
                    ((target (get edge :rel)) (n (get edge :neg)))
                    (let ((row (get reach head)))
                      (cond
                        ((has-key? row target)
                          (let ((cur (get row target)))
                            (dict-set!
                              row
                              target
                              {:any true :neg (or n (get cur :neg))})))
                        (else
                          (dict-set! row target {:any true :neg n}))))))
                (get graph head))))
          nodes)
        (for-each
          (fn
            (k)
            (for-each
              (fn
                (i)
                (let ((row-i (get reach i)))
                  (when
                    (has-key? row-i k)
                    (let ((ik (get row-i k)) (row-k (get reach k)))
                      (for-each
                        (fn
                          (j)
                          (when
                            (has-key? row-k j)
                            (let ((kj (get row-k j)))
                              (let
                                ((combined-neg (or (get ik :neg) (get kj :neg))))
                                (cond
                                  ((has-key? row-i j)
                                    (let ((cur (get row-i j)))
                                      (dict-set!
                                        row-i
                                        j
                                        {:any true
                                         :neg (or combined-neg (get cur :neg))})))
                                  (else
                                    (dict-set!
                                      row-i
                                      j
                                      {:any true :neg combined-neg})))))))
                        nodes)))))
              nodes))
          nodes)
        reach))))

;; Returns nil on success, or error message string on failure.
(define
  dl-check-stratifiable
  (fn
    (db)
    (let
      ((graph (dl-build-dep-graph db))
        (nodes (dl-all-relations db)))
      (let ((reach (dl-build-reach graph nodes)) (err nil))
        (do
          (for-each
            (fn
              (rule)
              (when
                (nil? err)
                (let ((head-rel (dl-rel-name (get rule :head))))
                  (for-each
                    (fn
                      (lit)
                      (cond
                        ((and (dict? lit) (has-key? lit :neg))
                          (let ((tgt (dl-rel-name (get lit :neg))))
                            (when
                              (and (not (nil? tgt))
                                   (dl-reach-cycle? reach head-rel tgt))
                              (set!
                                err
                                (str "non-stratifiable: relation " head-rel
                                     " transitively depends through negation on "
                                     tgt
                                     " which depends back on " head-rel)))))
                        ((dl-aggregate? lit)
                          (let ((edge (dl-aggregate-dep-edge lit)))
                            (when
                              (not (nil? edge))
                              (let ((tgt (get edge :rel)))
                                (when
                                  (and (not (nil? tgt))
                                       (dl-reach-cycle? reach head-rel tgt))
                                  (set!
                                    err
                                    (str "non-stratifiable: relation "
                                         head-rel
                                         " aggregates over " tgt
                                         " which depends back on "
                                         head-rel))))))))) 
                    (get rule :body)))))
            (dl-rules db))
          err)))))

(define
  dl-reach-cycle?
  (fn
    (reach a b)
    (and
      (dl-reach-row-has? reach b a)
      (dl-reach-row-has? reach a b))))

(define
  dl-reach-row-has?
  (fn
    (reach from to)
    (let ((row (get reach from)))
      (and (not (nil? row)) (has-key? row to)))))

;; Compute stratum per relation. Iteratively propagate from EDB roots.
;; Uses the per-relation max-stratum-of-deps formula. Stops when stable.
(define
  dl-compute-strata
  (fn
    (db)
    (let
      ((graph (dl-build-dep-graph db))
        (nodes (dl-all-relations db))
        (strata {}))
      (do
        (for-each (fn (n) (dict-set! strata n 0)) nodes)
        (let ((changed true))
          (do
            (define
              dl-cs-loop
              (fn
                ()
                (when
                  changed
                  (do
                    (set! changed false)
                    (for-each
                      (fn
                        (head)
                        (when
                          (has-key? graph head)
                          (for-each
                            (fn
                              (edge)
                              (let
                                ((tgt (get edge :rel))
                                  (n (get edge :neg)))
                                (let
                                  ((tgt-strat
                                     (if (has-key? strata tgt)
                                       (get strata tgt) 0))
                                    (cur (get strata head)))
                                  (let
                                    ((needed
                                       (if n (+ tgt-strat 1) tgt-strat)))
                                    (when
                                      (> needed cur)
                                      (do
                                        (dict-set! strata head needed)
                                        (set! changed true)))))))
                            (get graph head))))
                      nodes)
                    (dl-cs-loop)))))
            (dl-cs-loop)))
        strata))))

;; Group rules by their head's stratum. Returns dict {stratum-int -> rules}.
(define
  dl-group-rules-by-stratum
  (fn
    (db strata)
    (let ((groups {}) (max-s 0))
      (do
        (for-each
          (fn
            (rule)
            (let
              ((head-rel (dl-rel-name (get rule :head))))
              (let
                ((s (if (has-key? strata head-rel)
                       (get strata head-rel) 0)))
                (do
                  (when (> s max-s) (set! max-s s))
                  (let
                    ((sk (str s)))
                    (do
                      (when
                        (not (has-key? groups sk))
                        (dict-set! groups sk (list)))
                      (append! (get groups sk) rule)))))))
          (dl-rules db))
        {:groups groups :max max-s}))))