Merge pull request #156 from curist/take-and-drop

Take and drop

src/boot/boot.janet

@@ -702,13 +702,24 @@
   (def i (find-index pred ind))
   (if (= i nil) nil (get ind i)))
 
+(defn take
+  "Take first n elements in an indexed type. Returns new indexed instance."
+  [n ind]
+  (def use-str (bytes? ind))
+  (def f (if use-str string/slice tuple/slice))
+  # make sure end is in [0, len]
+  (def end (max 0 (min n (length ind))))
+  (f ind 0 end))
+
 (defn take-until
   "Same as (take-while (complement pred) ind)."
   [pred ind]
+  (def use-str (bytes? ind))
+  (def f (if use-str string/slice tuple/slice))
+  (def len (length ind))
   (def i (find-index pred ind))
-  (if i
-    (array/slice ind 0 i)
-    ind))
+  (def end (if (nil? i) len i))
+  (f ind 0 end))
 
 (defn take-while
   "Given a predicate, take only elements from an indexed type that satisfy
@@ -716,13 +727,24 @@
   [pred ind]
   (take-until (complement pred) ind))
 
+(defn drop
+  "Drop first n elements in an indexed type. Returns new indexed instance."
+  [n ind]
+  (def use-str (bytes? ind))
+  (def f (if use-str string/slice tuple/slice))
+  # make sure start is in [0, len]
+  (def start (max 0 (min n (length ind))))
+  (f ind start -1))
+
 (defn drop-until
   "Same as (drop-while (complement pred) ind)."
   [pred ind]
+  (def use-str (bytes? ind))
+  (def f (if use-str string/slice tuple/slice))
   (def i (find-index pred ind))
-  (if i
-    (array/slice ind i)
-    @[]))
+  (def len (length ind))
+  (def start (if (nil? i) len i))
+  (f ind start))
 
 (defn drop-while
   "Given a predicate, remove elements from an indexed type that satisfy

test/suite5.janet

@@ -80,13 +80,49 @@
 (assert-no-error "break 3" (for i 0 10 (if (> i 8) (break i))))
 (assert-no-error "break 4" ((fn [i] (if (> i 8) (break i))) 100))
 
+# take
+
+(assert (deep= (take 0 []) []) "take 1")
+(assert (deep= (take 10 []) []) "take 2")
+(assert (deep= (take 0 [1 2 3 4 5]) []) "take 3")
+(assert (deep= (take 10 [1 2 3]) [1 2 3]) "take 4")
+(assert (deep= (take -1 [:a :b :c]) []) "take 5")
+(assert-error :invalid-type (take 3 {}) "take 6")
+
+# take-until
+
+(assert (deep= (take-until pos? @[]) []) "take-until 1")
+(assert (deep= (take-until pos? @[1 2 3]) []) "take-until 2")
+(assert (deep= (take-until pos? @[-1 -2 -3]) [-1 -2 -3]) "take-until 3")
+(assert (deep= (take-until pos? @[-1 -2 3]) [-1 -2]) "take-until 4")
+(assert (deep= (take-until pos? @[-1 1 -2]) [-1]) "take-until 5")
+(assert (deep= (take-until |(= $ 115) "books") "book") "take-until 6")
+
+# take-while
+
+(assert (deep= (take-while neg? @[]) []) "take-while 1")
+(assert (deep= (take-while neg? @[1 2 3]) []) "take-while 2")
+(assert (deep= (take-while neg? @[-1 -2 -3]) [-1 -2 -3]) "take-while 3")
+(assert (deep= (take-while neg? @[-1 -2 3]) [-1 -2]) "take-while 4")
+(assert (deep= (take-while neg? @[-1 1 -2]) [-1]) "take-while 5")
+
+# drop
+
+(assert (deep= (drop 0 []) []) "drop 1")
+(assert (deep= (drop 10 []) []) "drop 2")
+(assert (deep= (drop 0 [1 2 3 4 5]) [1 2 3 4 5]) "drop 3")
+(assert (deep= (drop 10 [1 2 3]) []) "drop 4")
+(assert (deep= (drop -2 [:a :b :c]) [:a :b :c]) "drop 5")
+(assert-error :invalid-type (drop 3 {}) "drop 6")
+
 # drop-until
 
-(assert (deep= (drop-until pos? @[]) @[]) "drop-until 1")
-(assert (deep= (drop-until pos? @[1 2 3]) @[1 2 3]) "drop-until 2")
-(assert (deep= (drop-until pos? @[-1 -2 -3]) @[]) "drop-until 3")
-(assert (deep= (drop-until pos? @[-1 -2 3]) @[3]) "drop-until 4")
-(assert (deep= (drop-until pos? @[-1 1 -2]) @[1 -2]) "drop-until 5")
+(assert (deep= (drop-until pos? @[]) []) "drop-until 1")
+(assert (deep= (drop-until pos? @[1 2 3]) [1 2 3]) "drop-until 2")
+(assert (deep= (drop-until pos? @[-1 -2 -3]) []) "drop-until 3")
+(assert (deep= (drop-until pos? @[-1 -2 3]) [3]) "drop-until 4")
+(assert (deep= (drop-until pos? @[-1 1 -2]) [1 -2]) "drop-until 5")
+(assert (deep= (drop-until |(= $ 115) "books") "s") "drop-until 6")
 
 # Quasiquote bracketed tuples
 (assert (= (tuple/type ~[1 2 3]) (tuple/type '[1 2 3])) "quasiquote bracket tuples")