typesetting/csp/aima.rkt

#lang debug racket
(require racket/generator sugar/debug)

(struct $csp (variables domains neighbors constraints initial curr_domains nassigns) #:transparent #:mutable)
(define assignment? hash?)
(define variable? symbol?)
(define removal? (cons/c variable? any/c))

(define/contract (make-csp variables domains neighbors constraints)
  ((listof variable?) hash? hash? procedure? . -> . $csp?)
  ($csp variables domains neighbors constraints null #f 0))

(define/contract (assign csp var val assignment)
  ($csp? variable? any/c assignment? . -> . void?)
  ;; Add {var: val} to assignment; Discard the old value if any.
  (hash-set! assignment var val)
  (set-$csp-nassigns! csp (add1 ($csp-nassigns csp))))

(define/contract (unassign csp var assignment)
  ($csp? variable? assignment? . -> . void?)
  ;; Remove {var: val} from assignment.
  ;; DO NOT call this if you are changing a variable to a new value;
  ;; just call assign for that.
  (hash-remove! assignment var))

(define/contract (nconflicts csp var val assignment)
  ($csp? variable? any/c assignment? . -> . number?)
  ;; Return the number of conflicts var=val has with other variables."""
  ;; Subclasses may implement this more efficiently
  (for/sum ([v (in-list (hash-ref ($csp-neighbors csp) var))]
            #:when (hash-has-key? assignment v))
    (if (($csp-constraints csp) var val v (hash-ref assignment v)) 0 1)))

(define (display csp assignment)
  (displayln csp))

(define/contract (all-variables-assigned? csp assignment)
  ($csp? assignment? . -> . boolean?)
  (= (length (hash-keys assignment)) (length ($csp-variables csp))))

(define/contract (goal_test csp state)
  ($csp? assignment? . -> . boolean?)
  ;; The goal is to assign all variables, with all constraints satisfied.
  (define assignment state)
  (and (all-variables-assigned? csp assignment)
       (for/and ([variable ($csp-variables csp)])
         (zero? (nconflicts csp variable (hash-ref assignment variable) assignment)))))

;; These are for constraint propagation

(define/contract (support_pruning csp)
  ($csp? . -> . void?)
  ;; Make sure we can prune values from domains. (We want to pay
  ;; for this only if we use it.)
  (unless ($csp-curr_domains csp)
    (define h (make-hasheq))
    (for ([v ($csp-variables csp)])
      (hash-set! h v (hash-ref ($csp-domains csp) v)))
    (set-$csp-curr_domains! csp h)))
      
(define/contract (suppose csp var value)
  ($csp? variable? any/c . -> . (listof removal?))
  ;; Start accumulating inferences from assuming var=value
  (support_pruning csp)
  (define removals
    (for/list ([a (hash-ref ($csp-curr_domains csp) var)]
               #:when (not (equal? a value)))
      (cons var a)))
  (hash-set! ($csp-curr_domains csp) var (list value))
  removals)

;; todo: update uses of `prune` to be functional on removals
(define/contract (prune csp var value removals)
  ($csp? variable? any/c (or/c #f (listof removal?)) . -> . (listof removal?))
  ;; Rule out var=value
  (hash-update! ($csp-curr_domains csp) var (λ (vals) (remove value vals)))
  (and removals (append removals (list (cons var value)))))

(define/contract (choices csp var)
  ($csp? variable? . -> . (listof any/c))
  ;; Return all values for var that aren't currently ruled out.
  (hash-ref (or ($csp-curr_domains csp) ($csp-domains csp)) var))

(define/contract (infer_assignment csp)
  ($csp? . -> . assignment?)
  ;; Return the partial assignment implied by the current inferences.
  (support_pruning csp)
  (define assignment (make-hasheq))
  (for ([v (in-list ($csp-variables csp))])
    (match (hash-ref ($csp-curr_domains csp) v)
      [(list one-value) (hash-set! assignment v one-value)]
      [else #f]))
  assignment)

(define/contract (restore csp removals)
  ($csp? (listof removal?) . -> . void?)
  ;; Undo a supposition and all inferences from it.
  (for ([removal (in-list removals)])
    (match removal
      [(cons B b) (hash-update! ($csp-curr_domains csp) B
                                (λ (vals) (append vals (list b))))])))


;; ______________________________________________________________________________
;; CSP Backtracking Search

;; Variable ordering

(define/contract (first_unassigned_variable assignment csp)
  (assignment? $csp? . -> . (or/c #false variable?))
  ;; The default variable order.
  (for/first ([var (in-list ($csp-variables csp))]
              #:unless (hash-has-key? assignment var))
    var))

;; Value ordering

(define/contract (unordered_domain_values var assignment csp)
  (variable? assignment? $csp? . -> . (listof any/c))
  ;; The default value order.
  (choices csp var))

;; Inference

(define/contract (no_inference csp var value assignment removals)
  ($csp? variable? any/c assignment? (listof removal?) . -> . boolean?)
  #true)

(define/contract (backtracking_search csp
                                      [select_unassigned_variable first_unassigned_variable]
                                      [order_domain_values unordered_domain_values]
                                      [inference no_inference])
  (($csp?) (procedure? procedure? procedure?) . ->* . generator?)
  (generator ()
             ;; todo: incorporate `yield`
             (let backtrack ([assignment (make-hasheq)])
               (match (select_unassigned_variable assignment csp)
                 [#false (and (goal_test csp assignment) assignment)]
                 [var
                  (cond
                    [(for/or ([val (in-list (order_domain_values var assignment csp))]
                              #:when (zero? (nconflicts csp var val assignment)))
                       (assign csp var val assignment)
                       (define removals (suppose csp var val))
                       (cond
                         [(and (inference csp var val assignment removals) (backtrack assignment))]
                         [else (restore csp removals) #false]))]
                    [else (unassign csp var assignment) #false])]))))

(define/contract (solve* csp [solver backtracking_search] [finish-proc values]
                         #:count [solution-limit +inf.0])
  (($csp?) (procedure? procedure? #:count integer?) . ->* . (or/c #f (non-empty-listof any/c)))
  (match (for/list ([solution (in-producer (solver csp) (void))]
                    [idx (in-range solution-limit)])
           (finish-proc solution))
    [(? pair? solutions) solutions]
    [else #f]))

(define/contract (solve csp [solver backtracking_search] [finish-proc values])
  (($csp?) (procedure? procedure?) . ->* . any/c)
  (match (solve* csp solver finish-proc #:count 1)
    [(list solution) solution]
    [else #f]))

(require rackunit)
(define vs '(wa nsw t q nt v sa))
(define ds (for/hash ([k vs])
             (values k '(red green blue))))
(define ns (for*/hash ([(i ns) (in-dict
                                '((wa nt sa)
                                  (nt wa sa q)
                                  (q nt sa nsw)
                                  (nsw q sa v)
                                  (v sa nsw)
                                  (sa wa nt q nsw v)
                                  (t)))])
             (values i ns)))
(define csp (make-csp vs ds ns (λ (A a B b) (not (equal? a b)))))
(check-true ($csp? csp))
(define a (make-hasheq))
(assign csp 'key 42 a)
(check-equal? (hash-ref a 'key) 42)
(unassign csp 'key a)
(check-exn exn:fail? (λ () (hash-ref a 'key)))
(check-equal? 0 (nconflicts csp 'wa 'red (hasheq 'wa 42)))
(support_pruning csp)
(check-true (hash? ($csp-curr_domains csp)))

(check-equal? (suppose csp 'wa 'red) '((wa . green) (wa . blue)))
(check-equal?
 (hash-ref ($csp-curr_domains csp) 'wa) '(red))

(check-equal? (prune csp 'v 'red empty) '((v . red)))

(check-equal? (choices csp 'v) '(green blue))
(check-equal? (choices csp 'wa) '(red))
(check-equal? (infer_assignment csp)
              (make-hasheq '((wa . red))))
(check-equal? (suppose csp 'v 'blue) '((v . green)))
(check-equal? (infer_assignment csp)
              (make-hasheq '((v . blue) (wa . red))))
(restore csp '((wa . green)))
(check-equal? (infer_assignment csp)
              (make-hasheq '((v . blue))))
(restore csp '((v . blue)))
(check-equal? (infer_assignment csp) (make-hasheq))

(check-equal? (first_unassigned_variable (hash) csp) 'wa)
(check-equal? (unordered_domain_values 'wa (hash) csp) '(red green))

(set-$csp-curr_domains! csp #f) ; reset current domains

(check-equal? (solve csp)
              (make-hasheq '((nsw . green) (nt . green) (q . red) (sa . blue) (t . blue) (v . red) (wa . red))))

(set-$csp-curr_domains! csp #f)
(check-equal? (suppose csp 'nsw 'red) '((nsw . green) (nsw . blue)))
(check-equal? (solve csp)
              (make-hasheq '((nsw . red) (nt . red) (q . green) (sa . blue) (t . blue) (v . green) (wa . green))))