diff --git a/src/boot/boot.janet b/src/boot/boot.janet index 80251b76..570dd057 100644 --- a/src/boot/boot.janet +++ b/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 diff --git a/test/suite5.janet b/test/suite5.janet index d74eb263..7dced355 100644 --- a/test/suite5.janet +++ b/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")