amb.ss
#lang scheme

;;; Exceptions signalling a lack of alternatives for backtracking

(define-struct (exn:fail:amb exn:fail)
  ()
  #:transparent)

(define (raise-amb-error name)
  (raise
   (make-exn:fail:amb (format "~a: out of alternatives" name)
                      (current-continuation-marks))))

(provide/contract
 (struct (exn:fail:amb exn:fail)
   ([message string?]
    [continuation-marks continuation-mark-set?])))

;;; Primitive ambivalence control mechanisms

(define amb-restarts
  (make-parameter '()))

(define amb-prompt-tag
  (make-continuation-prompt-tag 'amb))

(define (call-with-amb-prompt
         thunk
         [failure-result (λ ()
                           (raise-amb-error 'call-with-amb-prompt))])
  (parameterize ([amb-restarts '()])
    (call-with-continuation-prompt
     thunk amb-prompt-tag failure-result)))

(define (amb-fail)
  (let ([restarts (amb-restarts)])
    (when (null? restarts)
      (abort-current-continuation amb-prompt-tag))
    (amb-restarts (cdr restarts))
    ((car restarts))))

(define-for-syntax ((make-amb-for-syntax for) stx)
  (syntax-case stx ()
    [(_ (for-clauses ...) body ...)
     (quasisyntax/loc stx
       (nest [(let ([restarts (amb-restarts)]))
              (let/ec return)]
         (#,for #,stx () (for-clauses ...)
                (call/cc
                 (λ (continue)
                   (amb-restarts (cons continue restarts))
                   (call-with-values
                    (λ ()
                      body ...)
                    return))
                 amb-prompt-tag))
         (amb-fail)))]))

(provide/contract
 [call-with-amb-prompt (->* ((-> any)) ((-> any)) any)]
 [amb-fail (-> any)])

;;; Syntax to provide alternatives

(define-syntax for/amb
  (make-amb-for-syntax #'for/fold/derived))

(define-syntax for*/amb
  (make-amb-for-syntax #'for*/fold/derived))

(define (amb-call . thunks)
  (for/amb ([thunk (in-list thunks)])
    (thunk)))

(define-syntax amb
  (syntax-rules ()
    [(amb)
     (amb-fail)]
    [(amb expr)
     expr]
    [(amb expr ...)
     (amb-call (λ () expr) ...)]))

(define-syntax amb-assert
  (syntax-rules ()
    [(amb-assert condition)
     (unless condition
       (amb-fail))]))

(provide/contract
 [amb-call (->* () () #:rest (listof (-> any)) any)])

(provide
 for/amb for*/amb amb amb-assert)

;;; Syntax to consume alternatives

(define-syntax amb-find
  (syntax-rules ()
    [(amb-find expr ...)
     (call-with-amb-prompt
      (λ ()
        expr ...)
      (λ ()
        (raise-amb-error 'amb-find)))]))

(define-syntax amb-collect
  (syntax-rules ()
    [(amb-collect expr ...)
     (let ([results '()])
       (call-with-amb-prompt
        (λ ()
          (set! results
            (cons (begin
                    expr ...)
                  results))
          (amb-fail))
        (λ ()
          (reverse results))))]))

(provide
 amb-find amb-collect)

;;; Utilities to use ambivalent expressions as sequence generators

(define make-amb-sequence
  (local [(define-struct pos
            (results restarts))
          (define (delay-pos generator+restarts)
            (delay
              (parameterize ([amb-restarts (cdr generator+restarts)])
                (call-with-continuation-prompt
                 (λ ()
                   (call-with-values
                    (car generator+restarts)
                    (λ results
                      (make-pos results (amb-restarts)))))
                 amb-prompt-tag
                 (λ ()
                   #f)))))
          (define (pos->results pos)
            (apply values (pos-results (force pos))))
          (define (pos-next pos)
            (delay-pos (pos-restarts (force pos))))
          (define (constantly-true . _)
            #t)]
    (λ (thunk)
      (make-do-sequence
       (λ ()
         (values
          pos->results
          pos-next
          (delay-pos (list thunk))
          force
          constantly-true
          constantly-true))))))

(define-syntax in-amb
  (syntax-rules ()
    [(in-amb expr ...)
     (make-amb-sequence (λ () (amb expr ...)))]))

(provide/contract
 [make-amb-sequence (-> (-> any) sequence?)])

(provide
 in-amb)