amb in scheme (call-with-current-continuation)

amb.ss

#lang r5rs
(define amb-fail '*)

(define initialize-amb-fail
  (lambda ()
    (set! amb-fail
      (lambda ()
        (display "amb tree exhausted")))))

(initialize-amb-fail)

(define call/cc call-with-current-continuation)

(define displayln
  (lambda (x)
    (display x) (newline)
))

(define-syntax amb
  (syntax-rules ()
    ((amb alt ...)
     (let ((prev-amb-fail amb-fail))
       (call/cc
        (lambda (sk)

          (call/cc
           (lambda (fk)
             (set! amb-fail
                   (lambda ()
                     (set! amb-fail prev-amb-fail)
                     (fk 'fail)))
             (sk alt))) ...
             
             (prev-amb-fail)))))))

; util fuctions
(define number-between
  (lambda (lo hi)
    (let loop ((i lo))
      (if (> i hi) (amb)
          (amb i (loop (+ i 1)))))))
; usage : (number-between 0 10) (amb) (amb)

(define assert
  (lambda (pred)
    (if (not pred) (amb))))


(define-syntax apply-amb
  (syntax-rules ()
    ((apply-amb ls)
     (eval `(amb ,@ls) (interaction-environment)))))

(define-syntax bag-of
  (syntax-rules ()
    ((bag-of e)
     (let ((prev-amb-fail amb-fail)
           (results '()))
       (if (call/cc
            (lambda (k)                                                
              (set! amb-fail (lambda () (k #f)))                ;<-----+
              (let ((v e))             ;amb-fail will be modified by e |
                (set! results (cons v results))                       ;|
                (k #t))))                                             ;|
           (amb-fail))                 ;so this amb-fail may not be ---+
       (set! amb-fail prev-amb-fail)
       (reverse results)))))

(define (distinct? . ls)
  (let loop ((lst (car ls)))
    (let ((first (car lst)) (rest (cdr lst)))
      (cond 
       ((null? rest) #t)
       ((member first rest) #f)
       (else (loop rest))))))



(define (del n ls)
  (let ((ls (reverse (reverse ls))))
    (cond ((null? ls) ls)
          ((eqv? n (car ls)) (cdr ls))
          (else 
           (let loop ((l (cdr ls)) (last ls))
             (cond ((null? l) ls)
                   ((equal? n (car l))
                    (set-cdr! last (cdr l))
                    ls)
                   (else (loop (cdr l) l))))))))


(define (prime? n)
  (call/cc
   (lambda (return)
     (do ((i 2 (+ i 1)))
         ((> i (sqrt n)) #t)
       (if (= (modulo n i) 0) 
           (return #f))))))

(define gen-prime
  (lambda (hi)
    (let ((i (number-between 2 hi)))
      (assert (prime? i))
      i)))

(displayln (bag-of (gen-prime 20)))

(bag-of
(let ((x (list (amb 2 1 -2 5 8 18) (amb 9 8 2 4 14 20))))
  (assert (= (+ (car x) (cadr x)) 10))
  (displayln x)))