#lang plait
(define yield-param (make-parameter
                     (lambda (v) (error 'yield! "outside generator"))))

(define (yield v) ((parameter-ref yield-param) v))
(define-syntax-rule (generator (arg) exprs ...)
  (let ([last-checkpoint (none)])
    (lambda (v)
      (let/cc dyn-k ;; generator call site
        (parameterize ([yield-param
                        (lambda (v)
                          (let/cc gen-k ;; yield call site
                            (begin
                              (set! last-checkpoint
                                    (some gen-k))
                              (dyn-k v))))])
          (type-case (Optionof ('a -> 'b)) last-checkpoint
              [(none) (let ([arg v]) (begin exprs ...))]
              [(some k) (k v)])))
    )))

(define g1
  (generator (v)
      (letrec ([loop (lambda (n)
                       (begin
                         (yield n)
                         (loop (+ n 1))))])
        (loop v))))

(g1 10)  (g1 10)  (g1 10)

(define g2
  (generator (v)
      (letrec ([loop (lambda (n)
                       (loop (+ (yield n) n)))])
        (loop v))))

(g2 10)  (g2 10)  (g2 10)