racket/generator: add optional #:arity' clause to in-generator'

Allows the use of `in-generator' to produce multiple values in a
position other than immediately within `for' (where the arity
can be inferred).

Closes PR 11662

collects/racket/generator.rkt

@@ -138,20 +138,39 @@
 (define stop-value (gensym))
 
 (define-sequence-syntax in-generator
-  (syntax-rules ()
-    [(_ body0 body ...)
-     (in-producer (generator () body0 body ... stop-value) stop-value)])
   (lambda (stx)
     (syntax-case stx ()
-      [(() (_ body0 body ...))
-       #'[()
-          (in-producer (generator () body0 body ... stop-value) stop-value)]]
-      [((id ...) (_ body0 body ...))
-       (with-syntax ([(stops ...) (syntax-case #'((id stop-value) ...) ()
-                                    [((x v) ...) #'(v ...)])])
-       #'[(id ...)
-          (in-producer (generator () body0 body ... (values stops ...))
-                       (lambda xs (eq? (car xs) stop-value)))])])))
+      [(_ #:arity n body0 body ...)
+       (if (exact-nonnegative-integer? (syntax-e #'n))
+           #'(in-producer (generator () body0 body ... (vector->values (make-vector n stop-value)))
+                          (lambda xs (eq? (car xs) stop-value)))
+           (raise-syntax-error #f 
+                               "expected a literal exact nonnegative integer"
+                               stx
+                               #'n))]
+      [(_ body0 body ...)
+       #'(in-producer (generator () body0 body ... stop-value) stop-value)]))
+  (lambda (stx)
+    (let loop ([stx stx])
+      (syntax-case stx ()
+        [((id ...) (_ #:arity n body0 body ...))
+         (and (exact-integer? #'n)
+              (= (syntax-e #'n (length (syntax->list #'(id ...))))))
+         ;; arity matches, so drop it:
+         (loop #'[((id ...) (_ body0 body ...))])]
+        [(() (_ body0 body ...))
+         #'[()
+            (in-producer (generator () body0 body ... stop-value) stop-value)]]
+        [((id ...) (_ body0 body ...))
+         (with-syntax ([(stops ...) (syntax-case #'((id stop-value) ...) ()
+                                      [((x v) ...) #'(v ...)])])
+           #'[(id ...)
+              (in-producer (generator () body0 body ... (values stops ...))
+                           (lambda xs (eq? (car xs) stop-value)))])]
+        [((id ...) expr)
+         ;; arity mismatch or other syntax error; fall back to expression mode:
+         #'[(id ...) (values expr)]]))))
+        
 
 (define (sequence->generator sequence)
   (generator () (for ([i sequence]) (yield i))))

collects/scribblings/reference/sequences.scrbl

@@ -907,10 +907,13 @@ values from the generator.
     (welcome)
     (welcome)]}
 
-@defform[(in-generator body ...+)]{
+@defform/subs[(in-generator maybe-arity body ...+)
+              ([maybe-arity code:blank 
+                            (code:line #:arity arity-k)])]{
 Produces a @tech{sequence} that encapsulates the @tech{generator} formed by
   @racket[(generator () body ...+)]. The values produced by the
-  generator form the elements of the sequence.
+  generator form the elements of the sequence, except for the last value 
+  produced by the generator (i.e., the values produced by returning).
 
   @examples[#:eval generator-eval
     (for/list ([i (in-generator
@@ -920,22 +923,35 @@ Produces a @tech{sequence} that encapsulates the @tech{generator} formed by
                         (loop (cdr x)))))])
       i)]
 
-  To use an existing generator as a sequence, use @racket[in-producer]
-  with a stop-value known for the generator.
+If @racket[in-generator] is used immediately with a @racket[for] (or
+ @racket[for/list], etc.) binding's right-hand side, then its result
+ arity (i.e., the number of values in each element of the sequence)
+ can be inferred. Otherwise, if the generator produces multiple values
+ for each element, its arity should be declared with an
+ @racket[#:arity arity-k] clause; the @racket[arity-k] must be a
+ literal, exact, non-negative integer.
 
-  @examples[#:eval generator-eval
+  To use an existing generator as a sequence, use @racket[in-producer]
+  with a stop-value known for the generator:
+
+  @interaction[#:eval generator-eval
+    (define abc-generator (generator ()
+                           (for ([x '(a b c)])
+                              (yield x))))
+    (for/list ([i (in-producer abc-generator (void))])
+      i)
     (define my-stop-value (gensym))
     (define my-generator (generator ()
-                           (let loop ([x '(a b c)])
+                           (let loop ([x (list 'a (void) 'c)])
                              (if (null? x)
                                  my-stop-value
                                  (begin
                                    (yield (car x))
                                    (loop (cdr x)))))))
-
     (for/list ([i (in-producer my-generator my-stop-value)])
       i)]}
 
+
 @defproc[(generator-state [g generator?]) symbol?]{
   Returns a symbol that describes the state of the generator.
 

collects/tests/racket/generator.rktl

@@ -138,5 +138,27 @@
   (test-values '(1 2) g)
   (test-values '(1 2) g))
 
+(let ()
+  ;; Yield 1 value, using `in-generator' returned from a function:
+  (define (in-gen-1)
+    (in-generator #:arity 1
+                  (for ([i (in-range 4)])
+                    (yield i))))
+  (test '(0 1 2 3)
+        'gen-1
+        (for/list ([x (in-gen-1)])
+          x)))
+
+(let ()
+  ;; Yield 2 values, using `in-generator' returned from a function:
+  (define (in-gen-2)
+    (in-generator #:arity 2
+                  (for ([i (in-range 4)])
+                    (yield i 0))))
+  (test '((0 . 0) (1 . 0) (2 . 0) (3 . 0))
+        'gen-2
+        (for/list ([(x y) (in-gen-2)])
+          (cons x y))))
+  
 (report-errs)