Make take/take-while/take-until fiber-aware

src/boot/boot.janet

@@ -1031,30 +1031,60 @@
     (set k (next ind k)))
   ret)
 
-(defn take
+(defn- resume-n
-  "Take the first n elements of an indexed or bytes type. Returns a new tuple or string, respectively."
+  [n fib]
-  [n ind]
+  (def res @[])
-  (def use-str (bytes? ind))
+  (var taken 0)
-  (def f (if use-str string/slice tuple/slice))
+  (while (and (< taken n) (fiber/can-resume? fib))
+    (let [elem (resume fib)]
+      (+= taken 1)
+      (array/push res elem)))
+  res)
+
+(defn- slice-n
+  [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))
 
-(defn take-until
+(defn take
-  "Same as `(take-while (complement pred) ind)`."
+  "Take the first n elements of a fiber, indexed or bytes type. Returns a new array, tuple or string, respectively."
-  [pred ind]
+  [n ind]
-  (def use-str (bytes? ind))
+  (cond
-  (def f (if use-str string/slice tuple/slice))
+    (fiber? ind) (resume-n n ind)
+    (bytes? ind) (slice-n string/slice n ind)
+    (slice-n tuple/slice n ind)))
+
+(defn- resume-until
+  [pred fib]
+  (def res @[])
+  (while (fiber/can-resume? fib)
+    (let [elem (resume fib)]
+      (if (pred elem)
+        (break)
+        (array/push res elem))))
+  res)
+
+(defn- slice-until
+  [f pred ind]
   (def len (length ind))
   (def i (find-index pred ind))
   (def end (if (nil? i) len i))
   (f ind 0 end))
 
+(defn take-until
+  "Same as `(take-while (complement pred) ind)`."
+  [pred ind]
+  (cond
+    (fiber? ind) (resume-until pred ind)
+    (bytes? ind) (slice-until string/slice pred ind)
+    (slice-until tuple/slice pred ind)))
+
 (defn take-while
-  `Given a predicate, take only elements from an indexed or bytes type that satisfy
+  `Given a predicate, take only elements from a fiber, indexed or bytes type that satisfy
-  the predicate, and abort on first failure. Returns a new tuple or string, respectively.`
+  the predicate, and abort on first failure. Returns a new array, tuple or string, respectively.`
   [pred ind]
   (take-until (complement pred) ind))
 

test/suite0005.janet

@@ -51,6 +51,10 @@
 (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 3 (generate [x :in [1 2 3 4 5]] x)) @[1 2 3]) "take from fiber")
+# NB: repeatedly resuming a fiber created with `generate` includes a `nil` as
+# the final element. Thus a generate of 2 elements will create an array of 3.
+(assert (= (length (take 4 (generate [x :in [1 2]] x))) 3) "take from short fiber")
 (assert-error :invalid-type (take 3 {}) "take 6")
 
 # take-until
@@ -61,6 +65,8 @@
 (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")
+(assert (deep= (take-until |(= $ 115) (generate [x :in "books"] x))
+               @[98 111 111 107]) "take-until from fiber")
 
 # take-while
 
@@ -69,6 +75,8 @@
 (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")
+(assert (deep= (take-while neg? (generate [x :in  @[-1 1 -2]] x))
+               @[-1]) "take-while from fiber")
 
 # drop