Add map-padded

Author: NoahTheDuke

src/medley/core.cljc

@@ -706,3 +706,31 @@
   (if (next ks)
     (-> (get-in m (butlast ks)) (find (last ks)))
     (find m (first ks))))
+
+(letfn [(first* [s v] (if s (first s) v))]
+  (defn map-padded
+    "Similar to `clojure.core/map`, except that it runs until all colls are
+    exhausted, using `val` as the missing value for each exhausted coll."
+    {:added "<<next>>"}
+    #_{:clj-kondo/ignore [:unused-binding]}
+    ([f val c1] (map f c1))
+    ([f val c1 c2]
+     (lazy-seq
+      (let [s1 (seq c1) s2 (seq c2)]
+        (when (or s1 s2)
+          (cons (f (first* s1 val) (first* s2 val))
+                (map-padded f val (rest s1) (rest s2)))))))
+    ([f val c1 c2 c3]
+     (lazy-seq
+      (let [s1 (seq c1) s2 (seq c2) s3 (seq c3)]
+        (when (or s1 s2 s3)
+          (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]
+                  (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))))))

test/medley/core_test.cljc

@@ -547,3 +547,28 @@
   (is (= [:b 2] (m/find-in {:a {:b 2}} [:a :b])))
   (is (= [:b {:c 3}] (m/find-in {:a {:b {:c 3}}} [:a :b])))
   (is (= [:c 3] (m/find-in {:a {:b {:c 3}}} [:a :b :c]))))
+
+(deftest test-map-padded
+  (is (= (map + (range 3) (range 4) (range 5) (range 10))
+         [0 4 8]
+         (take 3 (m/map-padded + 0 (range 3) (range 4) (range 5) (range 10)))))
+  (is (= [0 4 8 19 28 35 36 37 38 39]
+         (m/map-padded + 10 (range 3) (range 4) (range 5) (range 10))))
+  (is (= ()
+         (m/map-padded + 10 () () ())))
+  (testing "laziness"
+    (let [state (volatile! [])]
+      (is (= [0 4 8 19 28]
+             (take 5 (m/map-padded (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/map-padded vector :missing [nil nil nil] (range 4) (range 5) (range 10)))))))