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

 
  ;; We raise exception 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 ()))
  
  
  ;; 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))
  
  

  ;; 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?)))

  
  (define-syntax (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)
    (let
        ([saved-point (void)]
         [caller (void)]
         [caller-continuation-marks (void)])
      (letrec ([entry-point 
                (lambda ()
                  (set-caller-context!)
                  (let ([yielded
                         (call/cc (lambda (k)
                                    (set! caller k)
                                    (cond 
                                      [(restoring-from-save-point?) 
                                       (restore-from-save-point!)]
                                      [else (do-work)])))])
                    (match yielded
                      ((struct yielded-datum (datum)) datum)
                      ((struct yielded-exn (exn)) (raise exn)))))]
               [restoring-from-save-point?
                (lambda () (not (void? saved-point)))]
               [restore-from-save-point!
                (lambda ()
                  (saved-point 'resume))]
               [set-save-point!
                (lambda (f) (set! saved-point f))]
               [set-caller-context!
                (lambda () 
                  (set! caller-continuation-marks
                        (current-continuation-marks)))]
               [do-work 
                (lambda ()
                  (with-handlers ((exn:fail? 
                                   (lambda (exn)
                                     (forever (yield-exn exn)))))
                    (seed-function yield-datum))
                  (forever (yield-generator-exhausted)))]
               [yield-datum 
                (lambda (v)
                  (call/cc 
                   (lambda (k)
                     (set-save-point! k)
                     (caller (make-yielded-datum v)))))]
               [yield-exn
                (lambda (exn)
                  (call/cc
                   (lambda (k)
                     (set-save-point! k)
                     (caller (make-yielded-exn exn)))))]
               [yield-generator-exhausted 
                (lambda ()
                  (call/cc 
                   (lambda (k)
                     (set-save-point! k)
                     (caller
                      (make-yielded-exn
                       (make-exn:fail:generator-exhausted
                        (string->immutable-string 
                         (format "generator ~a exhausted" 
                                 (object-name seed-function)))
                        caller-continuation-marks))))))])
        (make-generator entry-point))))
  
 
  (provide/contract (generator-next 
                     (case->
                      [generator? (exn:fail? . -> . any) . -> . any]
                      [generator? . -> . any])))
  ;; generator-next: (generator-of X) (exn:fail -> Y) -> (union X Y)
  ;; generator-next: (generator-of X) -> X
  ;; Returns the next element in the generator.
  ;;
  ;; With the first form, if the generator is exhaused,
  ;; returns the value of applying the fail-function.
  ;; With the second form, if the generator is exhausted,
  ;; raises exn:fail:generator-exhausted.
  (define generator-next
    (case-lambda 
      [(gen exhausted-function)
       (with-handlers ([exn:fail:generator-exhausted?
                        (lambda (exn)
                          (exhausted-function exn))])
         (gen))]

      [(gen)
       (generator-next gen (lambda (exn) (raise exn)))]))
  
  
  ;; 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 (error 'leftmost-identifier "don't know how to handle ~a" stx)])) 
  
  
  
;  (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 
              (syntax/loc stx generator-name-or-curried-form))])
         (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)]))))