From b3499cde198c263b9c555fb564f2e3bd48b4c6e0 Mon Sep 17 00:00:00 2001
From: Neil Toronto <ntoronto@racket-lang.org>
Date: Sun, 14 Oct 2012 15:27:29 -0400
Subject: [PATCH] Added working for/vector:, for*/vector:, for/flvector:,
 for*/flvector:

Closes PR13185.

original commit: 6f52be186b83fdb4340fa4eb772a3e6754050e71
---
 .../tests/typed-racket/succeed/for-vector.rkt | 309 ++++++++++++++++++
 collects/typed-racket/base-env/prims.rkt      | 129 +++++++-
 2 files changed, 433 insertions(+), 5 deletions(-)
 create mode 100644 collects/tests/typed-racket/succeed/for-vector.rkt

diff --git a/collects/tests/typed-racket/succeed/for-vector.rkt b/collects/tests/typed-racket/succeed/for-vector.rkt
new file mode 100644
index 00000000..d1081c3e
--- /dev/null
+++ b/collects/tests/typed-racket/succeed/for-vector.rkt
@@ -0,0 +1,309 @@
+#lang typed/racket
+
+(require (for-syntax syntax/parse)
+         racket/unsafe/ops
+         typed/rackunit
+         racket/flonum)
+
+(check-equal? (for/vector: #:length 4 ([i  (in-range 4)]) : Float
+                (real->double-flonum i))
+              (vector 0.0 1.0 2.0 3.0))
+
+(check-exn exn? (λ () (for/vector: #:length 4 ([i  (in-range 0)]) : Float
+                        (real->double-flonum i))))
+
+(check-equal? (for/vector: #:length 4 () : Float
+                1.2)
+              (vector 1.2 1.2 1.2 1.2))
+
+(check-equal? (for/vector: #:length 4 ([i  (in-range 2)]) : Float
+                (real->double-flonum i))
+              (vector 0.0 1.0 0.0 0.0))
+
+(check-equal? (for/vector: #:length 4 ([i  (in-range 5)]) : Float
+                (real->double-flonum i))
+              (vector 0.0 1.0 2.0 3.0))
+
+(check-equal? (for/vector: #:length 0 ([i  (in-range 5)]) : Float
+                (real->double-flonum i))
+              (vector))
+
+(check-equal? (for/vector: ([i  (in-range 4)]) : Float
+                (real->double-flonum i))
+              (vector 0.0 1.0 2.0 3.0))
+
+(check-equal? (for/vector: () : Float 1.2)
+              (vector 1.2))
+
+(check-equal? (for/vector: ([i  (in-range 0)]) : Float
+                (real->double-flonum i))
+              (vector))
+
+(check-equal? (for/vector: #:length 4 ([x  (in-range 2)]
+                                       #:when #t
+                                       [y  (in-range 2)]) : Float
+                (real->double-flonum (+ x y)))
+              (vector 0.0 1.0 1.0 2.0))
+
+(check-exn exn? (λ () (for/vector: #:length 4 ([x  (in-range 0)]
+                                               #:when #t
+                                               [y  (in-range 2)]) : Float
+                        (real->double-flonum (+ x y)))))
+
+(check-equal? (for/vector: #:length 4 ([x  (in-range 2)]
+                                       #:when #t
+                                       [y  (in-range 1)]) : Float
+                (real->double-flonum (+ x y)))
+              (vector 0.0 1.0 0.0 0.0))
+
+(check-equal? (for/vector: #:length 4 ([x  (in-range 2)]
+                                       #:when #t
+                                       [y  (in-range 3)]) : Float
+                (real->double-flonum (+ x y)))
+              (vector 0.0 1.0 2.0 1.0))
+
+(check-equal? (for/vector: #:length 0 ([x  (in-range 2)]
+                                       #:when #t
+                                       [y  (in-range 3)]) : Float
+                (real->double-flonum (+ x y)))
+              (vector))
+
+(check-equal? (for/vector: ([x  (in-range 2)]
+                            #:when #t
+                            [y  (in-range 2)]) : Float
+                (real->double-flonum (+ x y)))
+              (vector 0.0 1.0 1.0 2.0))
+
+(check-equal? (for/vector: ([x  (in-range 0)]
+                            #:when #t
+                            [y  (in-range 2)]) : Float
+                (real->double-flonum (+ x y)))
+              (vector))
+
+;; ===================================================================================================
+;; for*/vector:
+
+(check-equal? (for*/vector: #:length 4 ([i  (in-range 4)]) : Float
+                (real->double-flonum i))
+              (vector 0.0 1.0 2.0 3.0))
+
+(check-exn exn? (λ () (for*/vector: #:length 4 ([i  (in-range 0)]) : Float
+                        (real->double-flonum i))))
+
+(check-equal? (for*/vector: #:length 4 () : Float
+                1.2)
+              (vector 1.2 1.2 1.2 1.2))
+
+(check-equal? (for*/vector: #:length 4 ([i  (in-range 2)]) : Float
+                (real->double-flonum i))
+              (vector 0.0 1.0 0.0 0.0))
+
+(check-equal? (for*/vector: #:length 4 ([i  (in-range 5)]) : Float
+                (real->double-flonum i))
+              (vector 0.0 1.0 2.0 3.0))
+
+(check-equal? (for*/vector: #:length 0 ([i  (in-range 5)]) : Float
+                (real->double-flonum i))
+              (vector))
+
+(check-equal? (for*/vector: ([i  (in-range 4)]) : Float
+                (real->double-flonum i))
+              (vector 0.0 1.0 2.0 3.0))
+
+(check-equal? (for*/vector: () : Float 1.2)
+              (vector 1.2))
+
+(check-equal? (for*/vector: ([i  (in-range 0)]) : Float
+                (real->double-flonum i))
+              (vector))
+
+(check-equal? (for*/vector: #:length 4 ([x  (in-range 2)]
+                                        [y  (in-range 2)]) : Float
+                (real->double-flonum (+ x y)))
+              (vector 0.0 1.0 1.0 2.0))
+
+(check-exn exn? (λ () (for*/vector: #:length 4 ([x  (in-range 0)]
+                                                [y  (in-range 2)]) : Float
+                        (real->double-flonum (+ x y)))))
+
+(check-equal? (for*/vector: #:length 4 ([x  (in-range 2)]
+                                        [y  (in-range 1)]) : Float
+                (real->double-flonum (+ x y)))
+              (vector 0.0 1.0 0.0 0.0))
+
+(check-equal? (for*/vector: #:length 4 ([x  (in-range 2)]
+                                        [y  (in-range 3)]) : Float
+                (real->double-flonum (+ x y)))
+              (vector 0.0 1.0 2.0 1.0))
+
+(check-equal? (for*/vector: #:length 0 ([x  (in-range 2)]
+                                        [y  (in-range 3)]) : Float
+                (real->double-flonum (+ x y)))
+              (vector))
+
+(check-equal? (for*/vector: ([x  (in-range 2)]
+                             [y  (in-range 2)]) : Float
+                (real->double-flonum (+ x y)))
+              (vector 0.0 1.0 1.0 2.0))
+
+(check-equal? (for*/vector: ([x  (in-range 0)]
+                             [y  (in-range 2)]) : Float
+                (real->double-flonum (+ x y)))
+              (vector))
+
+;; ===================================================================================================
+;; for/flvector:
+
+(define-syntax-rule (test-flvector a b) ; for some reason, check-equal? doesn't work below
+  (unless (equal? a b)
+    (error "bad")))
+
+(test-flvector (for/flvector: #:length 4 ([i  (in-range 4)])
+                (real->double-flonum i))
+              (flvector 0.0 1.0 2.0 3.0))
+
+(test-flvector (for/flvector: #:length 4 ([i  (in-range 0)])
+                (real->double-flonum i))
+              (flvector 0.0 0.0 0.0 0.0))
+
+(test-flvector (for/flvector: #:length 4 ()
+                1.2)
+              (flvector 1.2 0.0 0.0 0.0))
+
+(test-flvector (for/flvector: #:length 4 ([i  (in-range 2)])
+                (real->double-flonum i))
+              (flvector 0.0 1.0 0.0 0.0))
+
+(test-flvector (for/flvector: #:length 4 ([i  (in-range 5)])
+                (real->double-flonum i))
+              (flvector 0.0 1.0 2.0 3.0))
+
+(test-flvector (for/flvector: #:length 0 ([i  (in-range 5)])
+                (real->double-flonum i))
+              (flvector))
+
+(test-flvector (for/flvector: ([i  (in-range 4)])
+                (real->double-flonum i))
+              (flvector 0.0 1.0 2.0 3.0))
+
+(test-flvector (for/flvector: () 1.2)
+              (flvector 1.2))
+
+(test-flvector (for/flvector: ([i  (in-range 0)])
+                (real->double-flonum i))
+              (flvector))
+
+(test-flvector (for/flvector: #:length 4 ([x  (in-range 2)]
+                                         #:when #t
+                                         [y  (in-range 2)])
+                (real->double-flonum (+ x y)))
+              (flvector 0.0 1.0 1.0 2.0))
+
+(test-flvector (for/flvector: #:length 4 ([x  (in-range 0)]
+                                         #:when #t
+                                         [y  (in-range 2)])
+                (real->double-flonum (+ x y)))
+              (flvector 0.0 0.0 0.0 0.0))
+
+(test-flvector (for/flvector: #:length 4 ([x  (in-range 2)]
+                                         #:when #t
+                                         [y  (in-range 1)])
+                (real->double-flonum (+ x y)))
+              (flvector 0.0 1.0 0.0 0.0))
+
+(test-flvector (for/flvector: #:length 4 ([x  (in-range 2)]
+                                         #:when #t
+                                         [y  (in-range 3)])
+                (real->double-flonum (+ x y)))
+              (flvector 0.0 1.0 2.0 1.0))
+
+(test-flvector (for/flvector: #:length 0 ([x  (in-range 2)]
+                                         #:when #t
+                                         [y  (in-range 3)])
+                (real->double-flonum (+ x y)))
+              (flvector))
+
+(test-flvector (for/flvector: ([x  (in-range 2)]
+                              #:when #t
+                              [y  (in-range 2)])
+                (real->double-flonum (+ x y)))
+              (flvector 0.0 1.0 1.0 2.0))
+
+(test-flvector (for/flvector: ([x  (in-range 0)]
+                              #:when #t
+                              [y  (in-range 2)])
+                (real->double-flonum (+ x y)))
+              (flvector))
+
+;; ===================================================================================================
+;; for*/flvector:
+
+(test-flvector (for*/flvector: #:length 4 ([i  (in-range 4)])
+                (real->double-flonum i))
+              (flvector 0.0 1.0 2.0 3.0))
+
+(test-flvector (for*/flvector: #:length 4 ([i  (in-range 0)])
+                (real->double-flonum i))
+              (flvector 0.0 0.0 0.0 0.0))
+
+(test-flvector (for*/flvector: #:length 4 ()
+                1.2)
+              (flvector 1.2 0.0 0.0 0.0))
+
+(test-flvector (for*/flvector: #:length 4 ([i  (in-range 2)])
+                (real->double-flonum i))
+              (flvector 0.0 1.0 0.0 0.0))
+
+(test-flvector (for*/flvector: #:length 4 ([i  (in-range 5)])
+                (real->double-flonum i))
+              (flvector 0.0 1.0 2.0 3.0))
+
+(test-flvector (for*/flvector: #:length 0 ([i  (in-range 5)])
+                (real->double-flonum i))
+              (flvector))
+
+(test-flvector (for*/flvector: ([i  (in-range 4)])
+                (real->double-flonum i))
+              (flvector 0.0 1.0 2.0 3.0))
+
+(test-flvector (for*/flvector: () 1.2)
+              (flvector 1.2))
+
+(test-flvector (for*/flvector: ([i  (in-range 0)])
+                (real->double-flonum i))
+              (flvector))
+
+(test-flvector (for*/flvector: #:length 4 ([x  (in-range 2)]
+                                          [y  (in-range 2)])
+                (real->double-flonum (+ x y)))
+              (flvector 0.0 1.0 1.0 2.0))
+
+(test-flvector (for*/flvector: #:length 4 ([x  (in-range 0)]
+                                          [y  (in-range 2)])
+                (real->double-flonum (+ x y)))
+              (flvector 0.0 0.0 0.0 0.0))
+
+(test-flvector (for*/flvector: #:length 4 ([x  (in-range 2)]
+                                          [y  (in-range 1)])
+                (real->double-flonum (+ x y)))
+              (flvector 0.0 1.0 0.0 0.0))
+
+(test-flvector (for*/flvector: #:length 4 ([x  (in-range 2)]
+                                          [y  (in-range 3)])
+                (real->double-flonum (+ x y)))
+              (flvector 0.0 1.0 2.0 1.0))
+
+(test-flvector (for*/flvector: #:length 0 ([x  (in-range 2)]
+                                          [y  (in-range 3)])
+                (real->double-flonum (+ x y)))
+              (flvector))
+
+(test-flvector (for*/flvector: ([x  (in-range 2)]
+                               [y  (in-range 2)])
+                (real->double-flonum (+ x y)))
+              (flvector 0.0 1.0 1.0 2.0))
+
+(test-flvector (for*/flvector: ([x  (in-range 0)]
+                               [y  (in-range 2)])
+                (real->double-flonum (+ x y)))
+              (flvector))
diff --git a/collects/typed-racket/base-env/prims.rkt b/collects/typed-racket/base-env/prims.rkt
index 6abde03f..afb906c7 100644
--- a/collects/typed-racket/base-env/prims.rkt
+++ b/collects/typed-racket/base-env/prims.rkt
@@ -52,7 +52,9 @@ This file defines two sorts of primitives. All of them are provided into any mod
           "../utils/tc-utils.rkt"
           "../types/utils.rkt"
           "for-clauses.rkt")
-         "../types/numeric-predicates.rkt")
+         "../types/numeric-predicates.rkt"
+         racket/unsafe/ops
+         racket/vector)
 (provide index?) ; useful for assert, and racket doesn't have it
 
 (begin-for-syntax 
@@ -752,8 +754,6 @@ This file defines two sorts of primitives. All of them are provided into any mod
   (for/or: for/or)
   (for/first: for/first)
   (for/last: for/last)
-  (for/vector: for/vector)
-  (for/flvector: for/flvector)
   (for/product: for/product))
 
 ;; Unlike with the above, the inferencer can handle any number of #:when
@@ -843,8 +843,6 @@ This file defines two sorts of primitives. All of them are provided into any mod
   (for*/or: for*/or)
   (for*/first: for*/first)
   (for*/last: for*/last)
-  (for*/vector: for*/vector)
-  (for*/flvector: for*/flvector)
   (for*/product: for*/product))
 
 ;; Like for/lists: and for/fold:, the inferencer can handle these correctly.
@@ -984,3 +982,124 @@ This file defines two sorts of primitives. All of them are provided into any mod
      #'(quote-syntax (typecheck-fail-internal orig "Incomplete case coverage" var))]
     [(_ orig)
      #'(quote-syntax (typecheck-fail-internal orig "Incomplete case coverage" #f))]))
+
+(define-syntax (base-for/vector stx)
+  (syntax-case stx ()
+    [(name for ann T K #:length n-expr #:fill fill-expr (clauses ...) body-expr)
+     (syntax/loc stx
+       (call/ec
+        (ann (λ (break)
+               (define n n-expr)
+               (define vs (ann (make-vector n fill-expr) T))
+               (define i 0)
+               (for (clauses ...)
+                 (unsafe-vector-set! vs i body-expr)
+                 (set! i (unsafe-fx+ i 1))
+                 (when (i . unsafe-fx>= . n) (break vs)))
+               vs)
+             K)))]
+    [(name for ann T K #:length n-expr (clauses ...) body-expr)
+     (syntax/loc stx
+       (let ([n n-expr])
+         (define vs
+           (call/ec
+            (ann (λ (break)
+                   (define vs (ann (vector) T))
+                   (define i 0)
+                   (for (clauses ...)
+                     (define v body-expr)
+                     (cond [(unsafe-fx= i 0)  (define new-vs (ann (make-vector n v) T))
+                                              (set! vs new-vs)]
+                           [else  (unsafe-vector-set! vs i v)])
+                     (set! i (unsafe-fx+ i 1))
+                     (when (i . unsafe-fx>= . n) (break vs)))
+                   vs)
+                 K)))
+         (cond [(= (vector-length vs) n)  vs]
+               [else
+                ;; Only happens when n > 0 and vs = (vector)
+                (raise-result-error 'name (format "~e-element vector" n) vs)])))]
+    [(_ for ann T K (clauses ...) body-expr)
+     (syntax/loc stx
+       (let ()
+         (define n 0)
+         (define vs (ann (vector) T))
+         (define i 0)
+         (for (clauses ...)
+           (define v body-expr)
+           (cond [(unsafe-fx= i n)  (define new-n (max 4 (unsafe-fx* 2 n)))
+                                    (define new-vs (ann (make-vector new-n v) T))
+                                    (vector-copy! new-vs 0 vs)
+                                    (set! n new-n)
+                                    (set! vs new-vs)]
+                 [else  (unsafe-vector-set! vs i v)])
+           (set! i (unsafe-fx+ i 1)))
+         (vector-copy vs 0 i)))]))
+
+(define-for-syntax (base-for/vector: stx for:)
+  (syntax-parse stx #:literals (:)
+    [(name (~optional (~seq : T:expr))
+           (~optional (~seq #:length n-expr:expr))
+           (~optional (~seq #:fill fill-expr:expr))
+           (clauses ...)
+           (~optional (~seq : A:expr))
+           body ...+)
+     (let ([T  (attribute T)]
+           [A  (attribute A)])
+       (with-syntax ([(maybe-length ...)  (if (attribute n-expr) #'(#:length n-expr) #'())]
+                     [(maybe-fill ...)  (if (attribute fill-expr) #'(#:fill fill-expr) #'())]
+                     [body-expr  (if A #`(ann (let () body ...) #,A) #'(let () body ...))]
+                     [T  (cond [(and T A)  #`(U #,T (Vectorof #,A))]
+                               [T  T]
+                               [A  #`(Vectorof #,A)]
+                               [else  #'(Vectorof Any)])])
+         (quasisyntax/loc stx
+           (base-for/vector #,for: ann T ((T -> Nothing) -> T)
+                            maybe-length ... maybe-fill ... (clauses ...) body-expr))))]))
+
+(define-syntax (for/vector: stx)
+  (base-for/vector: stx #'for:))
+
+(define-syntax (for*/vector: stx)
+  (base-for/vector: stx #'for*:))
+
+(define-syntax (base-for/flvector: stx)
+  (syntax-parse stx
+    [(_ for: #:length n-expr:expr (clauses ...) body ...+)
+     (syntax/loc stx
+       (let: ([n : Integer  n-expr])
+         (cond [(n . > . 0)
+                (define xs (make-flvector n))
+                (define: i : Nonnegative-Fixnum 0)
+                (let/ec: break : Void
+                  (for: (clauses ...)
+                    (unsafe-flvector-set! xs i (let () body ...))
+                    (set! i (unsafe-fx+ i 1))
+                    (when (i . unsafe-fx>= . n) (break (void)))))
+                xs]
+               [else  (flvector)])))]
+    [(_ for: (clauses ...) body ...+)
+     (syntax/loc stx
+       (let ()
+         (define n 4)
+         (define xs (make-flvector 4))
+         (define i 0)
+         (for: (clauses ...)
+           (let: ([x : Float  (let () body ...)])
+             (cond [(unsafe-fx= i n)  (define new-n (unsafe-fx* 2 n))
+                                      (define new-xs (make-flvector new-n x))
+                                      (let: loop : Void ([i : Nonnegative-Fixnum 0])
+                                        (when (i . unsafe-fx< . n)
+                                          (unsafe-flvector-set! new-xs i (unsafe-flvector-ref xs i))
+                                          (loop (unsafe-fx+ i 1))))
+                                      (set! n new-n)
+                                      (set! xs new-xs)]
+                   [else  (unsafe-flvector-set! xs i x)]))
+           (set! i (unsafe-fx+ i 1)))
+         (flvector-copy xs 0 i)))]))
+
+(define-syntax-rule (for/flvector: e ...)
+  (base-for/flvector: for: e ...))
+
+(define-syntax-rule (for*/flvector: e ...)
+  (base-for/flvector: for*: e ...))