;;;; 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)
(import (chicken base))
(import (chicken syntax))
(import (chicken type))
(import (only (chicken condition) signal make-property-condition make-composite-condition))
(import (only (chicken string) conc))
(import (only (chicken bitwise) integer-length))
(import (only (chicken sort) sort!))
(import (only (chicken random) pseudo-random-integer))

;;(from miscmacros.scm)

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

;;

(: vector-shuffle! ((vector-of *) #!optional (procedure (fixnum) fixnum) -> void))
(: shuffle ((list-of *) #!optional (procedure (fixnum) fixnum) -> (list-of *)))
(: amb-thunks ((list-of procedure) -> *))
(: amb-thunks-shuffled ((list-of procedure) #!rest (list procedure) -> *))
(: amb-find-thunk (procedure #!rest (list procedure) -> *))
(: amb-collect-thunk (procedure -> *))

;;

(define platform-random pseudo-random-integer)

(define (vector-shuffle! vec #!optional (rnd platform-random))
  (let ((len (vector-length vec)))
    (define (swap-adj! i)
      (let ((i+1 (modulo (add1 i) len))
            (tmp (vector-ref vec i)) )
            (vector-set! vec i (vector-ref vec i+1))
            (vector-set! vec i+1 tmp) ) )
    (do ((n (integer-length len) (sub1 n)))
        ((= n 0))
      (swap-adj! (rnd len)) ) ) )

(define (shuffle ls #!optional (rnd platform-random))
  (let ((vec (list->vector ls)))
    (vector-shuffle! vec rnd)
    (vector->list vec) ) )

;;

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

(define make-amb-exhausted-condition
  (let (
    (cached-amb-exhausted-condition
      (make-composite-condition
        (make-property-condition 'exn
          'location 'amb 'message "expression tree exhausted" 'arguments '())
        (make-property-condition '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 (conc "(amb-failure-continuation) bad argument - not a procedure: " x))
        (amb-failure-continuation) ) ) ) )

;;

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

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

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

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

(define-syntax amb-assert
  (syntax-rules ()
    ((amb-assert ?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 . opts)
  (let ((rand (optional opts (amb-random-function))))
    (amb-thunks (shuffle thunks rand)) ) )

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

(define (amb-collect-thunk thunk)
  (let ((afc (amb-failure-continuation)))
    (dynamic-wind
      ;
      void
      ;
      (lambda ()
        (let/cc return
          (let* ((root (the list (list #f))) ;ovverride strict-types assumption
                 (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