Make zip variadic (#2041)

crates/typst/src/eval/args.rs

@@ -41,6 +41,11 @@ impl Args {
         Self { span, items }
     }
 
+    /// Returns the number of remaining positional arguments.
+    pub fn remaining(&self) -> usize {
+        self.items.iter().filter(|slot| slot.name.is_none()).count()
+    }
+
     /// Push a positional argument.
     pub fn push(&mut self, span: Span, value: Value) {
         self.items.push(Arg {

crates/typst/src/eval/array.rs

@@ -45,6 +45,11 @@ impl Array {
         Self::default()
     }
 
+    /// Creates a new vec, with a known capacity.
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self(EcoVec::with_capacity(capacity))
+    }
+
     /// Return `true` if the length is 0.
     pub fn is_empty(&self) -> bool {
         self.0.len() == 0
@@ -312,14 +317,45 @@ impl Array {
         Array(vec)
     }
 
-    /// Zips the array with another array. If the two arrays are of unequal length, it will only
-    /// zip up until the last element of the smaller array and the remaining elements will be
-    /// ignored. The return value is an array where each element is yet another array of size 2.
-    pub fn zip(&self, other: Array) -> Array {
-        self.iter()
-            .zip(other)
-            .map(|(first, second)| array![first.clone(), second].into_value())
-            .collect()
+    /// The method `array.zip`, depending on the arguments, it automatically
+    /// detects whether it should use the single zip operator, which depends
+    /// on the standard library's implementation and can therefore be faster.
+    /// Or it zips using a manual implementation which allows for zipping more
+    /// than two arrays at once.
+    pub fn zip(&self, args: &mut Args) -> SourceResult<Self> {
+        // Fast path for just two arrays.
+        if args.remaining() <= 1 {
+            return Ok(self
+                .iter()
+                .zip(args.expect::<Array>("others")?)
+                .map(|(first, second)| array![first.clone(), second].into_value())
+                .collect());
+        }
+
+        // If there is more than one array, we use the manual method.
+        let mut out = Self::with_capacity(self.len());
+        let mut iterators = args
+            .all::<Array>()?
+            .into_iter()
+            .map(|i| i.into_iter())
+            .collect::<Vec<_>>();
+
+        for this in self.iter() {
+            let mut row = Self::with_capacity(1 + iterators.len());
+            row.push(this.clone());
+
+            for iterator in &mut iterators {
+                let Some(item) = iterator.next() else {
+                    return Ok(out);
+                };
+
+                row.push(item);
+            }
+
+            out.push(row.into_value());
+        }
+
+        Ok(out)
     }
 
     /// Return a sorted version of this array, optionally by a given key function.

crates/typst/src/eval/methods.rs

@@ -179,7 +179,7 @@ pub fn call(
             }
             "intersperse" => array.intersperse(args.expect("separator")?).into_value(),
             "sorted" => array.sorted(vm, span, args.named("key")?)?.into_value(),
-            "zip" => array.zip(args.expect("other")?).into_value(),
+            "zip" => array.zip(&mut args)?.into_value(),
             "enumerate" => array
                 .enumerate(args.named("start")?.unwrap_or(0))
                 .at(span)?

docs/reference/types.md

@@ -1033,13 +1033,17 @@ for loop.
 - returns: array
 
 ### zip()
-Zips the array with another array. If the two arrays are of unequal length, it
-will only zip up until the last element of the smaller array and the remaining
-elements will be ignored. The return value is an array where each element is yet
-another array of size 2.
+Zips the array with other arrays. If the arrays are of unequal length, it will
+only zip up until the last element of the shortest array and the remaining
+elements will be ignored. The return value is an array where each element is
+yet another array, the size of each of those is the number of zipped arrays.
 
-- other: array (positional, required)
-  The other array which should be zipped with the current one.
+This method is variadic, meaning that you can zip multiple arrays together at
+once: `(1, 2, 3).zip((3, 4, 5), (6, 7, 8))` returning:
+`((1, 3, 6), (2, 4, 7), (3, 5, 8))`.
+
+- others: array (variadic)
+  The other arrays which should be zipped with the current one.
 - returns: array
 
 ### fold()

tests/typ/compiler/array.typ

@@ -244,6 +244,9 @@
 #test((1, 2, 3, 4).zip((5, 6)), ((1, 5), (2, 6)))
 #test(((1, 2), 3).zip((4, 5)), (((1, 2), 4), (3, 5)))
 #test((1, "hi").zip((true, false)), ((1, true), ("hi", false)))
+#test((1, 2, 3).zip((3, 4, 5), (6, 7, 8)), ((1, 3, 6), (2, 4, 7), (3, 5, 8)))
+#test(().zip((), ()), ())
+#test((1,).zip((2,), (3,)), ((1, 2, 3),))
 
 ---
 // Test the `enumerate` method.