From c4c53ab52ead548427d66963585b9fa596e85308 Mon Sep 17 00:00:00 2001
From: T0mstone <39707032+T0mstone@users.noreply.github.com>
Date: Mon, 6 May 2024 16:02:27 +0200
Subject: [PATCH] Add `exact` argument to `array.zip` (#4030)

---
 crates/typst/src/foundations/array.rs | 52 +++++++++++++++++++++------
 tests/suite/foundations/array.typ     | 10 ++++++
 2 files changed, 52 insertions(+), 10 deletions(-)

diff --git a/crates/typst/src/foundations/array.rs b/crates/typst/src/foundations/array.rs
index 6bde7d6c7..123dc1f5f 100644
--- a/crates/typst/src/foundations/array.rs
+++ b/crates/typst/src/foundations/array.rs
@@ -8,14 +8,14 @@ use ecow::{eco_format, EcoString, EcoVec};
 use serde::{Deserialize, Serialize};
 use smallvec::SmallVec;
 
-use crate::diag::{bail, At, SourceResult, StrResult};
+use crate::diag::{bail, At, SourceDiagnostic, SourceResult, StrResult};
 use crate::engine::Engine;
 use crate::eval::ops;
 use crate::foundations::{
     cast, func, repr, scope, ty, Args, Bytes, CastInfo, Context, Dict, FromValue, Func,
     IntoValue, Reflect, Repr, Str, Value, Version,
 };
-use crate::syntax::Span;
+use crate::syntax::{Span, Spanned};
 
 /// Create a new [`Array`] from values.
 #[macro_export]
@@ -482,9 +482,14 @@ impl Array {
     #[func]
     pub fn zip(
         self,
-        /// The real arguments (the other arguments are just for the docs, this
-        /// function is a bit involved, so we parse the arguments manually).
+        /// The real arguments (the `others` arguments are just for the docs, this
+        /// function is a bit involved, so we parse the positional arguments manually).
         args: &mut Args,
+        /// Whether all arrays have to have the same length.
+        /// For example, `(1, 2).zip((1, 2, 3), exact: true)` produces an error.
+        #[named]
+        #[default(false)]
+        exact: bool,
         /// The arrays to zip with.
         #[external]
         #[variadic]
@@ -499,7 +504,16 @@ impl Array {
 
         // Fast path for just two arrays.
         if remaining == 1 {
-            let other = args.expect::<Array>("others")?;
+            let Spanned { v: other, span: other_span } =
+                args.expect::<Spanned<Array>>("others")?;
+            if exact && self.len() != other.len() {
+                bail!(
+                    other_span,
+                    "second array has different length ({}) from first array ({})",
+                    other.len(),
+                    self.len()
+                );
+            }
             return Ok(self
                 .into_iter()
                 .zip(other)
@@ -509,11 +523,29 @@ impl Array {
 
         // 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<_>>();
+        let arrays = args.all::<Spanned<Array>>()?;
+        if exact {
+            let errs = arrays
+                .iter()
+                .filter(|sp| sp.v.len() != self.len())
+                .map(|Spanned { v, span }| {
+                    SourceDiagnostic::error(
+                        *span,
+                        eco_format!(
+                            "array has different length ({}) from first array ({})",
+                            v.len(),
+                            self.len()
+                        ),
+                    )
+                })
+                .collect::<EcoVec<_>>();
+            if !errs.is_empty() {
+                return Err(errs);
+            }
+        }
+
+        let mut iterators =
+            arrays.into_iter().map(|i| i.v.into_iter()).collect::<Vec<_>>();
 
         for this in self {
             let mut row = Self::with_capacity(1 + iterators.len());
diff --git a/tests/suite/foundations/array.typ b/tests/suite/foundations/array.typ
index 7b4c574ab..24dad1c1f 100644
--- a/tests/suite/foundations/array.typ
+++ b/tests/suite/foundations/array.typ
@@ -350,6 +350,7 @@
 #test((1,).zip(()), ())
 #test((1,).zip((2,)), ((1, 2),))
 #test((1, 2).zip((3, 4)), ((1, 3), (2, 4)))
+#test((1, 2).zip((3, 4), exact: true), ((1, 3), (2, 4)))
 #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)))
@@ -359,6 +360,15 @@
 #test((1, 2, 3).zip(), ((1,), (2,), (3,)))
 #test(array.zip(()), ())
 
+--- array-zip-exact-error ---
+// Error: 13-22 second array has different length (3) from first array (2)
+#(1, 2).zip((1, 2, 3), exact: true)
+
+--- array-zip-exact-multi-error ---
+// Error: 13-22 array has different length (3) from first array (2)
+// Error: 24-36 array has different length (4) from first array (2)
+#(1, 2).zip((1, 2, 3), (1, 2, 3, 4), exact: true)
+
 --- array-enumerate ---
 // Test the `enumerate` method.
 #test(().enumerate(), ())