;; lib/minikanren/clpfd.sx — Phase 6: native CLP(FD) on miniKanren.
;;
;; The substitution dict carries an extra reserved key "_fd" that holds a
;; constraint-store record:
;;
;;   {:domains    {var-name -> sorted-int-list}
;;    :constraints (... pending constraint closures ...)}
;;
;; Domains are sorted SX lists of ints (no duplicates).
;; Constraints are functions s -> s-or-nil that propagate / re-check.
;; They are re-fired after every label binding via fd-fire-store.

(define fd-key "_fd")

;; --- domain primitives ---

(define
  fd-dom-rev
  (fn
    (xs acc)
    (cond
      ((empty? xs) acc)
      (:else (fd-dom-rev (rest xs) (cons (first xs) acc))))))

(define
  fd-dom-insert
  (fn
    (x desc)
    (cond
      ((empty? desc) (list x))
      ((= x (first desc)) desc)
      ((> x (first desc)) (cons x desc))
      (:else (cons (first desc) (fd-dom-insert x (rest desc)))))))

(define
  fd-dom-sort-dedupe
  (fn
    (xs acc)
    (cond
      ((empty? xs) (fd-dom-rev acc (list)))
      (:else (fd-dom-sort-dedupe (rest xs) (fd-dom-insert (first xs) acc))))))

(define fd-dom-from-list (fn (xs) (fd-dom-sort-dedupe xs (list))))

(define fd-dom-empty? (fn (d) (empty? d)))
(define
  fd-dom-singleton?
  (fn (d) (and (not (empty? d)) (empty? (rest d)))))
(define fd-dom-min (fn (d) (first d)))

(define
  fd-dom-last
  (fn
    (d)
    (cond ((empty? (rest d)) (first d)) (:else (fd-dom-last (rest d))))))

(define fd-dom-max (fn (d) (fd-dom-last d)))
(define fd-dom-member? (fn (x d) (some (fn (y) (= x y)) d)))

(define
  fd-dom-intersect
  (fn
    (a b)
    (cond
      ((empty? a) (list))
      ((empty? b) (list))
      ((= (first a) (first b))
        (cons (first a) (fd-dom-intersect (rest a) (rest b))))
      ((< (first a) (first b)) (fd-dom-intersect (rest a) b))
      (:else (fd-dom-intersect a (rest b))))))

(define
  fd-dom-without
  (fn
    (x d)
    (cond
      ((empty? d) (list))
      ((= (first d) x) (rest d))
      ((> (first d) x) d)
      (:else (cons (first d) (fd-dom-without x (rest d)))))))

(define
  fd-dom-range
  (fn
    (lo hi)
    (cond
      ((> lo hi) (list))
      (:else (cons lo (fd-dom-range (+ lo 1) hi))))))

;; --- constraint store accessors ---

(define fd-store-empty (fn () {:domains {} :constraints (list)}))

(define
  fd-store-of
  (fn
    (s)
    (cond ((has-key? s fd-key) (get s fd-key)) (:else (fd-store-empty)))))

(define fd-domains-of (fn (s) (get (fd-store-of s) :domains)))
(define fd-with-store (fn (s store) (assoc s fd-key store)))

(define
  fd-domain-of
  (fn
    (s var-name)
    (let
      ((doms (fd-domains-of s)))
      (cond ((has-key? doms var-name) (get doms var-name)) (:else nil)))))

(define
  fd-set-domain
  (fn
    (s var-name d)
    (cond
      ((fd-dom-empty? d) nil)
      (:else
        (let
          ((store (fd-store-of s)))
          (let
            ((doms-prime (assoc (get store :domains) var-name d)))
            (let
              ((store-prime (assoc store :domains doms-prime)))
              (fd-with-store s store-prime))))))))

(define
  fd-add-constraint
  (fn
    (s c)
    (let
      ((store (fd-store-of s)))
      (let
        ((cs-prime (cons c (get store :constraints))))
        (let
          ((store-prime (assoc store :constraints cs-prime)))
          (fd-with-store s store-prime))))))

(define
  fd-fire-list
  (fn
    (cs s)
    (cond
      ((empty? cs) s)
      (:else
        (let
          ((s2 ((first cs) s)))
          (cond ((= s2 nil) nil) (:else (fd-fire-list (rest cs) s2))))))))

(define
  fd-fire-store
  (fn
    (s)
    (let ((cs (get (fd-store-of s) :constraints))) (fd-fire-list cs s))))

;; --- user-facing goals ---

(define
  fd-in
  (fn
    (x dom-list)
    (fn
      (s)
      (let
        ((new-dom (fd-dom-from-list dom-list)))
        (let
          ((wx (mk-walk x s)))
          (cond
            ((number? wx)
              (cond ((fd-dom-member? wx new-dom) (unit s)) (:else mzero)))
            ((is-var? wx)
              (let
                ((existing (fd-domain-of s (var-name wx))))
                (let
                  ((narrowed (cond ((= existing nil) new-dom) (:else (fd-dom-intersect existing new-dom)))))
                  (let
                    ((s2 (fd-set-domain s (var-name wx) narrowed)))
                    (cond ((= s2 nil) mzero) (:else (unit s2)))))))
            (:else mzero)))))))

;; --- fd-neq ---

(define
  fd-neq-prop
  (fn
    (x y s)
    (let
      ((wx (mk-walk x s)) (wy (mk-walk y s)))
      (cond
        ((and (number? wx) (number? wy))
          (cond ((= wx wy) nil) (:else s)))
        ((and (number? wx) (is-var? wy))
          (let
            ((y-dom (fd-domain-of s (var-name wy))))
            (cond
              ((= y-dom nil) s)
              (:else
                (fd-set-domain s (var-name wy) (fd-dom-without wx y-dom))))))
        ((and (number? wy) (is-var? wx))
          (let
            ((x-dom (fd-domain-of s (var-name wx))))
            (cond
              ((= x-dom nil) s)
              (:else
                (fd-set-domain s (var-name wx) (fd-dom-without wy x-dom))))))
        (:else s)))))

(define
  fd-neq
  (fn
    (x y)
    (fn
      (s)
      (let
        ((c (fn (s-prime) (fd-neq-prop x y s-prime))))
        (let
          ((s2 (fd-add-constraint s c)))
          (let
            ((s3 (c s2)))
            (cond ((= s3 nil) mzero) (:else (unit s3)))))))))

;; --- fd-lt ---

(define
  fd-lt-prop
  (fn
    (x y s)
    (let
      ((wx (mk-walk x s)) (wy (mk-walk y s)))
      (cond
        ((and (number? wx) (number? wy))
          (cond ((< wx wy) s) (:else nil)))
        ((and (number? wx) (is-var? wy))
          (let
            ((yd (fd-domain-of s (var-name wy))))
            (cond
              ((= yd nil) s)
              (:else
                (fd-set-domain
                  s
                  (var-name wy)
                  (filter (fn (v) (> v wx)) yd))))))
        ((and (is-var? wx) (number? wy))
          (let
            ((xd (fd-domain-of s (var-name wx))))
            (cond
              ((= xd nil) s)
              (:else
                (fd-set-domain
                  s
                  (var-name wx)
                  (filter (fn (v) (< v wy)) xd))))))
        ((and (is-var? wx) (is-var? wy))
          (let
            ((xd (fd-domain-of s (var-name wx)))
              (yd (fd-domain-of s (var-name wy))))
            (cond
              ((or (= xd nil) (= yd nil)) s)
              (:else
                (let
                  ((xd-prime (filter (fn (v) (< v (fd-dom-max yd))) xd)))
                  (let
                    ((s2 (fd-set-domain s (var-name wx) xd-prime)))
                    (cond
                      ((= s2 nil) nil)
                      (:else
                        (let
                          ((yd-prime (filter (fn (v) (> v (fd-dom-min xd-prime))) yd)))
                          (fd-set-domain s2 (var-name wy) yd-prime))))))))))
        (:else s)))))

(define
  fd-lt
  (fn
    (x y)
    (fn
      (s)
      (let
        ((c (fn (sp) (fd-lt-prop x y sp))))
        (let
          ((s2 (fd-add-constraint s c)))
          (let
            ((s3 (c s2)))
            (cond ((= s3 nil) mzero) (:else (unit s3)))))))))

;; --- fd-lte ---

(define
  fd-lte-prop
  (fn
    (x y s)
    (let
      ((wx (mk-walk x s)) (wy (mk-walk y s)))
      (cond
        ((and (number? wx) (number? wy))
          (cond ((<= wx wy) s) (:else nil)))
        ((and (number? wx) (is-var? wy))
          (let
            ((yd (fd-domain-of s (var-name wy))))
            (cond
              ((= yd nil) s)
              (:else
                (fd-set-domain
                  s
                  (var-name wy)
                  (filter (fn (v) (>= v wx)) yd))))))
        ((and (is-var? wx) (number? wy))
          (let
            ((xd (fd-domain-of s (var-name wx))))
            (cond
              ((= xd nil) s)
              (:else
                (fd-set-domain
                  s
                  (var-name wx)
                  (filter (fn (v) (<= v wy)) xd))))))
        ((and (is-var? wx) (is-var? wy))
          (let
            ((xd (fd-domain-of s (var-name wx)))
              (yd (fd-domain-of s (var-name wy))))
            (cond
              ((or (= xd nil) (= yd nil)) s)
              (:else
                (let
                  ((xd-prime (filter (fn (v) (<= v (fd-dom-max yd))) xd)))
                  (let
                    ((s2 (fd-set-domain s (var-name wx) xd-prime)))
                    (cond
                      ((= s2 nil) nil)
                      (:else
                        (let
                          ((yd-prime (filter (fn (v) (>= v (fd-dom-min xd-prime))) yd)))
                          (fd-set-domain s2 (var-name wy) yd-prime))))))))))
        (:else s)))))

(define
  fd-lte
  (fn
    (x y)
    (fn
      (s)
      (let
        ((c (fn (sp) (fd-lte-prop x y sp))))
        (let
          ((s2 (fd-add-constraint s c)))
          (let
            ((s3 (c s2)))
            (cond ((= s3 nil) mzero) (:else (unit s3)))))))))

;; --- fd-eq ---

(define
  fd-eq-prop
  (fn
    (x y s)
    (let
      ((wx (mk-walk x s)) (wy (mk-walk y s)))
      (cond
        ((and (number? wx) (number? wy))
          (cond ((= wx wy) s) (:else nil)))
        ((and (number? wx) (is-var? wy))
          (let
            ((yd (fd-domain-of s (var-name wy))))
            (cond
              ((and (not (= yd nil)) (not (fd-dom-member? wx yd))) nil)
              (:else
                (let
                  ((s2 (mk-unify wy wx s)))
                  (cond ((= s2 nil) nil) (:else s2)))))))
        ((and (is-var? wx) (number? wy))
          (let
            ((xd (fd-domain-of s (var-name wx))))
            (cond
              ((and (not (= xd nil)) (not (fd-dom-member? wy xd))) nil)
              (:else
                (let
                  ((s2 (mk-unify wx wy s)))
                  (cond ((= s2 nil) nil) (:else s2)))))))
        ((and (is-var? wx) (is-var? wy))
          (let
            ((xd (fd-domain-of s (var-name wx)))
              (yd (fd-domain-of s (var-name wy))))
            (cond
              ((and (= xd nil) (= yd nil))
                (let
                  ((s2 (mk-unify wx wy s)))
                  (cond ((= s2 nil) nil) (:else s2))))
              (:else
                (let
                  ((shared (cond ((= xd nil) yd) ((= yd nil) xd) (:else (fd-dom-intersect xd yd)))))
                  (cond
                    ((fd-dom-empty? shared) nil)
                    (:else
                      (let
                        ((s2 (fd-set-domain s (var-name wx) shared)))
                        (cond
                          ((= s2 nil) nil)
                          (:else
                            (let
                              ((s3 (fd-set-domain s2 (var-name wy) shared)))
                              (cond
                                ((= s3 nil) nil)
                                (:else (mk-unify wx wy s3))))))))))))))
        (:else s)))))

(define
  fd-eq
  (fn
    (x y)
    (fn
      (s)
      (let
        ((c (fn (sp) (fd-eq-prop x y sp))))
        (let
          ((s2 (fd-add-constraint s c)))
          (let
            ((s3 (c s2)))
            (cond ((= s3 nil) mzero) (:else (unit s3)))))))))

;; --- labelling ---

(define
  fd-try-each-value
  (fn
    (x dom s)
    (cond
      ((empty? dom) mzero)
      (:else
        (let
          ((s2 (mk-unify x (first dom) s)))
          (let
            ((s3 (cond ((= s2 nil) nil) (:else (fd-fire-store s2)))))
            (let
              ((this-stream (cond ((= s3 nil) mzero) (:else (unit s3))))
                (rest-stream (fd-try-each-value x (rest dom) s)))
              (mk-mplus this-stream rest-stream))))))))

(define
  fd-label-one
  (fn
    (x)
    (fn
      (s)
      (let
        ((wx (mk-walk x s)))
        (cond
          ((number? wx) (unit s))
          ((is-var? wx)
            (let
              ((dom (fd-domain-of s (var-name wx))))
              (cond
                ((= dom nil) mzero)
                (:else (fd-try-each-value wx dom s)))))
          (:else mzero))))))

(define
  fd-label
  (fn
    (vars)
    (cond
      ((empty? vars) succeed)
      (:else (mk-conj (fd-label-one (first vars)) (fd-label (rest vars)))))))