(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)]))))