take-drop symmetry

Allow `take` from the end of bytes or indexed (as `drop` does).
Allow `drop` from fibers (as `take` does).

src/boot/boot.janet

@@ -1080,42 +1080,30 @@
     (set k (next ind k)))
   ret)
 
-(defn- take-n-fallback
-  [n xs]
-  (def res @[])
-  (when (> n 0)
-    (var left n)
-    (each x xs
-      (array/push res x)
-      (-- left)
-      (if (= 0 left) (break))))
-  res)
-
-(defn- take-until-fallback
-  [pred xs]
-  (def res @[])
-  (each x xs
-    (if (pred x) (break))
-    (array/push res x))
-  res)
-
-(defn- slice-n
+(defn- take-n-slice
   [f n ind]
   (def len (length ind))
-  # make sure end is in [0, len]
-  (def m (if (> n 0) n 0))
-  (def end (if (> m len) len m))
-  (f ind 0 end))
+  (def negn (< n 0))
+  (def start (if negn (max 0 (+ len n)) 0))
+  (def end (if negn len (min n len)))
+  (f ind start end))
 
 (defn take
-  "Take the first n elements of a fiber, indexed or bytes type. Returns a new array, tuple or string, respectively."
+  ``Take the first n elements of a fiber, indexed or bytes type. Returns a new array, tuple or string,
+  respectively. If `n` is negative, takes the last `n` elements instead.``
   [n ind]
   (cond
-    (bytes? ind) (slice-n string/slice n ind)
-    (indexed? ind) (slice-n tuple/slice n ind)
-    (take-n-fallback n ind)))
+    (bytes? ind) (take-n-slice string/slice n ind)
+    (indexed? ind) (take-n-slice tuple/slice n ind)
+    (do
+      (def res @[])
+      (var key nil)
+      (repeat n
+        (if (= nil (set key (next ind key))) (break))
+        (array/push res (in ind key)))
+      res)))
 
-(defn- slice-until
+(defn- take-until-slice
   [f pred ind]
   (def len (length ind))
   (def i (find-index pred ind))
@@ -1126,9 +1114,9 @@
   "Same as `(take-while (complement pred) ind)`."
   [pred ind]
   (cond
-    (bytes? ind) (slice-until string/slice pred ind)
-    (indexed? ind) (slice-until tuple/slice pred ind)
-    (take-until-fallback pred ind)))
+    (bytes? ind) (take-until-slice string/slice pred ind)
+    (indexed? ind) (take-until-slice tuple/slice pred ind)
+    (seq [x :in ind :until (pred x)] x)))
 
 (defn take-while
   `Given a predicate, take only elements from a fiber, indexed, or bytes type that satisfy
@@ -1136,27 +1124,41 @@
   [pred ind]
   (take-until (complement pred) ind))
 
+(defn- drop-n-slice
+  [f n ind]
+  (def len (length ind))
+  (def negn (< n 0))
+  (def start (if negn 0 (min n len)))
+  (def end (if negn (max 0 (+ len n)) len))
+  (f ind start end))
+
 (defn drop
-  ``Drop the first `n elements in an indexed or bytes type. Returns a new tuple or string
+  ``Drop the first `n` elements in an indexed or bytes type. Returns a new tuple or string
   instance, respectively. If `n` is negative, drops the last `n` elements instead.``
   [n ind]
-  (def use-str (bytes? ind))
-  (def f (if use-str string/slice tuple/slice))
+  (cond
+    (bytes? ind) (drop-n-slice string/slice n ind)
+    (indexed? ind) (drop-n-slice tuple/slice n ind)
+    (do
+      (var key nil)
+      (repeat n
+        (if (= nil (set key (next ind key))) (break)))
+      ind)))
+
+(defn- drop-until-slice
+  [f pred ind]
   (def len (length ind))
-  (def negn (>= n 0))
-  (def start (if negn (min n len) 0))
-  (def end (if negn len (max 0 (+ len n))))
-  (f ind start end))
+  (def i (find-index pred ind))
+  (def start (if (nil? i) len i))
+  (f ind start))
 
 (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))
-  (def len (length ind))
-  (def start (if (nil? i) len i))
-  (f ind start))
+  (cond
+    (bytes? ind) (drop-until-slice string/slice pred ind)
+    (indexed? ind) (drop-until-slice tuple/slice pred ind)
+    (do (find pred ind) ind)))
 
 (defn drop-while
   `Given a predicate, remove elements from an indexed or bytes type that satisfy
@@ -4166,7 +4168,8 @@
 
   (defn do-one-file
     [fname]
-    (if-not (has-value? boot/args "image-only") (do
+    (if-not (has-value? boot/args "image-only")
+      (do
         (print "\n/* " fname " */")
         (print "#line 0 \"" fname "\"\n")
         (def source (slurp fname))

test/suite-boot.janet

@@ -439,7 +439,7 @@
 (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 (deep= (take -1 [:a :b :c]) [:c]) "take 5")
 # 34019222c
 (assert (deep= (take 3 (generate [x :in [1 2 3 4 5]] x)) @[1 2 3])
         "take from fiber")
@@ -482,7 +482,6 @@
 (assert (deep= (drop 10 "abc") "") "drop 8")
 (assert (deep= (drop -1 "abc") "ab") "drop 9")
 (assert (deep= (drop -10 "abc") "") "drop 10")
-(assert-error :invalid-type (drop 3 {}) "drop 11")
 
 # drop-until
 # 75dc08f
@@ -493,6 +492,27 @@
 (assert (deep= (drop-until pos? @[-1 1 -2]) [1 -2]) "drop-until 5")
 (assert (deep= (drop-until |(= $ 115) "books") "s") "drop-until 6")
 
+# take-drop symmetry #1178
+(def items-list ['abcde :abcde "abcde" @"abcde" [1 2 3 4 5] @[1 2 3 4 5]])
+
+(each items items-list
+  (def len (length items))
+  (for i 0 (+ len 1)
+    (assert (deep= (take i items) (drop (- i len) items)) (string/format "take-drop symmetry %q %d" items i))
+    (assert (deep= (take (- i) items) (drop (- len i) items)) (string/format "take-drop symmetry %q %d" items i))))
+
+(defn squares []
+  (coro
+    (var [a b] [0 1])
+    (forever (yield a) (+= a b) (+= b 2))))
+
+(def sqr1 (squares))
+(assert (deep= (take 10 sqr1) @[0 1 4 9 16 25 36 49 64 81]))
+(assert (deep= (take 1 sqr1) @[100]) "take fiber next value")
+
+(def sqr2 (drop 10 (squares)))
+(assert (deep= (take 1 sqr2) @[100]) "drop fiber next value")
+
 # Comment macro
 # issue #110 - 698e89aba
 (comment 1)