From 8d0273ab6117c694e5da716ad06fd9c4944e2f73 Mon Sep 17 00:00:00 2001
From: Nat Quayle Nelson <natquaylenelson@gmail.com>
Date: Mon, 12 Jul 2021 16:17:38 -0600
Subject: [PATCH] variadic zip. Close #9

---
 src/kiss/Kiss.hx                  |  1 -
 src/kiss/Macros.hx                | 14 +++++----
 src/kiss/Prelude.hx               | 48 ++++++++++++++++++++-----------
 src/test/cases/BasicTestCase.kiss |  3 +-
 4 files changed, 42 insertions(+), 24 deletions(-)

diff --git a/src/kiss/Kiss.hx b/src/kiss/Kiss.hx
index 2340921..e9554d3 100644
--- a/src/kiss/Kiss.hx
+++ b/src/kiss/Kiss.hx
@@ -55,7 +55,6 @@ class Kiss {
                 "print" => Symbol("Prelude.print"),
                 "sort" => Symbol("Prelude.sort"),
                 "groups" => Symbol("Prelude.groups"),
-                "zip" => Symbol("Prelude.zip"),
                 "pairs" => Symbol("Prelude.pairs"), // TODO test pairs
                 "reversed" => Symbol("Prelude.reversed"), // TODO test reversed
                 "memoize" => Symbol("Prelude.memoize"), // TODO test memoize
diff --git a/src/kiss/Macros.hx b/src/kiss/Macros.hx
index 66770f5..cb4228e 100644
--- a/src/kiss/Macros.hx
+++ b/src/kiss/Macros.hx
@@ -124,7 +124,8 @@ class Macros {
             "<=" => "Prelude.lesserEqual",
             "=" => "Prelude.areEqual",
             "max" => "Prelude.max",
-            "min" => "Prelude.min"
+            "min" => "Prelude.min",
+            "zip" => "Prelude.zip"
         ];
 
         macros["apply"] = (wholeExp:ReaderExp, exps:Array<ReaderExp>, k) -> {
@@ -166,10 +167,13 @@ class Macros {
             var min = if (exps.length > 1) exps[0] else b.symbol("0");
             var max = if (exps.length > 1) exps[1] else exps[0];
             var step = if (exps.length > 2) exps[2] else b.symbol("1");
-            b.call(
-                b.symbol("Prelude.range"), [
-                    min, max, step
-                ]);
+            b.callSymbol("Prelude.range", [min, max, step]);
+        };
+
+        macros["zip"] = (wholeExp:ReaderExp, exps:kiss.List<ReaderExp>, k) -> {
+            wholeExp.checkNumArgs(2, null, '(zip [listA] [listB] [moreLists...])');
+            var b = wholeExp.expBuilder();
+            return b.callSymbol("Prelude.zip", [b.list(exps.slice(0, -1)), exps[-1]]);
         };
 
         // Most conditional compilation macros are based on this macro:
diff --git a/src/kiss/Prelude.hx b/src/kiss/Prelude.hx
index 01b71ae..14f5c60 100644
--- a/src/kiss/Prelude.hx
+++ b/src/kiss/Prelude.hx
@@ -178,31 +178,45 @@ class Prelude {
         return fullGroups;
     }
 
-    public static function zip(a:Array<Dynamic>, b:Array<Dynamic>, extraHandling = Drop):kiss.List<kiss.List<Dynamic>> {
-        var max = Math.floor(if (a.length != b.length) {
+    public static function zip(arrays:Array<Array<Dynamic>>, extraHandling = Drop):kiss.List<kiss.List<Dynamic>> {
+        var lengthsAreEqual = true;
+        var lengths = [for (arr in arrays) arr.length];
+        var lengthOperands = [for (length in lengths) Operand.fromDynamic(length)];
+        for (idx in 1...lengths.length) {
+            if (lengths[idx] != lengths[idx - 1]) {
+                lengthsAreEqual = false;
+                break;
+            }
+        }
+        var max = Math.floor(if (!lengthsAreEqual) {
             switch (extraHandling) {
                 case Throw:
-                    throw 'zip was given lists of mis-matched size: $a, $b';
+                    throw 'zip was given lists of mis-matched size: $arrays';
                 case Keep:
-                    Math.max(a.length, b.length);
+                    Operand.toDynamic(Prelude.max(lengthOperands));
                 case Drop:
-                    Math.min(a.length, b.length);
+                    Operand.toDynamic(Prelude.min(lengthOperands));
             }
         } else {
-            a.length;
+            lengths[0];
         });
 
         return [
-            for (idx in 0...max) [
-                if (idx < a.length)
-                    a[idx]
-                else
-                    null,
-                if (idx < b.length)
-                    b[idx]
-                else
-                    null
-            ]
+            for (idx in 0...max) {
+                var zipped:Array<Dynamic> = [];
+
+                for (arr in arrays) {
+                    zipped.push(
+                        if (idx < arr.length) {
+                            arr[idx];
+                        } else {
+                            null;
+                        }
+                    );
+                }
+
+                zipped;
+            }
         ];
     }
 
@@ -213,7 +227,7 @@ class Prelude {
             l1.push(l[-1]);
             l2.unshift(l[0]);
         }
-        return zip(l1, l2);
+        return zip([l1, l2]);
     }
 
     public static function reversed<T>(l:kiss.List<T>):kiss.List<T> {
diff --git a/src/test/cases/BasicTestCase.kiss b/src/test/cases/BasicTestCase.kiss
index 6cfcf01..84a6194 100644
--- a/src/test/cases/BasicTestCase.kiss
+++ b/src/test/cases/BasicTestCase.kiss
@@ -122,7 +122,8 @@
     (try (begin (zip [1 3] [2] Throw) (Assert.fail))
         (catch [error] (Assert.pass)))
     (try (begin (zip [1] [2 4] Throw) (Assert.fail))
-        (catch [error] (Assert.pass))))
+        (catch [error] (Assert.pass)))
+    (Assert.equals (.toString [[1 2 3] [2 null 3]]) (.toString (zip [1 2] [2] [3 3] Keep))))
 
 (defun _testLet []
     (let [a 5