Better support for #:when clauses in the for: macros.

collects/typed-scheme/private/for-clauses.rkt

@@ -1,17 +1,44 @@
-#lang scheme/base
+#lang racket/base
 
 (require syntax/parse
-         "annotate-classes.rkt")
+         "annotate-classes.rkt"
+         (for-template racket/base))
 
-(provide for-clause)
+(provide convert-for-clauses)
 
-(define-splicing-syntax-class for-clause
-  ;; single-valued seq-expr
-  (pattern (var:annotated-name seq-expr:expr)
-           #:with (expand ...) (list #'(var.ann-name seq-expr)))
-  ;; multi-valued seq-expr
-  (pattern ((var:annotated-name ...) seq-expr:expr)
-           #:with (expand ...) (list #'((var.ann-name ...) seq-expr)))
-  ;; when clause
-  (pattern (~seq #:when guard:expr)
-           #:with (expand ...) (list #'#:when #'guard)))
+;; we need handle #:when clauses manually because we need to annotate
+;; the type of each nested for
+(define (convert-for-clauses name clauses body ty)
+  (let loop ((clauses clauses))
+    (define-splicing-syntax-class for-clause
+    ;; single-valued seq-expr
+    (pattern (var:annotated-name seq-expr:expr)
+             #:with expand #'(var.ann-name seq-expr))
+    ;; multi-valued seq-expr
+    (pattern ((v:annotated-name ...) seq-expr:expr)
+             #:with expand #'((v.ann-name ...) seq-expr)))
+    (syntax-parse clauses
+      [(head:for-clause next:for-clause ... #:when rest ...)
+       (syntax-property
+        (quasisyntax/loc clauses
+          (#,name
+           (head.expand next.expand ...)
+           #,(loop #'(#:when rest ...))))
+        'type-ascription
+        ty)]
+      [(head:for-clause ...) ; we reached the end
+       (syntax-property
+        (quasisyntax/loc clauses
+          (#,name
+           (head.expand ...)
+           #,@body))
+        'type-ascription
+        ty)]
+      [(#:when guard) ; we end on a #:when clause
+       (quasisyntax/loc clauses
+         (when guard
+           #,@body))]
+      [(#:when guard rest ...)
+       (quasisyntax/loc clauses
+         (when guard
+           #,(loop #'(rest ...))))])))

collects/typed-scheme/private/prims.rkt

@@ -383,14 +383,9 @@ This file defines two sorts of primitives. All of them are provided into any mod
   (lambda (stx)
     (syntax-parse stx #:literals (:)
       [(_ : ty
-          (clause:for-clause ...)
+          clauses
           c:expr ...)
-       (quasisyntax/loc
-           stx
-         (ann (#,name
-               (clause.expand ... ...)
-               c ...)
-              ty))])))
+       (convert-for-clauses name #'clauses #'(c ...) #'ty)])))
 (define-syntax (define-for-variants stx)
   (syntax-parse stx
     [(_ (name untyped-name) ...)