mirror of
https://github.com/janet-lang/janet
synced 2024-12-25 07:50:27 +00:00
Make take/take-while/take-until fiber-aware
This commit is contained in:
parent
61769c8f16
commit
34019222c2
@ -1031,30 +1031,60 @@
|
||||
(set k (next ind k)))
|
||||
ret)
|
||||
|
||||
(defn take
|
||||
"Take the first n elements of an indexed or bytes type. Returns a new tuple or string, respectively."
|
||||
[n ind]
|
||||
(def use-str (bytes? ind))
|
||||
(def f (if use-str string/slice tuple/slice))
|
||||
(defn- resume-n
|
||||
[n fib]
|
||||
(def res @[])
|
||||
(var taken 0)
|
||||
(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
|
||||
"Same as `(take-while (complement pred) ind)`."
|
||||
[pred ind]
|
||||
(def use-str (bytes? ind))
|
||||
(def f (if use-str string/slice tuple/slice))
|
||||
(defn take
|
||||
"Take the first n elements of a fiber, indexed or bytes type. Returns a new array, tuple or string, respectively."
|
||||
[n ind]
|
||||
(cond
|
||||
(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
|
||||
the predicate, and abort on first failure. Returns a new tuple or string, respectively.`
|
||||
`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 array, tuple or string, respectively.`
|
||||
[pred ind]
|
||||
(take-until (complement pred) ind))
|
||||
|
||||
|
@ -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
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user