;;;; amb.scm
;;;; The fundamental non-deterministic backtracking operator
;;;; Chicken 4 Port: Kon Lovett, Mar '09

;;; Module `amb'

(module amb

  (;export
    amb
    amb/random
    amb-find
    amb-collect
    amb-assert
    amb-failure-continuation
    amb-thunks
    amb-find-thunk
    amb-collect-thunk)

  (import
    scheme
    chicken
    (only data-structures shuffle)
    (only extras random)
    (only miscmacros let/cc define-parameter)
    (only type-errors warning-argument-type)
    (only condition-utils make-exn-condition+))

  (require-library data-structures extras miscmacros type-errors condition-utils)

;;

(define-syntax amb
  (syntax-rules ()
    ((_)            ((amb-failure-continuation)) )
    ((_ ?expr0 ...) (amb-thunks (list (lambda () ?expr0) ...)) ) ) )

(define-syntax amb/random
  (syntax-rules ()
    ((_)            ((amb-failure-continuation)) )
    ((_ ?expr0 ...) (amb-thunks (shuffle (list (lambda () ?expr0) ...) random)) ) ) )

(define-syntax amb-find
  (syntax-rules ()
    ((_ ?expr)        (amb-find-thunk (lambda () ?expr)) )
    ((_ ?expr ?fail)  (amb-find-thunk (lambda () ?expr) (lambda () ?fail)) ) ) )

(define-syntax amb-collect
  (syntax-rules ()
    ((_ ?expr) (amb-collect-thunk (lambda () ?expr)) ) ) )

(define-syntax amb-assert
  (syntax-rules ()
    ((_ ?expr) (unless ?expr ((amb-failure-continuation))) ) ) )

;;

(define (make-amb-exhausted-condition)
  (make-exn-condition+ 'amb "expression tree exhausted" '() 'amb) )

(define (amb-exhausted) (signal (make-amb-exhausted-condition)))

;;

(define-parameter amb-failure-continuation amb-exhausted
  (lambda (x)
    (cond ((procedure? x) x)
          (else
           (warning-argument-type 'amb-failure-continuation x 'procedure)
           (amb-failure-continuation) ) ) ) )

;;

(define (amb-thunks thunks)
  (let ((afc (amb-failure-continuation)))
    (let/cc return
      (let loop ((tt thunks))
        (cond ((null? tt)
               (amb-failure-continuation afc)
               (afc) )
              (else
               (amb-failure-continuation (lambda () (loop (cdr tt))))
               (return ((car tt))) ) ) ) ) ) )

(define (amb-find-thunk thunk #!optional (failure amb-exhausted))
  (let/cc return
    (parameterize ((amb-failure-continuation (lambda () (return (failure)))))
      (thunk) ) ) )

(define (amb-collect-thunk thunk)
  (let ((afc #f))
    (dynamic-wind
      (lambda () (set! afc (amb-failure-continuation)) )
      (lambda ()
        (let/cc return
          (let* ((root (list #f))
                 (head root))
            (amb-failure-continuation (lambda () (return (cdr root))))
            (set-cdr! head (list (thunk)))
            (set! head (cdr head))
            ((amb-failure-continuation))) ) )
      (lambda () (amb-failure-continuation afc) ) ) ) )

) ;module amb