;;;; amb.scm
;;;; The fundamental non-deterministic backtracking operator

;;;; Chicken 4 Port: Kon Lovett, Mar '09
;;;; Kon Lovett, May '17

(module amb

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

(import scheme)

(import chicken)
(import
  (only data-structures sort!)
  (only extras random) )

(import (only (srfi 1) map!))
(require-library (srfi 1))

(import
  (only miscmacros let/cc define-parameter)
  (only type-errors warning-argument-type)
  (only condition-utils make-exn-condition+))
(require-library
  miscmacros type-errors condition-utils)

;;;

;;

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

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

;;

(define-parameter amb-random-function random
  (lambda (x)
    (if (procedure? x)
      x
      (begin
	(warning 'amb-random-function "not a procedure" x)
        (amb-random-function) ) ) ) )

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

;;

(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-shuffled (list (lambda () ?expr0) ...)) ) ) )

(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 (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-thunks-shuffled thunks #!optional (rand (amb-random-function)))
  (amb-thunks (shuffle thunks rand)) )

(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) ) ) ) )

;;;

(define (shuffle ls random)
  (let ((len (length ls)))
    (map!
      cdr
      (sort!
        (map (lambda (x) (cons (random len) x)) ls)
        (lambda (x y) (< (car x) (car y)))) ) ) )

) ;module amb