generator.ss
(module generator mzscheme
(require (lib "plt-match.ss")
(lib "list.ss")
(lib "contract.ss")
(lib "stxparam.ss")
(lib "etc.ss")
(lib "control.ss"))
(require-for-syntax (lib "stx.ss" "syntax"))

;; We define a generator as
;;
;;    (make-generator thunk)
;;
;; where thunk is a zero-arity procedure.
;;
;; The structure (make-generator thunk) is callable
;; as a procedure.  (See Section 4.6 of MzScheme language
;; manual for details on this.)
(define-values
(struct:generator make-generator generator?
generator-ref generator-set)
(make-struct-type 'generator #f 1 0 #f null #f 0))
(provide/contract (generator? (-> any/c boolean?)))

;; We raise exn:fail:generator-exhausted
;; when we reach the end of a generator.
(define-struct (exn:fail:generator-exhausted exn:fail) ())

(provide (struct exn:fail:generator-exhausted ()))

(define-syntax (loop-forever stx)
(syntax-case stx ()
[(_ e1 e2 ...)
(syntax/loc stx
(let loop ()
e1 e2 ... (loop)))]))

(provide/contract
(rename ext:make-generator make-generator
(((-> any/c any) . -> . any) . -> . generator?)))
;; make-generator: ((X -> any) -> any) -> (generator-of X)
;; Creates a generator, given a seed function that accepts
;; a yield function.  Whenever the seed-function wants to
;; pass values back, it should call the yield function with
;; that value.
;;
;; For example: (make-generator
;;                (lambda (yield) (yield 1) (yield 2)))
;; produces a generator that yields 1 and 2.
(define (ext:make-generator seed-function)
(local (

;; A yielded-value is one of the following:
;;     1.  (make-yielded-datum datum)
;;     2.  (make-yielded-exn exn)
(define-struct yielded-datum (datum))
(define-struct yielded-exn (exn))

(define my-prompt (new-prompt))

(define toplevel
(case-lambda
[()
(toplevel (void))]
[(resumption-val)
;; We save the continuation marks before the prompt
;; to give a cleaner stack trace when the generator
;; exhausts.
(set! cont-marks (current-continuation-marks))
(call-with-continuation-barrier
(lambda ()
(match (prompt-at my-prompt (resume-here resumption-val))
[(struct yielded-datum (datum)) datum]
[(struct yielded-exn (exn)) (raise exn)])))]))

(define cont-marks (current-continuation-marks))

;; Initially, resume-here is bound to a function that does
;; the seed-function.  As we generate new values, resume-here
;; will be rebound to the remainder of the work.
(define (resume-here resume-value)
(with-handlers ((exn:fail?
(lambda (exn)
(loop-forever (yield-exn exn)))))
(seed-function yield-datum)
(loop-forever (yield-generator-exhausted))))

;; yield-datum: any/c -> void
;; Suspends the generator computation, and return
;; the yielded value back.
(define (yield-datum v)
(control-at my-prompt new-save-point
(set! resume-here new-save-point)
(make-yielded-datum v)))

;; yield-exn: exn:fail? -> void
;; Suspends computation, signals a user error.
(define (yield-exn exn)
(control-at my-prompt new-save-point
(set! resume-here new-save-point)
(make-yielded-exn exn)))

;; yield-exn: -> void
;; Suspends computation, signals generator exhaustion.
(define (yield-generator-exhausted)
(control-at my-prompt new-save-point
(set! resume-here new-save-point)
(make-yielded-exn
(make-exn:fail:generator-exhausted
(string->immutable-string
(format "generator ~a exhausted"
(object-name seed-function)))
cont-marks)))))

(make-generator toplevel)))

(provide/contract (generator-next
(case->
[generator? . -> . any]
[generator? (exn:fail? . -> . any) . -> . any]
[generator? (exn:fail? . -> . any) any/c . -> . any])))
;; generator-next: (generator-of X) -> X
;; generator-next: (generator-of X) (exn:fail -> Y) -> (union X Y)
;; generator-next: (generator-of X) (exn:fail -> Y) any -> (union X Y)
;;
;; Returns the next element in the generator.
;;
;; With the first form, if the generator is exhausted,
;; raises exn:fail:generator-exhausted.
;;
;; With the second form, if the generator is exhaused,
;; returns the value of applying the fail-function.
;;
;; The third form also allows one to communicate a value back to
;; the generator at the resumption point.
(define generator-next
(case-lambda
[(gen)
(generator-next gen (lambda (exn) (raise exn)) (void))]

[(gen exhausted-function)
(generator-next gen exhausted-function (void))]

[(gen exhausted-function resumption-val)
(with-handlers ([exn:fail:generator-exhausted?
(lambda (exn)
(exhausted-function exn))])
(gen resumption-val))]))

;; leftmost-identifier: syntax -> (union #f syntax)
;; Looks for the first leftmost identifier syntax object in stx.
;; If one can't be found, returns #f.
(define-for-syntax (leftmost-identifier stx)
(cond
[(identifier? stx) stx]
[(stx-null? stx) #f]
[(stx-pair? stx)
(or (leftmost-identifier (stx-car stx))
(leftmost-identifier (stx-cdr stx)))]
[else #f]))

;  (provide define-generator*)
;  ;; Hygienic macro for writing generator functions.
;  ;;
;  ;; For example,
;  ;;
;  ;; ;; nums: number -> (generator-of number)
;  ;; (define-generator* (nums n) yield
;  ;;   (let loop ((n n))
;  ;;     (yield n)
;  ;;     (loop (add1 n))))
;  ;;
;  ;; produces a generator function. when called, returns a
;  ;; generator that yields successive numbers, starting from n.
;  ;;
;  ;;
;  ;; We also support mzscheme's curry notation, so:
;  ;;
;  ;; (define-generator* ((foo x) y) ba!
;  ;;   (ba! x)
;  ;;   (ba! y))
;  ;;
;  ;; will also do the expected thing.
;  (define-syntax (define-generator* stx)
;    (syntax-case stx ()
;      [(_ (generator-name-or-curried-form args ...) yield-kw
;          body body-rest ...)
;       (unless (identifier? #'yield-kw)
;         (raise-syntax-error #f "yield-kw must be an identifier"
;                             stx #'yield-kw))
;       (with-syntax
;           ([function-name
;             (leftmost-identifier
;              (syntax/loc stx generator-name-or-curried-form))])
;         (syntax/loc stx
;           (define (generator-name-or-curried-form args ...)
;             ;; we wrap an additional LET here so that OBJECT-NAME
;             ;; can derive a good name.
;             (let ([function-name
;                    (lambda (yield-kw)
;                      body body-rest ...)])
;               (make-generator
;                function-name)))))]
;      [else (raise-syntax-error
;             #f
;             "expected (define-generator* (name args ...) yield-kw body ...)"
;             stx)]))

(provide yield)
;; YIELD syntactic parameter; we syntax-parameterize this in
;; define-generator.
(define-syntax-parameter yield
(lambda (stx)
(raise-syntax-error
#f
"used outside of the context of a DEFINE-GENERATOR"
stx)))

(provide define-generator)
;; Macro for writing generator functions with convenient
;; yield keyword.  Cooperates with the YIELD syntax parameter
;; above.
;;
;; For example,
;;
;; ;; nums: number -> (generator-of number)
;; (define-generator (nums n)
;;   (let loop ((n n))
;;     (yield n)
;;     (loop (add1 n))))
;;
;; produces a generator function.  When called, returns a
;; generator that yields successive numbers, starting from n.
;;
;; We also support mzscheme's curry notation, so:
;;
;; (define-generator ((flip-flop x) y)
;;   (let loop ()
;;     (yield x)
;;     (yield y)
;;     (loop)))
;;
;; will also do the expected thing.
(define-syntax (define-generator stx)
(syntax-case stx ()
[(_ (name-or-curried . args) body body-rest ...)
(with-syntax
([function-name (leftmost-identifier #'name-or-curried)])
(syntax/loc stx
(define (name-or-curried . args)
;; we wrap an additional LET here so that OBJECT-NAME
;; can derive a good name.
(let
([function-name
(lambda (real-yield)
(with-yield-rebound
real-yield
body body-rest ...))])
(ext:make-generator function-name)))))]
[else (raise-syntax-error
#f
"expected (define-generator (name args ...) body ...)"
stx)]))

;; with-yield-rebound
;; in the dynamic context of with-yield-rebound, real-yield
;; replaces yield in the body.
(define-syntax (with-yield-rebound stx)
(syntax-case stx ()
[(_ real-yield body ...)
;; Within the context of this body, we establish
;; the YIELD syntax parameter to be redirected to
;; REAL-YIELD.
(syntax/loc stx
(syntax-parameterize
[(yield (lambda (stx)
(syntax-case stx ()
[_ (identifier? stx)
(syntax/loc stx real-yield)]
[(yield value)
(syntax/loc stx (real-yield value))])))]
body ...))]))

(provide/contract
(generator-fold ((any/c any/c . -> . any) any/c generator?
. -> . any)))
;; generator-fold: (X Y -> Y) Y (generator-of X)
;; Applies a fold of f across all the elements in the generator gen,
;; using the initial accumulator acc.
(define (generator-fold f initial-acc gen)
(let loop ((acc initial-acc))
(let-values ([(finished? new-acc)
(with-handlers
([exn:fail:generator-exhausted?
(lambda (_) (values #t acc))])
(values #f (f (gen) acc)))])
(cond
[finished? acc]
[else (loop new-acc)]))))

(provide/contract
(generator-for-each ((any/c . -> . any) generator? . -> . any)))
;; generator-for-each: (X -> any) (generator-of X) -> any
;; Applies f along every elements in generator gen.
(define (generator-for-each f gen)
(generator-fold (lambda (x acc) (f x)) (void) gen))

(provide/contract (list/gen (-> (listof any/c) generator?)))
;; list/gen: (listof A) -> (-> A)
;; Given a list, returns a function iterates over
;; the elements of lst.
;; When the list is exhausted, raises exn:fail:generator-exhausted.
(define-generator (list/gen lst)
(for-each yield lst))

(provide/contract (list->flattened/gen (-> (listof any/c) generator?)))
;; list->flattened/gen: (listof X) -> (generator-of datum)
;; Given a nested list of datums,
;; returns a generator that only returns datums.
(define-generator (list->flattened/gen lst)
(let loop ((datum lst))
(cond
[(empty? datum) 'done]
[(pair? datum) (loop (first datum))
(loop (rest datum))]
[else (yield datum)]))))