Merge pull request #1160 from primo-ppcg/mapcat-et-al

Allow mapcat et al to accept multiple iterable arguments

src/boot/boot.janet

@@ -923,67 +923,68 @@
     (set k (next ind k)))
   ret)
 
+(defmacro- map-aggregator
+  `Aggregation logic for various map functions.`
+  [maptype res val]
+  (case maptype
+    :map ~(array/push ,res ,val)
+    :mapcat ~(array/concat ,res ,val)
+    :keep ~(if (def y ,val) (array/push ,res y))
+    :count ~(if ,val (++ ,res))
+    :some ~(if (def y ,val) (do (set ,res y) (break)))
+    :all ~(if (def y ,val) nil (do (set ,res y) (break)))))
+
+(defmacro- map-n
+  `Generates efficient map logic for a specific number of
+  indexed beyond the first.`
+  [n maptype res f ind inds]
+  ~(do
+    (def ,(seq [k :range [0 n]] (symbol 'ind k)) ,inds)
+    ,;(seq [k :range [0 n]] ~(var ,(symbol 'key k) nil))
+    (each x ,ind
+      ,;(seq [k :range [0 n]]
+        ~(if (= nil (set ,(symbol 'key k) (next ,(symbol 'ind k) ,(symbol 'key k)))) (break)))
+      (map-aggregator ,maptype ,res (,f x ,;(seq [k :range [0 n]] ~(in ,(symbol 'ind k) ,(symbol 'key k))))))))
+
+(defmacro- map-template
+  [maptype res f ind inds]
+  ~(do
+    (def ninds (length ,inds))
+    (case ninds
+      0 (each x ,ind (map-aggregator ,maptype ,res (,f x)))
+      1 (map-n 1 ,maptype ,res ,f ,ind ,inds)
+      2 (map-n 2 ,maptype ,res ,f ,ind ,inds)
+      3 (map-n 3 ,maptype ,res ,f ,ind ,inds)
+      4 (map-n 4 ,maptype ,res ,f ,ind ,inds)
+      (do
+        (def iter-keys (array/new-filled ninds))
+        (def call-buffer (array/new-filled ninds))
+        (var done false)
+        (each x ,ind
+          (forv i 0 ninds
+            (let [old-key (in iter-keys i)
+                  ii (in ,inds i)
+                  new-key (next ii old-key)]
+              (if (= nil new-key)
+                (do (set done true) (break))
+                (do (set (iter-keys i) new-key) (set (call-buffer i) (in ii new-key))))))
+          (if done (break))
+          (map-aggregator ,maptype ,res (,f x ;call-buffer)))))))
+
 (defn map
   `Map a function over every value in a data structure and
   return an array of the results.`
-  [f & inds]
+  [f ind & inds]
-  (def ninds (length inds))
-  (if (= 0 ninds) (error "expected at least 1 indexed collection"))
   (def res @[])
-  (def [i1 i2 i3 i4] inds)
+  (map-template :map res f ind inds)
-  (case ninds
-    1 (each x i1 (array/push res (f x)))
-    2 (do
-        (var k1 nil)
-        (var k2 nil)
-        (while true
-          (if (= nil (set k1 (next i1 k1))) (break))
-          (if (= nil (set k2 (next i2 k2))) (break))
-          (array/push res (f (in i1 k1) (in i2 k2)))))
-    3 (do
-        (var k1 nil)
-        (var k2 nil)
-        (var k3 nil)
-        (while true
-          (if (= nil (set k1 (next i1 k1))) (break))
-          (if (= nil (set k2 (next i2 k2))) (break))
-          (if (= nil (set k3 (next i3 k3))) (break))
-          (array/push res (f (in i1 k1) (in i2 k2) (in i3 k3)))))
-    4 (do
-        (var k1 nil)
-        (var k2 nil)
-        (var k3 nil)
-        (var k4 nil)
-        (while true
-          (if (= nil (set k1 (next i1 k1))) (break))
-          (if (= nil (set k2 (next i2 k2))) (break))
-          (if (= nil (set k3 (next i3 k3))) (break))
-          (if (= nil (set k4 (next i4 k4))) (break))
-          (array/push res (f (in i1 k1) (in i2 k2) (in i3 k3) (in i4 k4)))))
-    (do
-      (def iterkeys (array/new-filled ninds))
-      (var done false)
-      (def call-buffer @[])
-      (while true
-        (forv i 0 ninds
-          (let [old-key (in iterkeys i)
-                ii (in inds i)
-                new-key (next ii old-key)]
-            (if (= nil new-key)
-              (do (set done true) (break))
-              (do (set (iterkeys i) new-key) (array/push call-buffer (in ii new-key))))))
-        (if done (break))
-        (array/push res (f ;call-buffer))
-        (array/clear call-buffer))))
   res)
 
 (defn mapcat
   ``Map a function over every element in an array or tuple and
   use `array/concat` to concatenate the results.``
-  [f ind]
+  [f ind & inds]
   (def res @[])
-  (each x ind
+  (map-template :mapcat res f ind inds)
-    (array/concat res (f x)))
   res)
 
 (defn filter
@@ -999,23 +1000,19 @@
 (defn count
   ``Count the number of items in `ind` for which `(pred item)`
   is true.``
-  [pred ind]
+  [pred ind & inds]
-  (var counter 0)
+  (var res 0)
-  (each item ind
+  (map-template :count res pred ind inds)
-    (if (pred item)
+  res)
-      (++ counter)))
-  counter)
 
 (defn keep
   ``Given a predicate `pred`, return a new array containing the truthy results
   of applying `pred` to each element in the indexed collection `ind`. This is
   different from `filter` which returns an array of the original elements where
   the predicate is truthy.``
-  [pred ind]
+  [pred ind & inds]
   (def res @[])
-  (each item ind
+  (map-template :keep res pred ind inds)
-    (if-let [y (pred item)]
-      (array/push res y)))
   res)
 
 (defn range
@@ -2090,21 +2087,21 @@
   ret)
 
 (defn all
-  ``Returns true if `(pred item)` returns a truthy value for every item in `xs`.
+  ``Returns true if `(pred item)` is truthy for every item in `ind`.
-  Otherwise, returns the first falsey `(pred item)` result encountered.
+  Otherwise, returns the first falsey result encountered.
-  Returns true if `xs` is empty.``
+  Returns true if `ind` is empty.``
-  [pred xs]
+  [pred ind & inds]
-  (var ret true)
+  (var res true)
-  (loop [x :in xs :while ret] (set ret (pred x)))
+  (map-template :all res pred ind inds)
-  ret)
+  res)
 
 (defn some
-  ``Returns nil if all `xs` are false or nil, otherwise returns the result of the
+  ``Returns nil if `(pred item)` is false or nil for every item in `ind`.
-  first truthy predicate, `(pred x)`.``
+  Otherwise, returns the first truthy result encountered.``
-  [pred xs]
+  [pred ind & inds]
-  (var ret nil)
+  (var res nil)
-  (loop [x :in xs :while (not ret)] (if-let [y (pred x)] (set ret y)))
+  (map-template :some res pred ind inds)
-  ret)
+  res)
 
 (defn deep-not=
   ``Like `not=`, but mutable types (arrays, tables, buffers) are considered

test/suite0001.janet

@@ -323,11 +323,25 @@
 (assert (deep= (map + [1 2 3] [10 20 30] [100 200 300]) @[111 222 333]))
 (assert (deep= (map + [1 2 3] [10 20 30] [100 200 300] [1000 2000 3000]) @[1111 2222 3333]))
 (assert (deep= (map + [1 2 3] [10 20 30] [100 200 300] [1000 2000 3000] [10000 20000 30000]) @[11111 22222 33333]))
+(assert (deep= (map + [1 2 3] [10 20 30] [100 200 300] [1000 2000 3000] [10000 20000 30000] [100000 200000 300000]) @[111111 222222 333333]))
 
 # Mapping uses the shortest sequence
 (assert (deep= (map + [1 2 3 4] [10 20 30]) @[11 22 33]))
 (assert (deep= (map + [1 2 3 4] [10 20 30] [100 200]) @[111 222]))
 (assert (deep= (map + [1 2 3 4] [10 20 30] [100 200] [1000]) @[1111]))
+(assert (deep= (map + [1 2 3 4] [10 20 30] [100 200] [1000] []) @[]))
+
+# Variadic arguments to map-like functions
+(assert (deep= (mapcat tuple [1 2 3 4] [5 6 7 8]) @[1 5 2 6 3 7 4 8]))
+(assert (deep= (keep |(if (> $1 0) (/ $0 $1)) [1 2 3 4 5] [1 2 1 0 1]) @[1 1 3 5]))
+
+(assert (= (count = [1 3 2 4 3 5 4 2 1] [1 2 3 4 5 4 3 2 1]) 4))
+
+(assert (= (some not= (range 5) (range 5)) nil))
+(assert (= (some = [1 2 3 4 5] [5 4 3 2 1]) true))
+
+(assert (= (all = (range 5) (range 5)) true))
+(assert (= (all not= [1 2 3 4 5] [5 4 3 2 1]) false))
 
 (assert (= false (deep-not= [1] [1])) "issue #1149")