#lang plait

(define-syntax-rule (generator yield-id body)
  (make-generator (lambda (yield-id) ; binds in `body`
                    (lambda (vd) ; dummy void argument
                      body))))

(define (make-generator proc)
  (local [(define yield ; <- set! on each demand
            identity) ; dummy initial value
          (define go ; <- set! to generator's state
            ;; initial state is the body function,
            ;; which needs a way to call yield:
            (proc (lambda (v) (yield v))))]
    ;; represent a generator as this function:
    (lambda ()
      ;; generator's escape to yield jumps here:
      (let/cc escape
        (begin
          ;; install `yield` for this time:
          (set! yield
                (lambda (v)
                  ;; save generator's continuation:
                  (let/cc k
                    (begin
                      (set! go k)
                      ;; escape with yielded value:
                      (escape v)))))
          ;; run generator until yield:
          (go (void)))))))

;; ----------------------------------------

(define (make-numbers start-n)
  (generator yield ; <- binds for use below
             (local [(define (numbers n)
                       (begin
                         (yield n) ; <- yield a value
                         (numbers (+ n 1))))]
               (numbers start-n))))

(define g (make-numbers 0))
(g) ; => 0
(g) ; => 1
(g) ; => 2
; ...