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

(define-type random-fixnum (fixnum -> fixnum))
(define-type thunk (-> *))

(: amb-random-function (#!optional random-fixnum -> random-fixnum))
(: amb-failure-continuation (#!optional thunk -> thunk))

(: amb-thunks ((list-of thunk) -> *))
(: amb-thunks-shuffled ((list-of thunk) #!optional random-fixnum -> *))
(: amb-find-thunk (thunk #!optional thunk -> *))
(: amb-collect-thunk (thunk -> *))

(: vector-shuffle! ((vector-of *) random-fixnum -> void))
(: list-shuffle ((list-of *) random-fixnum -> (list-of *)))

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

;;

(define-inline (vector-swap! vec i j)
  (let ((tmp (vector-ref vec i)) )
    (vector-set! vec i (vector-ref vec j))
    (vector-set! vec j tmp) ) )

;#; ;fast
(define (vector-shuffle! vec rnd)
  (let ((len (vector-length vec)))
    (define (swap-adj! i) (vector-swap! vec i (fxmod (fx+ i 1) len)))
    (do ((n (integer-length len) (fx- n 1)))
        ((fx= n 0))
      (swap-adj! (rnd len)) ) ) )

#; ;Fisher Yates shuffle
(define (vector-shuffle! vec rnd)
  (let ((len (vector-length vec)))
    ;backward compatible w/ integer random (fp ~= but pseudo-random-real faster)
    (define rnd-real
      (let ((near-1 (- 1 (/ 1 len))))
        (lambda ()
          (let ((n (rnd len)))
            (cond ((fx= n 0) 0)
                  ((fx= n 1) near-1)
                  (else (/ 1 n))) ) ) ) )
    (do ((i (fx- len 1) (fx- i 1)))
        ((fx= i 0))
      (vector-swap! vec i (floor (* (rnd-real) (fx+ i 1)))) ) ) )

(define (list-shuffle ls rnd)
  (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 (list-shuffle thunks rand)) ) )

(define (amb-find-thunk thunk . opts)
  (let ((failure (optional opts amb-exhausted)))
    (let/cc return
      (parameterize ((amb-failure-continuation
                      (lambda () (return (failure)))))
        (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