diff --git a/ndc_stdlib/src/sequence.rs b/ndc_stdlib/src/sequence.rs index 56a3c0cd..2ecc80e0 100644 --- a/ndc_stdlib/src/sequence.rs +++ b/ndc_stdlib/src/sequence.rs @@ -1,6 +1,7 @@ #![allow(clippy::ptr_arg)] use anyhow::anyhow; +use itertools::Itertools; use ndc_macros::export_module; use ndc_vm::VmCallable; use ndc_vm::value::{Object, SeqValue, Value}; @@ -701,26 +702,21 @@ mod inner { } /// Transposes a sequence of sequences, turning rows into columns, and returns the result as a list of lists. - #[function(signature = "(Iterable) -> List>")] + #[function(signature = "(Iterable>) -> List>")] pub fn transposed(seq: SeqValue) -> anyhow::Result { - let main: Vec = seq - .try_into_iter() - .ok_or_else(|| anyhow!("transposed requires a sequence"))? - .collect(); - let mut iterators: Vec>> = Vec::new(); - for elem in main { - iterators.push(Box::new(elem.try_into_iter().ok_or_else(|| { - anyhow!("elements of transposed sequence must be iterable") - })?)); - } - let mut out = Vec::new(); - loop { - let row: Vec = iterators.iter_mut().filter_map(|i| i.next()).collect(); - if row.is_empty() { - return Ok(Value::list(out)); - } - out.push(Value::list(row)); - } + transpose(seq) + } + + /// Transposes a sequence of strings, turning rows of characters into columns. + #[function(name = "transposed", signature = "(Iterable) -> List>")] + pub fn transposed_strings(seq: SeqValue) -> anyhow::Result { + transpose(seq) + } + + /// Transposes a sequence of sequences whose elements have different types. + #[function(name = "transposed", signature = "(Iterable) -> List>")] + pub fn transposed_any(seq: SeqValue) -> anyhow::Result { + transpose(seq) } /// Return a list of all windows, wrapping back to the first elements when the window would otherwise exceed the length of source list, producing tuples of size 2. @@ -840,26 +836,18 @@ mod inner { /// [3,"c",false] /// ] /// ``` - #[function(signature = "(Iterable) -> List>")] + #[function(signature = "(Iterable>) -> List>")] pub fn multi_cartesian_product(seq: SeqValue) -> anyhow::Result { - let mut iterators = Vec::new(); - for elem in seq - .try_into_iter() - .ok_or_else(|| anyhow!("multi_cartesian_product requires a sequence"))? - { - let inner: Vec = elem - .try_into_iter() - .ok_or_else(|| anyhow!("elements of sequence must be iterable"))? - .collect_vec(); - iterators.push(inner.into_iter()); - } - Ok(Value::list( - iterators - .into_iter() - .multi_cartesian_product() - .map(Value::list) - .collect_vec(), - )) + cartesian_product(seq) + } + + /// The cartesian product of sequences whose elements have different types. + #[function( + name = "multi_cartesian_product", + signature = "(Iterable) -> List>" + )] + pub fn multi_cartesian_product_any(seq: SeqValue) -> anyhow::Result { + cartesian_product(seq) } /// Split the input sequence into evenly sized chunks. If the input length of the sequence @@ -896,6 +884,48 @@ mod inner { } } +fn transpose(seq: SeqValue) -> anyhow::Result { + let main: Vec = seq + .try_into_iter() + .ok_or_else(|| anyhow!("transposed requires a sequence"))? + .collect(); + let mut iterators: Vec>> = Vec::new(); + for elem in main { + iterators.push(Box::new(elem.try_into_iter().ok_or_else(|| { + anyhow!("elements of transposed sequence must be iterable") + })?)); + } + let mut out = Vec::new(); + loop { + let row: Vec = iterators.iter_mut().filter_map(|i| i.next()).collect(); + if row.is_empty() { + return Ok(Value::list(out)); + } + out.push(Value::list(row)); + } +} + +fn cartesian_product(seq: SeqValue) -> anyhow::Result { + let mut iterators = Vec::new(); + for elem in seq + .try_into_iter() + .ok_or_else(|| anyhow!("multi_cartesian_product requires a sequence"))? + { + let inner: Vec = elem + .try_into_iter() + .ok_or_else(|| anyhow!("elements of sequence must be iterable"))? + .collect_vec(); + iterators.push(inner.into_iter()); + } + Ok(Value::list( + iterators + .into_iter() + .multi_cartesian_product() + .map(Value::list) + .collect_vec(), + )) +} + fn by_key(seq: Value, func: &mut VmCallable<'_>, better: Ordering) -> anyhow::Result { let mut best_value: Option = None; let mut best_key: Option = None; @@ -965,8 +995,11 @@ pub mod extra { pub fn register(env: &mut FunctionRegistry>) { for signature in [ - // Two sequences zip into pairs of their elements. + // Two sequences zip into pairs of their elements, three into + // triples, and four into quadruples. "(Iterable, Iterable) -> List<(A, B)>", + "(Iterable, Iterable, Iterable) -> List<(A, B, C)>", + "(Iterable, Iterable, Iterable, Iterable) -> List<(A, B, C, D)>", // More make tuples as long as the number of arguments. "(*) -> List", ] { diff --git a/tests/functional/programs/603_stdlib_seq/026_transposed_and_cartesian_products_keep_the_element_type.ndc b/tests/functional/programs/603_stdlib_seq/026_transposed_and_cartesian_products_keep_the_element_type.ndc new file mode 100644 index 00000000..2e2a3661 --- /dev/null +++ b/tests/functional/programs/603_stdlib_seq/026_transposed_and_cartesian_products_keep_the_element_type.ndc @@ -0,0 +1,6 @@ +// `transposed` and `multi_cartesian_product` give lists of the elements the +// rows hold, and sequences of mixed elements give lists of `Any`. +// expect-output: [[1,3],[2,4]] [[1,3],[2,3]] [[1,"a"],[2,"a"]] +let columns: List> = [[1, 2], [3, 4]].transposed(); +let products: List> = [[1, 2], [3]].multi_cartesian_product(); +print(columns, products, [[1, 2], ["a"]].multi_cartesian_product()); diff --git a/tests/functional/programs/603_stdlib_seq/027_zip_of_three_or_four_sequences_makes_typed_tuples.ndc b/tests/functional/programs/603_stdlib_seq/027_zip_of_three_or_four_sequences_makes_typed_tuples.ndc new file mode 100644 index 00000000..8894b392 --- /dev/null +++ b/tests/functional/programs/603_stdlib_seq/027_zip_of_three_or_four_sequences_makes_typed_tuples.ndc @@ -0,0 +1,5 @@ +// Zipping three or four sequences makes tuples of their element types. +// expect-output: [(1,"a",1.5)] [(1,"a",true,2)] +let triples: List<(Int, String, Float)> = zip([1], ["a"], [1.5]); +let quadruples: List<(Int, String, Bool, Int)> = zip([1], ["a"], [true], [2]); +print(triples, quadruples);