From 82927b6d171c8d0e2d7947417d608b5e8af78657 Mon Sep 17 00:00:00 2001
From: ben <benjaminlgreenman@gmail.com>
Date: Thu, 3 Mar 2016 15:18:36 -0500
Subject: [PATCH] [vector] take+drop

---
 test/vector-fail.rkt | 21 ++++++++++++
 test/vector-pass.rkt | 36 ++++++++++++++++++++
 vector.rkt           | 79 +++++++++++++++++++++++++++++++++-----------
 3 files changed, 117 insertions(+), 19 deletions(-)

diff --git a/test/vector-fail.rkt b/test/vector-fail.rkt
index 15d0053..090a719 100644
--- a/test/vector-fail.rkt
+++ b/test/vector-fail.rkt
@@ -19,10 +19,14 @@
     (let-vector: ([v2 (vector v1)])
       (vector-ref: (vector-ref: v2 0) 1)))
 
+  (vector-set!: (vector 0) -1 0)
+
   (vector-set!: (vector 0) 0 "hello") ;; Strong update
 
   (vector-ref: (vector-map: (lambda (x) x) (vector #t "ha")) 20)
 
+  (vector-ref: (vector 0) -5)
+
   (vector-ref:
     (vector-map: add1 (vector-map: add1 (vector-map: add1 (vector 0 0 0))))
     3)
@@ -39,6 +43,23 @@
     (vector-ref: (vector-append: v2 v) 8))
 
   (vector-ref: (vector->immutable-vector: (vector 1 2 1)) 3)
+
+  (vector-take: (vector) 1)
+  (vector-take: (vector 1 2) 4)
+  (vector-take: (vector 'a) -1)
+
+  (vector-take-right: (vector) 1)
+  (vector-take-right: (vector 1 2) 4)
+  (vector-take-right: (vector 'a) -1)
+
+  (vector-drop: (vector) 1)
+  (vector-drop: (vector 1 2) 4)
+  (vector-drop: (vector 'a) -1)
+
+  (vector-drop-right: (vector) 1)
+  (vector-drop-right: (vector 1 2) 4)
+  (vector-drop-right: (vector 'a) -1)
+
 )))
 
 ;; -----------------------------------------------------------------------------
diff --git a/test/vector-pass.rkt b/test/vector-pass.rkt
index 0fb3105..f511e4f 100644
--- a/test/vector-pass.rkt
+++ b/test/vector-pass.rkt
@@ -262,6 +262,42 @@
         (check-equal? (vector-ref v 2) 9)))
   ;)
 
+  ;(test-suite "vector-take"
+    (test-case "take basis"
+      (let-vector: ([v (vector 2 3 1)])
+        (check-equal? (vector-take: v 3) v)
+        (check-equal? (vector-take: v 2) (vector 2 3))
+        (check-equal? (vector-take: v 1) (vector 2))
+        (check-equal? (vector-take: v 0) (vector))))
+  ;)
+
+  ;(test-suite "vector-take-right"
+    (test-case "take-right basic"
+      (let-vector: ([v (vector 2 3 1)])
+        (check-equal? (vector-take-right: v 3) v)
+        (check-equal? (vector-take-right: v 2) (vector 3 1))
+        (check-equal? (vector-take-right: v 1) (vector 1))
+        (check-equal? (vector-take-right: v 0) (vector))))
+  ;)
+
+  ;(test-suite "vector-drop-right"
+    (test-case "drop-right basic"
+      (let-vector: ([v (vector 2 3 1)])
+        (check-equal? (vector-drop-right: v 0) v)
+        (check-equal? (vector-drop-right: v 1) (vector 2 3))
+        (check-equal? (vector-drop-right: v 2) (vector 2))
+        (check-equal? (vector-drop-right: v 3) (vector))))
+  ;)
+
+  ;(test-suite "vector-drop"
+    (test-case "drop basic"
+      (let-vector: ([v (vector 2 3 1)])
+        (check-equal? (vector-drop: v 0) v)
+        (check-equal? (vector-drop: v 1) (vector 3 1))
+        (check-equal? (vector-drop: v 2) (vector 1))
+        (check-equal? (vector-drop: v 3) (vector))))
+  ;)
+
   ;; -- define-vector:
   (let ()
     (define-vector: v (vector 1 1 2 2))
diff --git a/vector.rkt b/vector.rkt
index fb939d2..8887741 100644
--- a/vector.rkt
+++ b/vector.rkt
@@ -1,6 +1,6 @@
 #lang typed/racket/base
 
-;; TODO integrate with trivial/math to get ints from identifiers
+;; TOOD abstract some designs
 
 (provide
   define-vector:
@@ -14,12 +14,12 @@
   vector->list:
   vector->immutable-vector:
   vector-fill!:
-  ;vector-take
-  ;vector-take-right
-  ;vector-drop
-  ;vector-drop-right
-  ;vector-split-at
-  ;vector-split-at-right
+  vector-take:
+  vector-take-right:
+  vector-drop:
+  vector-drop-right:
+;  vector-split-at:
+;  vector-split-at-right:
 
   ;; --- private
   (for-syntax parse-vector-length)
@@ -32,8 +32,10 @@
     unsafe-vector-set!
     unsafe-vector-ref)
   racket/vector
+  trivial/math
   (for-syntax
     typed/racket/base
+    racket/syntax
     syntax/id-table
     syntax/parse
     syntax/stx
@@ -94,10 +96,10 @@
 
 (define-syntax (vector-ref: stx)
   (syntax-parse stx
-   [(_ v:vector/length i:nat)
-    (unless (< (syntax-e #'i) (syntax-e #'v.length))
-      (vector-bounds-error 'vector-ref: #'v (syntax-e #'i)))
-    (syntax/loc stx (unsafe-vector-ref v.expanded i))]
+   [(_ v:vector/length i:nat/expand)
+    (unless (< (syntax-e #'i.expanded) (syntax-e #'v.length))
+      (vector-bounds-error 'vector-ref: #'v (syntax-e #'i.expanded)))
+    (syntax/loc stx (unsafe-vector-ref v.expanded 'i.expanded))]
    [_:id
     (syntax/loc stx vector-ref)]
    [(_ e* ...)
@@ -105,10 +107,10 @@
 
 (define-syntax (vector-set!: stx)
   (syntax-parse stx
-   [(_ v:vector/length i:nat val)
-    (unless (< (syntax-e #'i) (syntax-e #'v.length))
-      (vector-bounds-error 'vector-set!: #'v (syntax-e #'i)))
-    (syntax/loc stx (unsafe-vector-set! v.expanded i val))]
+   [(_ v:vector/length i:nat/expand val)
+    (unless (< (syntax-e #'i.expanded) (syntax-e #'v.length))
+      (vector-bounds-error 'vector-set!: #'v (syntax-e #'i.expanded)))
+    (syntax/loc stx (unsafe-vector-set! v.expanded 'i.expanded val))]
    [_:id
     (syntax/loc stx vector-set!)]
    [(_ e* ...)
@@ -117,14 +119,14 @@
 (define-syntax (vector-map: stx)
   (syntax-parse stx
    [(_ f v:vector/length)
-    #:with (i* ...) (for/list ([i (in-range (syntax-e #'v.length))]) i)
     #:with f+ (gensym 'f)
     #:with v+ (gensym 'v)
     #:with v++ (syntax-property
                  (if (small-vector-size? (syntax-e #'v.length))
-                   (syntax/loc stx
-                     (let ([f+ f] [v+ v.expanded])
-                       (vector (f+ (unsafe-vector-ref v+ 'i*)) ...)))
+                   (with-syntax ([(i* ...) (for/list ([i (in-range (syntax-e #'v.length))]) i)])
+                     (syntax/loc stx
+                       (let ([f+ f] [v+ v.expanded])
+                         (vector (f+ (unsafe-vector-ref v+ 'i*)) ...))))
                    (syntax/loc stx
                      (let ([f+ f] [v+ v.expanded])
                        (build-vector 'v.length (lambda ([i : Integer])
@@ -236,6 +238,45 @@
    [(_ e* ...)
     (syntax/loc stx (vector->fill! e* ...))]))
 
+(begin-for-syntax (define-syntax-rule (make-slice-op op-name left? take?)
+  (lambda (stx)
+    (syntax-parse stx
+     [(_ v:vector/length n:nat/expand)
+      #:with (lo hi)
+        (if 'take?
+          (if 'left?
+            (list 0 (syntax-e #'n.expanded))
+            (list
+              (- (syntax-e #'v.length) (syntax-e #'n.expanded))
+              (syntax-e #'v.length)))
+          (if 'left?
+            (list (syntax-e #'n.expanded) (syntax-e #'v.length))
+            (list 0 (- (syntax-e #'v.length) (syntax-e #'n.expanded)))))
+      #:with n+ (gensym 'n)
+      #:with v+ (gensym 'v)
+      (unless (<= (syntax-e #'n.expanded) (syntax-e #'v.length))
+        (vector-bounds-error 'op-name #'v
+          (if 'take? (if 'left? (syntax-e #'hi) (syntax-e #'lo))
+                     (if 'left? (syntax-e #'lo) (syntax-e #'hi)))))
+      (syntax-property
+        (syntax/loc stx
+          (let ([v+ v.expanded]
+                [n+ (-: 'hi 'lo)])
+            (build-vector n+ (lambda ([i : Integer]) (unsafe-vector-ref v+ (+: i 'lo))))))
+        vector-length-key
+        (syntax-e #'v.length))]
+     [(_ v n:int/expand)
+      (vector-bounds-error 'op-name #'v (syntax-e #'n.expanded))]
+     [_:id
+      (syntax/loc stx op-name)]
+     [(_ e* (... ...))
+      (syntax/loc stx (op-name e* (... ...)))]))))
+
+(define-syntax vector-take:       (make-slice-op vector-take       #t #t))
+(define-syntax vector-take-right: (make-slice-op vector-take-right #f #t))
+(define-syntax vector-drop-right: (make-slice-op vector-drop-right #f #f))
+(define-syntax vector-drop:       (make-slice-op vector-drop       #t #f))
+
 ;; -----------------------------------------------------------------------------
 
 (define-for-syntax (small-vector-size? n)