Array combinations

Cargo.toml

@@ -24,6 +24,7 @@ test = false
 
 [dependencies]
 either = { version = "1.0", default-features = false }
+array-init = "2.0"
 
 [dev-dependencies]
 rand = "0.7"

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);

src/adaptors/mod.rs

@@ -812,6 +812,72 @@ 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);
 
+/// An iterator to iterate through all combinations in an iterator that produces arrays
+/// of a specific size.
+///
+/// See [`.array_combinations()`](crate::Itertools::array_combinations) for more
+/// information.
+#[derive(Debug, Clone)]
+pub struct ArrayCombinations<I: Iterator, const R: usize> {
+    iter: I,
+    buf: Vec<I::Item>,
+    indices: [usize; R],
+}
+
+impl<I: Iterator, const R: usize> ArrayCombinations<I, R> {
+    /// Create a new `ArrayCombinations` from an iterator.
+    pub fn new(iter: I) -> Self {
+        let indices = array_init::array_init(|i| i);
+        let buf = Vec::new();
+
+        Self { iter, buf, indices }
+    }
+}
+
+impl<I: Iterator, const R: usize> Iterator for ArrayCombinations<I, R>
+where
+    I::Item: Clone,
+{
+    type Item = [I::Item; R];
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.buf.is_empty() {
+            // If the buffer is empty, this is the first invocation of next
+            for _ in 0..R {
+                // If the source iter returns None, we won't have enough data
+                // for even 1 complete combination. So we can bail.
+                self.buf.push(self.iter.next()?);
+            }
+        } else if self.indices[0] + R == self.buf.len() {
+            // If the first index is as close to the end as possible
+            // eg: [0, 1, 2, 3, 4, 5]
+            //               ^  ^  ^
+            // then we can try get some more data. If there's no more data left
+            // then we've gone over all combinations of the underlying
+            // and we can bail
+            self.buf.push(self.iter.next()?);
+
+            // Reset the indices
+            for i in 0..R - 1 {
+                self.indices[i] = i;
+            }
+            self.indices[R - 1] += 1;
+        } else {
+            let mut i = R - 2;
+            while i > 0 && self.indices[i] + R == self.buf.len() + i {
+                i -= 1;
+            }
+
+            self.indices[i] += 1;
+            for j in i + 1..R-1 {
+                self.indices[j] = self.indices[j - 1] + 1;
+            }
+        }
+
+        Some(array_init::array_init(|i| self.buf[self.indices[i]].clone()))
+    }
+}
+
 /// An iterator adapter to filter values within a nested `Result::Ok`.
 ///
 /// See [`.filter_ok()`](crate::Itertools::filter_ok) for more information.

src/lib.rs

@@ -105,6 +105,7 @@ pub mod structs {
         WhileSome,
         Coalesce,
         TupleCombinations,
+        ArrayCombinations,
         Positions,
         Update,
     };
@@ -1449,6 +1450,45 @@ pub trait Itertools : Iterator {
         adaptors::tuple_combinations(self)
     }
 
+    /// Return an iterator adaptor that iterates over the combinations of the
+    /// elements from an iterator.
+    ///
+    /// Iterator element can be any array of type `Self::Item`
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// let mut v = Vec::new();
+    /// for [a, b] in (1..5).array_combinations() {
+    ///     v.push((a, b));
+    /// }
+    /// assert_eq!(v, vec![(1, 2), (1, 3), (2, 3), (1, 4), (2, 4), (3, 4)]);
+    ///
+    /// let mut it = (1..5).array_combinations();
+    /// assert_eq!(Some([1, 2, 3]), it.next());
+    /// assert_eq!(Some([1, 2, 4]), it.next());
+    /// assert_eq!(Some([1, 3, 4]), it.next());
+    /// assert_eq!(Some([2, 3, 4]), it.next());
+    /// assert_eq!(None, it.next());
+    ///
+    /// // this requires a type hint
+    /// let it = (1..5).array_combinations::<3>();
+    /// itertools::assert_equal(it, vec![[1, 2, 3], [1, 2, 4], [1, 3, 4], [2, 3, 4]]);
+    ///
+    /// // you can also specify the complete type
+    /// use itertools::ArrayCombinations;
+    /// use std::ops::RangeFrom;
+    ///
+    /// let it: ArrayCombinations<RangeFrom<u32>, 2> = (1..).array_combinations();
+    /// itertools::assert_equal(it.take(6), vec![[1, 2], [1, 3], [2, 3], [1, 4], [2, 4], [3, 4]]);
+    /// ```
+    fn array_combinations<const R: usize>(self) -> ArrayCombinations<Self, R>
+        where Self: Sized,
+              Self::Item: Clone,
+    {
+        ArrayCombinations::new(self)
+    }
+
     /// Return an iterator adaptor that iterates over the `k`-length combinations of
     /// the elements from an iterator.
     ///

tests/quick.rs

@@ -895,6 +895,22 @@ quickcheck! {
     }
 }
 
+quickcheck! {
+    fn equal_combinations_array(it: Iter<i16>) -> bool {
+        let values = it.clone().collect_vec();
+        let mut cmb = it.array_combinations();
+        for j in 1..values.len() {
+            for i in 0..j {
+                let pair = [values[i], values[j]];
+                if pair != cmb.next().unwrap() {
+                    return false;
+                }
+            }
+        }
+        cmb.next() == None
+    }
+}
+
 quickcheck! {
     fn size_pad_tail(it: Iter<i8>, pad: u8) -> bool {
         correct_size_hint(it.clone().pad_using(pad as usize, |_| 0)) &&