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

;;;; 4 CHICKEN
;;;; Kon Lovett, Jul '18
;;;; Kon Lovett, Aug '17
;;;; Kon Lovett, May '17
;;;; Kon Lovett, Mar '09

(module amb

(;export
  ;
  shuffle
  ;
  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
  (chicken base)
  (chicken fixnum)
  (chicken syntax)
  (chicken type)
  (only (chicken sort) sort!)
  (only (chicken random) pseudo-random-integer)
  (only (srfi 1) map!)
  (srfi 12)
  (only type-errors warning-argument-type)
  (only exn-condition make-exn-condition+))

;;; data-structures

(: shuffle ((list-of *) procedure -> (list-of *)))
;
(define (shuffle ls rand)
  (let* (
    (len (length ls))
    (tagged-ls (map (lambda (x) (cons (rand len) x)) ls)) )
    (map! cdr (sort! tagged-ls (lambda (a b) (fx< (car a) (car b)))) ) ) )

;;; miscmacros

(define-syntax let/cc
  (syntax-rules ()
    ((let/cc k e0 e1 ...)
     (call-with-current-continuation
      (lambda (k) e0 e1 ...)))))

(define-syntax define-parameter
  (syntax-rules ()
    ((define-parameter name value guard)
     (define name (make-parameter value guard)))
    ((define-parameter name value)
     (define name (make-parameter value)))
    ((define-parameter name)
     (define name (make-parameter (void))))))

;;;

;;

(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 pseudo-random-integer
  (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))) ) ) )

;;

(: amb-thunks ((list-of procedure) -> *))
;
(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))) ) ) ) ) ) )

(: amb-thunks-shuffled ((list-of procedure) #!rest (list procedure) -> *))
;
(define (amb-thunks-shuffled thunks . opts)
  (let ((rand (optional opts (amb-random-function))))
    (amb-thunks (shuffle thunks rand)) ) )

(: amb-find-thunk (procedure #!rest (list procedure) -> *))
;
(define (amb-find-thunk thunk . opts)
  (let ((failure (optional opts amb-exhausted)))
    (let/cc return
      (let ((fail-k (lambda () (return (failure)))))
        (parameterize ((amb-failure-continuation fail-k))
          (thunk) ) ) ) ) )

(: amb-collect-thunk (procedure -> *))
;
(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