Add sequence-padded

Author: NoahTheDuke

src/medley/core.cljc

@@ -727,10 +727,30 @@
           (cons (f (first* s1 val) (first* s2 val) (first* s3 val))
                 (map-padded f val (rest s1) (rest s2) (rest s3)))))))
     ([f val c1 c2 c3 & colls]
-     (let [step (fn step [cs]
+     (let [colls (vec (conj colls c3 c2 c1))
+           step (fn step [cs]
                   (lazy-seq
                    (let [ss (mapv seq cs)]
                      (when (some identity ss)
                        (cons (apply f (mapv #(first* % val) ss))
                              (step (mapv rest ss)))))))]
-        (step (conj colls c3 c2 c1))))))
+        (step colls))))
+
+  (defn sequence-padded
+    "Similar to `clojure.core/sequence`, except that it runs until all colls are
+    exhausted, using `val` as the missing value for each exhausted coll."
+    #_{:clj-kondo/ignore [:unused-binding]}
+    ([xform val c1] (sequence xform c1))
+    ([xform val c1 & colls]
+     (let [colls (vec (cons c1 colls))
+           f (xform (completing #(cons %2 %1)))
+           step (fn step [cs]
+                  (let [ss (mapv seq cs)]
+                    (if (some identity ss)
+                      (let [res (apply f nil (mapv #(first* % val) ss))]
+                        (cond (reduced? res) (f (deref res))
+                              (seq? res) (concat res (lazy-seq (step (mapv rest ss))))
+                              :else (step (mapv rest ss))))
+                      (f nil))))]
+       (or (step colls)
+           ())))))

test/medley/core_test.cljc

@@ -572,3 +572,29 @@
             [:missing :missing 4 4]
             [:missing :missing :missing 5]]
            (take 6 (m/map-padded vector :missing [nil nil nil] (range 4) (range 5) (range 10)))))))
+
+(deftest test-sequence-padded
+  (is (= (map + (range 3) (range 4) (range 5) (range 10))
+         [0 4 8]
+         (take 3 (m/sequence-padded (map +) 0 (range 3) (range 4) (range 5) (range 10)))))
+  (is (= [0 4 8 19 28 35 36 37 38 39]
+         (m/sequence-padded (map +) 10 (range 3) (range 4) (range 5) (range 10))))
+  (is (= ()
+         (m/sequence-padded (map +) 10 () () ())))
+  (testing "laziness"
+    (let [state (volatile! [])]
+      (is (= [0 4 8 19 28]
+             (take 5 (m/sequence-padded
+                      (map (fn [a b c d]
+                             (vswap! state conj [a b c d])
+                             (+ a b c d)))
+                      10
+                      (range 3) (range 4) (range 5) (range 10)))))
+      (is (= [[0 0 0 0] [1 1 1 1] [2 2 2 2] [10 3 3 3] [10 10 4 4]]
+             @state))))
+  (testing "handles sequences with nils"
+    (is (= [[nil 0 0 0] [nil 1 1 1] [nil 2 2 2]
+            [:missing 3 3 3]
+            [:missing :missing 4 4]
+            [:missing :missing :missing 5]]
+           (take 6 (m/sequence-padded (map vector) :missing [nil nil nil] (range 4) (range 5) (range 10)))))))