diff --git a/benches/tuple_combinations.rs b/benches/tuple_combinations.rs index 4e26b282e..a9255bf83 100644 --- a/benches/tuple_combinations.rs +++ b/benches/tuple_combinations.rs @@ -99,6 +99,51 @@ fn tuple_comb_c4(c: &mut Criterion) { }); } + +fn array_comb_c1(c: &mut Criterion) { + c.bench_function("array comb c1", move |b| { + b.iter(|| { + for [i] in (0..N1).array_combinations() { + black_box(i); + } + }) + }); +} + + +fn array_comb_c2(c: &mut Criterion) { + c.bench_function("array comb c2", move |b| { + b.iter(|| { + for [i, j] in (0..N2).array_combinations() { + black_box(i + j); + } + }) + }); +} + + +fn array_comb_c3(c: &mut Criterion) { + c.bench_function("array comb c3", move |b| { + b.iter(|| { + for [i, j, k] in (0..N3).array_combinations() { + black_box(i + j + k); + } + }) + }); +} + + +fn array_comb_c4(c: &mut Criterion) { + c.bench_function("array comb c4", move |b| { + b.iter(|| { + for [i, j, k, l] in (0..N4).array_combinations() { + black_box(i + j + k + l); + } + }) + }); +} + + criterion_group!( benches, tuple_comb_for1, @@ -109,5 +154,10 @@ criterion_group!( tuple_comb_c2, tuple_comb_c3, tuple_comb_c4, + array_comb_c1, + array_comb_c2, + array_comb_c3, + array_comb_c4, ); + criterion_main!(benches); diff --git a/src/adaptors/mod.rs b/src/adaptors/mod.rs index dfc68978f..eecf4df6d 100644 --- a/src/adaptors/mod.rs +++ b/src/adaptors/mod.rs @@ -690,7 +690,7 @@ pub trait HasCombination: Sized { /// Create a new `TupleCombinations` from a clonable iterator. pub fn tuple_combinations(iter: I) -> TupleCombinations - where I: Iterator + Clone, + where I: Iterator, I::Item: Clone, T: HasCombination, { @@ -812,6 +812,80 @@ impl_tuple_combination!(Tuple10Combination Tuple9Combination; a b c d e f g h i) impl_tuple_combination!(Tuple11Combination Tuple10Combination; a b c d e f g h i j); impl_tuple_combination!(Tuple12Combination Tuple11Combination; a b c d e f g h i j k); +#[derive(Debug, Clone)] +pub struct ArrayCombinations { + slice: Vec, + indicies: [usize; R], +} + +impl ArrayCombinations { + pub fn new(slice: Vec) -> Self { + debug_assert!(slice.len() >= R); + + let mut indicies = [0; R]; + for i in 0..R { + indicies[i] = i; + } + + Self { + slice, + indicies, + } + } +} + +impl Iterator for ArrayCombinations { + type Item = [T; R]; + + fn next(&mut self) -> Option { + if self.indicies[R-1] == 0 { + return None; + } + + // SAFETY: uninitialized data is never read + let output = unsafe { + let mut output: [T; R] = std::mem::uninitialized(); + for i in 0..R { + output[i] = self.slice[self.indicies[i]].clone(); + } + output + }; + + // // The below uses currently unstable rust. Once they are stable + // // we can replace the deprecated uninitialized + + // let mut output = std::mem::MaybeUninit::uninit_array::(); + // for i in 0..R { + // // SAFETY: only writing so no UB + // unsafe { + // *output[i].as_mut_ptr() = self.slice[self.indicies[i]].clone(); + // } + // } + // // SAFETY: initialised above + // let output = unsafe { std::mem::MaybeUninit::array_assume_init(output) }; + + let mut x = R; + for i in (0..R).rev() { + self.indicies[i] += 1; + if self.indicies[i] == self.slice.len() + i - R + 1 { + x = i; + } else { + break + } + } + + if x == 0 { + self.indicies[R-1] = 0; + } else { + for i in x..R { + self.indicies[i] = self.indicies[i-1] + 1; + } + } + + Some(output) + } +} + /// An iterator adapter to filter values within a nested `Result::Ok`. /// /// See [`.filter_ok()`](crate::Itertools::filter_ok) for more information. diff --git a/src/lib.rs b/src/lib.rs index 1a6632e20..2b3d8b50c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -105,6 +105,7 @@ pub mod structs { WhileSome, Coalesce, TupleCombinations, + ArrayCombinations, Positions, Update, }; @@ -1442,13 +1443,20 @@ pub trait Itertools : Iterator { /// itertools::assert_equal(it, vec![(1, 2, 3), (1, 2, 4), (1, 3, 4), (2, 3, 4)]); /// ``` fn tuple_combinations(self) -> TupleCombinations - where Self: Sized + Clone, + where Self: Sized, Self::Item: Clone, T: adaptors::HasCombination, { adaptors::tuple_combinations(self) } + fn array_combinations(self) -> ArrayCombinations + where Self: Sized, + Self::Item: Clone + std::fmt::Debug, + { + ArrayCombinations::new(self.collect()) + } + /// Return an iterator adaptor that iterates over the `k`-length combinations of /// the elements from an iterator. /// diff --git a/tests/quick.rs b/tests/quick.rs index 7769cb432..4d7d6a890 100644 --- a/tests/quick.rs +++ b/tests/quick.rs @@ -895,6 +895,26 @@ quickcheck! { } } +quickcheck! { + fn equal_combinations_array(it: Iter) -> bool { + let values = it.clone().collect_vec(); + if values.len() < 2 { + return true; + } + + let mut cmb = it.array_combinations(); + for i in 0..values.len() { + for j in i+1..values.len() { + let pair = [values[i], values[j]]; + if pair != cmb.next().unwrap() { + return false; + } + } + } + cmb.next() == None + } +} + quickcheck! { fn size_pad_tail(it: Iter, pad: u8) -> bool { correct_size_hint(it.clone().pad_using(pad as usize, |_| 0)) &&