From f52d43e600acf91839106299c2f53534266f7cfe Mon Sep 17 00:00:00 2001
From: Alexis King <lexi.lambda@gmail.com>
Date: Mon, 28 Dec 2015 20:53:47 -0800
Subject: [PATCH] Add for/stream and for*/stream comprehensions to
 racket/stream

Closes #664
---
 .../scribblings/reference/sequences.scrbl     | 20 ++++++++++++++
 .../racket-test-core/tests/racket/stream.rktl |  5 ++++
 racket/collects/racket/stream.rkt             | 26 ++++++++++++++++++-
 3 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/pkgs/racket-doc/scribblings/reference/sequences.scrbl b/pkgs/racket-doc/scribblings/reference/sequences.scrbl
index 321450ec6d..f6aea18d0d 100644
--- a/pkgs/racket-doc/scribblings/reference/sequences.scrbl
+++ b/pkgs/racket-doc/scribblings/reference/sequences.scrbl
@@ -1030,6 +1030,26 @@ stream, but plain lists can be used as streams, and functions such as
   new stream is constructed lazily.
 }
 
+@deftogether[(@defform[(for/stream (for-clause ...) body-or-break ... body)]
+              @defform[(for*/stream (for-clause ...) body-or-break ... body)])]{
+  Iterates like @racket[for/list] and @racket[for*/list], respectively, but the
+  results are lazily collected into a @tech{stream} instead of a list.
+
+  Unlike most @racket[for] forms, these forms are evaluated lazily, so each
+  @racket[body] will not be evaluated until the resulting stream is forced. This
+  allows @racket[for/stream] and @racket[for*/stream] to iterate over infinite
+  sequences, unlike their finite counterparts.
+
+  @examples[#:eval sequence-evaluator
+    (for/stream ([i '(1 2 3)]) (* i i))
+    (stream->list (for/stream ([i '(1 2 3)]) (* i i)))
+    (stream-ref (for/stream ([i '(1 2 3)]) (displayln i) (* i i)) 1)
+    (stream-ref (for/stream ([i (in-naturals)]) (* i i)) 25)
+  ]
+
+  @history[#:added "6.3.0.9"]
+}
+
 @defthing[gen:stream any/c]{
   Associates three methods to a structure type to implement the
   @tech{generic interface} (see @secref["struct-generics"]) for
diff --git a/pkgs/racket-test-core/tests/racket/stream.rktl b/pkgs/racket-test-core/tests/racket/stream.rktl
index 3b54420953..22f419f1b9 100644
--- a/pkgs/racket-test-core/tests/racket/stream.rktl
+++ b/pkgs/racket-test-core/tests/racket/stream.rktl
@@ -57,4 +57,9 @@
 (test '(1 3) stream->list (stream-filter odd? '(1 2 3)))
 (test '(1 a 2 a 3) stream->list (stream-add-between '(1 2 3) 'a))
 
+(test 4 'for/stream (stream-ref (for/stream ([x '(1 2 3)]) (* x x)) 1))
+(test 6 'for*/stream (stream-ref (for*/stream ([x '(1 2 3)] [y '(1 2 3)]) (* x y)) 7))
+(test 1 'for/stream (stream-first (for*/stream ([x '(1 0)]) (/ x))))
+(test 625 'for/stream (stream-ref (for/stream ([x (in-naturals)]) (* x x)) 25))
+
 (report-errs)
diff --git a/racket/collects/racket/stream.rkt b/racket/collects/racket/stream.rkt
index 362b3941fa..6cafb4b97b 100644
--- a/racket/collects/racket/stream.rkt
+++ b/racket/collects/racket/stream.rkt
@@ -4,6 +4,7 @@
          racket/generic
          racket/contract/base
          racket/contract/combinator
+         racket/generator
          (rename-in "private/for.rkt"
                     [stream-ref stream-get-generics])
          "private/sequence.rkt"
@@ -42,7 +43,10 @@
          stream-add-between
          stream-count
          
-         stream/c)
+         stream/c
+
+         for/stream
+         for*/stream)
 
 (define-syntax gen:stream
   (make-generic-info (quote-syntax gen:stream)
@@ -300,3 +304,23 @@
   (if (chaperone-contract? ctc)
       (chaperone-stream/c ctc)
       (impersonator-stream/c ctc)))
+
+;; Stream comprehensions -----------------------------------------------------------------------------
+
+(define-syntaxes (for/stream for*/stream)
+  (let ()
+    (define ((make-for/stream derived-stx) stx)
+      (syntax-case stx ()
+        [(_ clauses . body)
+         (begin
+           (when (null? (syntax->list #'body))
+             (raise-syntax-error (syntax-e #'derived-stx)
+                                 "missing body expression after sequence bindings"
+                                 stx #'body))
+           #`(sequence->stream
+              (in-generator
+               (#,derived-stx #,stx () clauses
+                (yield (let () . body))
+                (values)))))]))
+    (values (make-for/stream #'for/fold/derived)
+            (make-for/stream #'for*/fold/derived))))