#lang plai

;; Just use a box to represent a logic variable.  If the box contains |false|,
;; the variable is unbound; otherwise, it's bound to the content of the box.
;; This means that there's no way to represent a variable bound to |false|...
(define var? box?)

;; Binds a collection of logic variables within a sequence of body expressions.
(define-syntax exists
  (syntax-rules ()
    [(_ (v ...) e ...)
     (let ([v (box #f)] ...)
       e ...)]))

;; Recursively "resolves" a value by chasing down the contents of all variables.
(define (resolve thing)
  (cond
    [(var? thing)
     (cond [(unbox thing) => resolve]
           [else thing])]
    [(cons? thing)
     ;; We use car and cdr here because they're more liberal about their
     ;; arguments: first and rest complain if you don't give them a proper list.
     (cons (resolve (car thing)) (resolve (cdr thing)))]
    [else s]))

;; Unifies two things, either of which may be a variable or not.
(define (unify thing1 thing2)
  (let ([r1 (resolve thing1)]
        [r2 (resolve thing2)])
    (cond
      [(var? r1) (bind-var r1 r2)]
      [(var? r2) (bind-var r2 r1)]
      [else      (unify-values r1 r2)])))

;; Unifies two things which are not variables (though they may contain
;; variables).
(define (unify-values v1 v2)
  (or
   (and (symbol? v1) (symbol? v2) (symbol=? v1 v2))
   (and (number? v1) (number? v2) (= v1 v2))
   (and (string? v1) (string? v2) (string=? v1 v2))
   (and (empty? v1) (empty? v2))
   (and (cons? v1) (cons? v2)
        (unify (car v1) (car v2))
        (unify (cdr v1) (cdr v2)))
   (fail)))

;; Stack of mutations to undo bindings caused by unification when backtracking.
(define unbindings empty)

;; Binds var to val.
(define (bind-var var val)
  (set! unbindings (cons (cons (list var (unbox var)) (first unbindings)) (rest unbindings)))
  (set-box! var val))

;; Stack of continuations for backtracking to try different alternatives.
(define conts empty)

;; Pushes a continuation onto the stack.
(define (push! cont)
  (set! unbindings (cons empty unbindings))
  (set! conts (cons cont conts)))

;; Pops a continuation off the stack and undoes whatever bindings occurred since
;; the point when that continuation was pushed.
(define (pop!)
  (begin0
    (first conts)
    (let ([to-unbind (first unbindings)])
      (for-each (λ (p) (set-box! (first p) (second p))) to-unbind))
    (when (not (empty? (rest conts)))
      (set! unbindings (rest unbindings))
      (set! conts (rest conts)))))

;; Resets the continuation and unbinding stacks.
(define (reset!)
  (set! unbindings (list empty))
  (set! conts (last-pair conts)))

;; A macro for "nondeterministic" choice.
(define-syntax amb
  (syntax-rules ()
    [(_) (fail)]
    [(_ e es ...)
     (let/cc esc
       (let/cc k
         (push! k)
         (esc e))
       (amb es ...))]))

(define (fail)
  ((pop!)))

;; Set up top level.
(amb (void) (printf "no~n"))

(define more fail)

;; A macro to run a top-level query with a fresh set of logic variables.
;; Prints bindings for successful executions (or "yes") for queries with no
;; variables.
(define-syntax query
  (syntax-rules ()
    [(_ (v ...) e ...)
     (exists (v ...)
       (reset!)
       e ...
       (when (empty? (list 'v ...))
         (printf "yes~n"))
       (printf "~a -> ~a~n" 'v (resolve v)) ...)]))

;; Definition of an "app[e]nd" relation between three lists:
;; |result| is the result of appending |l1| and |l2|.
(define (appnd l1 l2 result)
  (amb (and (unify l1 empty)
            (unify l2 result))
       (exists (l1-fst l1-rst result-rst)
         (unify l1 (cons l1-fst l1-rst))
         (unify result (cons l1-fst result-rst))
         (appnd l1-rst l2 result-rst))))