amb with ivars

amb.scm

;; Adam Foltzer
;; amb with ivars

(load "pmatch.scm")

;;;; Global scheduler queue and helpers; ugly!
(define *q* '())

(define push-right!
  (lambda (x)
    (set! *q* (append *q* `(,x)))))

(define peek-left
  (lambda ()
    (cond
      [(null? *q*) #f]
      [else (car *q*)])))

(define pop-left!
  (lambda ()
    (cond
      [(null? *q*) #f]
      [else (let ([a (car *q*)])
              (set! *q* (cdr *q*))
              a)])))

(define filter!
  (lambda (pred)
    (set! *q* (filter pred *q*))))

;; Representation of work units to be put in the queue

(define-record-type work (fields eng sk id))

;;;; ids

;; ids are lists of gensyms where parent/child relationships are
;; represented by shared prefixes

(define prefix?
  (lambda (p ls)
    (cond
      [(null? p) #t]
      [(null? ls) #f]
      [else (and (eq? (car p) (car ls))
                 (prefix? (cdr p) (cdr ls)))])))

;; Create a new id as a child of the currently active work unit, or
;; just a gensym if there's no current work

(define new-id
  (lambda ()
    (cond
      [(null? *q*) `(,(gensym))]
      [else (append (work-id (peek-left)) `(,(gensym)))])))

;;;; ivars

;; A variant of ivars where rather than just empty and full, we also
;; have a case where the ivar represents the top level of the
;; computation. This variant contains a continuation that breaks out
;; of the whole nested amb.

(define new-ivar
  (lambda ()
    (box `(empty))))

(define new-top-level-ivar
  (lambda (k)
    (box `(top-level ,k))))

;; get-ivar blocks until the ivar is filled in
(define get-ivar
  (lambda (iv)
    (let ([contents (unbox iv)])
      (pmatch contents
        [(empty) (begin (engine-block) (get-ivar iv))]
        [(full ,x) x]))))

;; since we're only calling set-ivar! during the sk of an engine
;; (rather than in the engine body), we don't have to worry about
;; interrupts as with most ivar implementations
(define set-ivar!
  (lambda (iv x)
    (let ([contents (unbox iv)])
      (pmatch contents
        [(empty) (set-box! iv `(full ,x))]
        [(full ,x) (warningf 'put-ivar "multiple put")]
        ;; when top-level, make sure we clear out the *q* before exiting
        [(top-level ,k) (begin (set! *q* '()) (k x))]))))

;;;; main implementation

(define-syntax amb
  (syntax-rules ()
    [(_ e1 e2) (amb-impl (lambda () e1) (lambda () e2))]))

(define amb-impl
  (lambda (th1 th2)
    (call/cc ;; grab a continuation in case we're calling non-nested
      (lambda (k)
        (let* ([id (new-id)]
               [top-level? (null? *q*)] ;; *q* is empty if we're top level
               [iv (if top-level? (new-top-level-ivar k) (new-ivar))]
               [sk (lambda (_ v)
                     (begin (set-ivar! iv v)
                            ;; if we make it here, the ivar was *not*
                            ;; top-level, so we have to shoot down any
                            ;; children of this amb, including the one
                            ;; that just succeeded
                            (shootdown! id)
                            ;; since we removed the current engine, we
                            ;; reenter the loop rather than
                            ;; rescheduling with the current engine
                            (enter-loop)))])
          ;; we're fiddling with the queue in code that might be run
          ;; in an engine, so we have to disable interrupts
          (with-interrupts-disabled
            (push-right! (make-work (make-engine th1) sk id))
            (push-right! (make-work (make-engine th2) sk id)))
          ;; if top-level, we just make a tail call to enter the loop;
          ;; the answer will be returned as the value of the whole
          ;; call/cc. Otherwise, we do a [blocking] get on the ivar we
          ;; just made, which will eventually return the right value
          (if top-level?
              (enter-loop)
              (get-ivar iv)))))))      

(define nticks (make-parameter 50))

;; reschedule! is the fk of engines, so we don't have to worry about interrupts
(define reschedule!
  (lambda (new-eng)
    (when (null? *q*) (errorf 'reschedule! "empty work queue!"))
    ;; push the current work unit (with new engine) on the right of the queue
    (let ([old (peek-left)])
      (push-right! (make-work new-eng (work-sk old) (work-id old)))
      ;; pop only after pushing, in case this is the only work unit
      ;; (although this shouldn't arise)
      (pop-left!)
      (let ([new (peek-left)])
        ((work-eng new) (nticks) (work-sk new) reschedule!)))))

(define enter-loop
  (lambda ()
    (let ([new (peek-left)])
      ((work-eng new) (nticks) (work-sk new) reschedule!))))

;; remove all work units which have the given id as a prefix
(define shootdown!
  (lambda (id)
    (filter! (lambda (w) (not (prefix? id (work-id w)))))))

;;;; Examples

(define f
  (lambda (x)
    (amb (f x) x)))

(define fact
  (lambda (n)
    (if (zero? n)
        1
        (* n (fact (sub1 n))))))

(define fact-1
  (lambda ()
    (amb (fact -1) (fact 5))))

(define fact-2
  (lambda ()
    (amb (fact -1)
         (amb ((lambda (x) (x x)) (lambda (x) (x x)))
              (fact 5)))))

(define fib
  (lambda (n)
    (if (< n 2)
        n
        (+ (fib (- n 1)) (fib (- n 2))))))

(define fib-1
  (lambda ()
    (amb (fib 100)
         (fib 25))))

(define fib/ticks
  (lambda (ticks)
    (parameterize ([nticks ticks])
      (printf "nticks=~a " ticks)
      (collect)
      (time (fib-1)))))

(define many-fibs
  (lambda ()
    (map fib/ticks '(1 2 4 8 16 32 64 128))))