1use super::util::ConstOrExpr;
2use crate::compiler::prelude::*;
3
4fn zip2(value0: Value, value1: Value) -> Resolved {
5 Ok(value0
6 .try_array()?
7 .into_iter()
8 .zip(value1.try_array()?)
9 .map(|(v0, v1)| Value::Array(vec![v0, v1]))
10 .collect())
11}
12
13fn zip_all(value: Value) -> Resolved {
14 Ok(MultiZip(
15 value
16 .try_array()?
17 .into_iter()
18 .map(|value| value.try_array().map(Vec::into_iter))
19 .collect::<Result<_, _>>()?,
20 )
21 .collect::<Vec<_>>()
22 .into())
23}
24
25struct MultiZip(Vec<std::vec::IntoIter<Value>>);
26
27impl Iterator for MultiZip {
28 type Item = Vec<Value>;
29 fn next(&mut self) -> Option<Self::Item> {
30 self.0.iter_mut().map(Iterator::next).collect()
31 }
32}
33
34#[derive(Clone, Copy, Debug)]
35pub struct Zip;
36
37impl Function for Zip {
38 fn identifier(&self) -> &'static str {
39 "zip"
40 }
41
42 fn usage(&self) -> &'static str {
43 indoc! {"
44 Iterate over several arrays in parallel, producing a new array containing arrays of items from each source.
45 The resulting array will be as long as the shortest input array, with all the remaining elements dropped.
46 This function is modeled from the `zip` function [in Python](https://docs.python.org/3/library/functions.html#zip),
47 but similar methods can be found in [Ruby](https://docs.ruby-lang.org/en/master/Array.html#method-i-zip)
48 and [Rust](https://doc.rust-lang.org/stable/std/iter/trait.Iterator.html#method.zip).
49
50 If a single parameter is given, it must contain an array of all the input arrays.
51 "}
52 }
53
54 fn category(&self) -> &'static str {
55 Category::Array.as_ref()
56 }
57
58 fn internal_failure_reasons(&self) -> &'static [&'static str] {
59 &["`array_0` and `array_1` must be arrays."]
60 }
61
62 fn return_kind(&self) -> u16 {
63 kind::ARRAY
64 }
65
66 fn return_rules(&self) -> &'static [&'static str] {
67 &[
68 "`zip` is considered fallible if any of the parameters is not an array, or if only the first parameter is present and it is not an array of arrays.",
69 ]
70 }
71
72 fn parameters(&self) -> &'static [Parameter] {
73 const PARAMETERS: &[Parameter] = &[
74 Parameter::required(
75 "array_0",
76 kind::ARRAY,
77 "The first array of elements, or the array of input arrays if no other parameter is present.",
78 ),
79 Parameter::optional(
80 "array_1",
81 kind::ARRAY,
82 "The second array of elements. If not present, the first parameter contains all the arrays.",
83 ),
84 ];
85 PARAMETERS
86 }
87
88 fn examples(&self) -> &'static [Example] {
89 &[
90 example! {
91 title: "Merge two arrays",
92 source: "zip([1, 2, 3], [4, 5, 6, 7])",
93 result: Ok("[[1, 4], [2, 5], [3, 6]]"),
94 },
95 example! {
96 title: "Merge three arrays",
97 source: r"zip([[1, 2], [3, 4], [5, 6]])",
98 result: Ok(r"[[1, 3, 5], [2, 4, 6]]"),
99 },
100 example! {
101 title: "Merge an array of three arrays into an array of 3-tuples",
102 source: r#"zip([["a", "b", "c"], [1, null, true], [4, 5, 6]])"#,
103 result: Ok(r#"[["a", 1, 4], ["b", null, 5], ["c", true, 6]]"#),
104 },
105 example! {
106 title: "Merge two array parameters",
107 source: "zip([1, 2, 3, 4], [5, 6, 7])",
108 result: Ok("[[1, 5], [2, 6], [3, 7]]"),
109 },
110 ]
111 }
112
113 fn compile(
114 &self,
115 state: &TypeState,
116 _ctx: &mut FunctionCompileContext,
117 arguments: ArgumentList,
118 ) -> Compiled {
119 let array_0 = ConstOrExpr::new(arguments.required("array_0"), state);
120 let array_1 = ConstOrExpr::optional(arguments.optional("array_1"), state);
121
122 Ok(ZipFn { array_0, array_1 }.as_expr())
123 }
124}
125
126#[derive(Clone, Debug)]
127struct ZipFn {
128 array_0: ConstOrExpr,
129 array_1: Option<ConstOrExpr>,
130}
131
132impl FunctionExpression for ZipFn {
133 fn resolve(&self, ctx: &mut Context) -> Resolved {
134 let array_0 = self.array_0.resolve(ctx)?;
135 match &self.array_1 {
136 None => zip_all(array_0),
137 Some(array_1) => zip2(array_0, array_1.resolve(ctx)?),
138 }
139 }
140
141 fn type_def(&self, _state: &TypeState) -> TypeDef {
142 TypeDef::array(Collection::any())
143 }
144}
145
146#[cfg(test)]
147mod tests {
148 use crate::value;
149
150 use super::*;
151
152 test_function![
153 zip => Zip;
154
155 zips_two_arrays {
156 args: func_args![array_0: value!([[1, 2, 3], [4, 5, 6]])],
157 want: Ok(value!([[1, 4], [2, 5], [3, 6]])),
158 tdef: TypeDef::array(Collection::any()),
159 }
160
161 zips_three_arrays {
162 args: func_args![array_0: value!([[1, 2, 3], [4, 5, 6], [7, 8, 9]])],
163 want: Ok(value!([[1, 4, 7], [2, 5, 8], [3, 6, 9]])),
164 tdef: TypeDef::array(Collection::any()),
165 }
166
167 zips_two_parameters {
168 args: func_args![array_0: value!([1, 2, 3]), array_1: value!([4, 5, 6])],
169 want: Ok(value!([[1, 4], [2, 5], [3, 6]])),
170 tdef: TypeDef::array(Collection::any()),
171 }
172
173 uses_shortest_length1 {
174 args: func_args![array_0: value!([[1, 2, 3], [4, 5]])],
175 want: Ok(value!([[1, 4], [2, 5]])),
176 tdef: TypeDef::array(Collection::any()),
177 }
178
179 uses_shortest_length2 {
180 args: func_args![array_0: value!([[1, 2], [4, 5, 6]])],
181 want: Ok(value!([[1, 4], [2, 5]])),
182 tdef: TypeDef::array(Collection::any()),
183 }
184
185 requires_outer_array {
186 args: func_args![array_0: 1],
187 want: Err("expected array, got integer"),
188 tdef: TypeDef::array(Collection::any()),
189 }
190
191 requires_inner_arrays1 {
192 args: func_args![array_0: value!([true, []])],
193 want: Err("expected array, got boolean"),
194 tdef: TypeDef::array(Collection::any()),
195 }
196
197 requires_inner_arrays2 {
198 args: func_args![array_0: value!([[], null])],
199 want: Err("expected array, got null"),
200 tdef: TypeDef::array(Collection::any()),
201 }
202 ];
203}