diff --git a/jolt-core/clojure/core/10-seq.clj b/jolt-core/clojure/core/10-seq.clj index 78a27a7..ac54c6a 100644 --- a/jolt-core/clojure/core/10-seq.clj +++ b/jolt-core/clojure/core/10-seq.clj @@ -23,7 +23,29 @@ (recur (conj ret (first s)) (next s)) (seq ret)))) +(defn- mapcat-step [rs cur] + (lazy-seq + (if cur + (let [s (seq cur)] + (if s + (cons (first s) (mapcat-step rs (rest s))) + (mapcat-step rs nil))) + (let [s (seq rs)] + (if s + (let [c (first s) + sc (seq c)] + (if sc + (cons (first sc) (mapcat-step (rest s) (rest sc))) + (mapcat-step (rest s) nil))) + nil))))) + (defn mapcat ([f] (comp (map f) cat)) - ([f & colls] - (apply concat (apply map f colls)))) + ([f coll] + (mapcat-step (map f coll) nil)) + ([f c1 c2] + (mapcat-step (map f c1 c2) nil)) + ([f c1 c2 c3] + (mapcat-step (map f c1 c2 c3) nil)) + ([f c1 c2 c3 & colls] + (mapcat-step (apply map f c1 c2 c3 colls) nil))) diff --git a/test/integration/lazy-infinite-test.janet b/test/integration/lazy-infinite-test.janet index 0b0948f..56aeb67 100644 --- a/test/integration/lazy-infinite-test.janet +++ b/test/integration/lazy-infinite-test.janet @@ -52,9 +52,7 @@ ["take 3 partition-all 2 range" "(quote ((0 1) (2 3) (4 5)))" "(take 3 (partition-all 2 (range)))"] ["take 3 map-indexed vector range" "(quote ([0 0] [1 1] [2 2]))" "(take 3 (map-indexed vector (range)))"] ["take 3 distinct cycle" "(quote (1 2 3))" "(take 3 (distinct (cycle [1 2 1 3 1])))"] - # NOTE: mapcat on infinite inputs hangs because apply forces the lazy map - # result. This requires Step 4 (fix apply to lazily spread lazy-seqs). - # ["take 6 mapcat dup range" "(quote (0 0 1 1 2 2))" "(take 6 (mapcat (fn [x] [x x]) (range)))"] + ["take 6 mapcat dup range" "(quote (0 0 1 1 2 2))" "(take 6 (mapcat (fn [x] [x x]) (range)))"] ["take 3 take-nth 2 range" "(quote (0 2 4))" "(take 3 (take-nth 2 (range)))"] ["take 3 interpose :x range" "(quote (0 :x 1))" "(take 3 (interpose :x (range)))"] ["take 3 map vector range iterate" "(quote ([0 100] [1 101] [2 102]))" "(take 3 (map vector (range) (iterate inc 100)))"]