Repository: rust-itertools/itertools
Pull request #991 (2024-09-10 14:03:36)

array_combinations using array::map
This is an implementation of `array_combinations` that imitates `combinations` except using `[_; K]` (with `K` a const generic) for the indices and the item type instead of `Vec`. Uses `array::from_fn` and `array::map` although the former could be easily avoided.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/rust-itertools/itertools/pull/991?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=rust-itertools) Report
Attention: Patch coverage is `99.14530%` with `1 line` in your changes missing coverage. Please review.
> Project coverage is 94.34%. Comparing base [(`6814180`)](https://app.codecov.io/gh/rust-itertools/itertools/commit/68141809aabec2816d3250820c8d6608f3b2293f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=rust-itertools) to head [(`3e0d131`)](https://app.codecov.io/gh/rust-itertools/itertools/commit/3e0d1316924205422efdc9df70d7968673ae4ce0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=rust-itertools).
> Report is 116 commits behind head on master.

| [Files with missing lines](https://app.codecov.io/gh/rust-itertools/itertools/pull/991?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=rust-itertools) | Patch % | Lines |
|---|---|---|
| [src/array\_combinations.rs](https://app.codecov.io/gh/rust-itertools/itertools/pull/991?src=pr&el=tree&filepath=src%2Farray_combinations.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=rust-itertools#diff-c3JjL2FycmF5X2NvbWJpbmF0aW9ucy5ycw==) | 98.76% | [1 Missing :warning: ](https://app.codecov.io/gh/rust-itertools/itertools/pull/991?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=rust-itertools) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master     #991      +/-   ##
==========================================
- Coverage   94.38%   94.34%   -0.05%     
==========================================
  Files          48       50       +2     
  Lines        6665     6826     +161     
==========================================
+ Hits         6291     6440     +149     
- Misses        374      386      +12     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/rust-itertools/itertools/pull/991?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=rust-itertools).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=rust-itertools).

ronnodas: Note that `tuple_combinations` does not need to allocate but does need to clone the input iterator (and not just its elements) because it is using a different algorithm, in a way recursing on the length of tuple. I do not see a straightforward way to do that with const generics.
phimuemue: Hi there, thanks for this.

As far as I can judge, this is pretty much a copy-paste of combinations.rs, right?

As such, I strongly suggest to try and generalize our existing combinations (that works on `Vec`) so that it can work on arrays, too. The algorithms are non-trivial and we should strive to have one source of truth.

In addition, I hope that we are able to not only generalize combinations to arrays, but also other algorithms.

@Philippe-Cholet You refactored lots of these algorithms. What's your take on this?
ronnodas: Hi, thanks for the suggestion! Yes, it's basically copy-paste from combinations. 

I will refactor to share as much of the code as makes sense.

I also hope that there will be more array versions soon but I think most of the other clear candidates would be easy to implement once `array::try_from_fn` is available.
ronnodas: Thanks for the extremely helpful guidance!
ronnodas: Is there something to do on my end about the semver check? I'm not sure why it's complaining about `ConsTuples`. `Combinations` did change from a `pub struct` to a `pub type` made of a more generic struct but I think with the same public API so I don't know if this counts as a breaking change.
phimuemue: @jswrenn Do you agree we can include this?

Patch:
diff --git a/src/combinations.rs b/src/combinations.rs
index 6bb2f3e..54a0275 100644
--- a/src/combinations.rs
+++ b/src/combinations.rs
@@ -1,3 +1,5 @@
+use core::array;
+use core::borrow::BorrowMut;
 use std::fmt;
 use std::iter::FusedIterator;
 
@@ -6,45 +8,101 @@ use alloc::vec::Vec;
 
 use crate::adaptors::checked_binomial;
 
+/// Iterator for `Vec` valued combinations returned by [`.combinations()`](crate::Itertools::combinations)
+pub type Combinations<I> = CombinationsGeneric<I, Vec<usize>>;
+/// Iterator for const generic combinations returned by [`.array_combinations()`](crate::Itertools::array_combinations)
+pub type ArrayCombinations<I, const K: usize> = CombinationsGeneric<I, [usize; K]>;
+
+/// Create a new `Combinations` from a clonable iterator.
+pub fn combinations<I: Iterator>(iter: I, k: usize) -> Combinations<I>
+where
+    I::Item: Clone,
+{
+    Combinations::new(iter, (0..k).collect())
+}
+
+/// Create a new `ArrayCombinations` from a clonable iterator.
+pub fn array_combinations<I: Iterator, const K: usize>(iter: I) -> ArrayCombinations<I, K>
+where
+    I::Item: Clone,
+{
+    ArrayCombinations::new(iter, array::from_fn(|i| i))
+}
+
 /// An iterator to iterate through all the `k`-length combinations in an iterator.
 ///
-/// See [`.combinations()`](crate::Itertools::combinations) for more information.
+/// See [`.combinations()`](crate::Itertools::combinations) and [`.array_combinations()`](crate::Itertools::array_combinations) for more information.
 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
-pub struct Combinations<I: Iterator> {
-    indices: Vec<usize>,
+pub struct CombinationsGeneric<I: Iterator, Idx> {
+    indices: Idx,
     pool: LazyBuffer<I>,
     first: bool,
 }
 
-impl<I> Clone for Combinations<I>
+/// A type holding indices of elements in a pool or buffer of items from an inner iterator
+/// and used to pick out different combinations in a generic way.
+pub trait PoolIndex<T>: BorrowMut<[usize]> {
+    type Item;
+
+    fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> Self::Item
+    where
+        T: Clone;
+
+    fn len(&self) -> usize {
+        self.borrow().len()
+    }
+}
+
+impl<T> PoolIndex<T> for Vec<usize> {
+    type Item = Vec<T>;
+
+    fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> Vec<T>
+    where
+        T: Clone,
+    {
+        pool.get_at(self)
+    }
+}
+
+impl<T, const K: usize> PoolIndex<T> for [usize; K] {
+    type Item = [T; K];
+
+    fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> [T; K]
+    where
+        T: Clone,
+    {
+        pool.get_array(*self)
+    }
+}
+
+impl<I, Idx> Clone for CombinationsGeneric<I, Idx>
 where
-    I: Clone + Iterator,
+    I: Iterator + Clone,
     I::Item: Clone,
+    Idx: Clone,
 {
     clone_fields!(indices, pool, first);
 }
 
-impl<I> fmt::Debug for Combinations<I>
+impl<I, Idx> fmt::Debug for CombinationsGeneric<I, Idx>
 where
     I: Iterator + fmt::Debug,
     I::Item: fmt::Debug,
+    Idx: fmt::Debug,
 {
     debug_fmt_fields!(Combinations, indices, pool, first);
 }
 
-/// Create a new `Combinations` from a clonable iterator.
-pub fn combinations<I>(iter: I, k: usize) -> Combinations<I>
-where
-    I: Iterator,
-{
-    Combinations {
-        indices: (0..k).collect(),
-        pool: LazyBuffer::new(iter),
-        first: true,
+impl<I: Iterator, Idx: PoolIndex<I::Item>> CombinationsGeneric<I, Idx> {
+    /// Constructor with arguments the inner iterator and the initial state for the indices.
+    fn new(iter: I, indices: Idx) -> Self {
+        Self {
+            indices,
+            pool: LazyBuffer::new(iter),
+            first: true,
+        }
     }
-}
 
-impl<I: Iterator> Combinations<I> {
     /// Returns the length of a combination produced by this iterator.
     #[inline]
     pub fn k(&self) -> usize {
@@ -64,27 +122,7 @@ impl<I: Iterator> Combinations<I> {
         &self.pool
     }
 
-    /// Resets this `Combinations` back to an initial state for combinations of length
-    /// `k` over the same pool data source. If `k` is larger than the current length
-    /// of the data pool an attempt is made to prefill the pool so that it holds `k`
-    /// elements.
-    pub(crate) fn reset(&mut self, k: usize) {
-        self.first = true;
-
-        if k < self.indices.len() {
-            self.indices.truncate(k);
-            for i in 0..k {
-                self.indices[i] = i;
-            }
-        } else {
-            for i in 0..self.indices.len() {
-                self.indices[i] = i;
-            }
-            self.indices.extend(self.indices.len()..k);
-            self.pool.prefill(k);
-        }
-    }
-
+    /// Return the length of the inner iterator and the count of remaining combinations.
     pub(crate) fn n_and_count(self) -> (usize, usize) {
         let Self {
             indices,
@@ -92,7 +130,7 @@ impl<I: Iterator> Combinations<I> {
             first,
         } = self;
         let n = pool.count();
-        (n, remaining_for(n, first, &indices).unwrap())
+        (n, remaining_for(n, first, indices.borrow()).unwrap())
     }
 
     /// Initialises the iterator by filling a buffer with elements from the
@@ -113,19 +151,21 @@ impl<I: Iterator> Combinations<I> {
     ///
     /// Returns true if we've run out of combinations, false otherwise.
     fn increment_indices(&mut self) -> bool {
-        if self.indices.is_empty() {
+        // Borrow once instead of noise each time it's indexed
+        let indices = self.indices.borrow_mut();
+
+        if indices.is_empty() {
             return true; // Done
         }
-
         // Scan from the end, looking for an index to increment
-        let mut i: usize = self.indices.len() - 1;
+        let mut i: usize = indices.len() - 1;
 
         // Check if we need to consume more from the iterator
-        if self.indices[i] == self.pool.len() - 1 {
+        if indices[i] == self.pool.len() - 1 {
             self.pool.get_next(); // may change pool size
         }
 
-        while self.indices[i] == i + self.pool.len() - self.indices.len() {
+        while indices[i] == i + self.pool.len() - indices.len() {
             if i > 0 {
                 i -= 1;
             } else {
@@ -135,11 +175,10 @@ impl<I: Iterator> Combinations<I> {
         }
 
         // Increment index, and reset the ones to its right
-        self.indices[i] += 1;
-        for j in i + 1..self.indices.len() {
-            self.indices[j] = self.indices[j - 1] + 1;
+        indices[i] += 1;
+        for j in i + 1..indices.len() {
+            indices[j] = indices[j - 1] + 1;
         }
-
         // If we've made it this far, we haven't run out of combos
         false
     }
@@ -147,6 +186,7 @@ impl<I: Iterator> Combinations<I> {
     /// Returns the n-th item or the number of successful steps.
     pub(crate) fn try_nth(&mut self, n: usize) -> Result<<Self as Iterator>::Item, usize>
     where
+        I: Iterator,
         I::Item: Clone,
     {
         let done = if self.first {
@@ -162,16 +202,17 @@ impl<I: Iterator> Combinations<I> {
                 return Err(i + 1);
             }
         }
-        Ok(self.pool.get_at(&self.indices))
+        Ok(self.indices.extract_item(&self.pool))
     }
 }
 
-impl<I> Iterator for Combinations<I>
+impl<I, Idx> Iterator for CombinationsGeneric<I, Idx>
 where
     I: Iterator,
     I::Item: Clone,
+    Idx: PoolIndex<I::Item>,
 {
-    type Item = Vec<I::Item>;
+    type Item = Idx::Item;
     fn next(&mut self) -> Option<Self::Item> {
         let done = if self.first {
             self.init()
@@ -183,7 +224,7 @@ where
             return None;
         }
 
-        Some(self.pool.get_at(&self.indices))
+        Some(self.indices.extract_item(&self.pool))
     }
 
     fn nth(&mut self, n: usize) -> Option<Self::Item> {
@@ -192,8 +233,8 @@ where
 
     fn size_hint(&self) -> (usize, Option<usize>) {
         let (mut low, mut upp) = self.pool.size_hint();
-        low = remaining_for(low, self.first, &self.indices).unwrap_or(usize::MAX);
-        upp = upp.and_then(|upp| remaining_for(upp, self.first, &self.indices));
+        low = remaining_for(low, self.first, self.indices.borrow()).unwrap_or(usize::MAX);
+        upp = upp.and_then(|upp| remaining_for(upp, self.first, self.indices.borrow()));
         (low, upp)
     }
 
@@ -203,13 +244,37 @@ where
     }
 }
 
-impl<I> FusedIterator for Combinations<I>
+impl<I, Idx> FusedIterator for CombinationsGeneric<I, Idx>
 where
     I: Iterator,
     I::Item: Clone,
+    Idx: PoolIndex<I::Item>,
 {
 }
 
+impl<I: Iterator> Combinations<I> {
+    /// Resets this `Combinations` back to an initial state for combinations of length
+    /// `k` over the same pool data source. If `k` is larger than the current length
+    /// of the data pool an attempt is made to prefill the pool so that it holds `k`
+    /// elements.
+    pub(crate) fn reset(&mut self, k: usize) {
+        self.first = true;
+
+        if k < self.indices.len() {
+            self.indices.truncate(k);
+            for i in 0..k {
+                self.indices[i] = i;
+            }
+        } else {
+            for i in 0..self.indices.len() {
+                self.indices[i] = i;
+            }
+            self.indices.extend(self.indices.len()..k);
+            self.pool.prefill(k);
+        }
+    }
+}
+
 /// For a given size `n`, return the count of remaining combinations or None if it would overflow.
 fn remaining_for(n: usize, first: bool, indices: &[usize]) -> Option<usize> {
     let k = indices.len();
diff --git a/src/lazy_buffer.rs b/src/lazy_buffer.rs
index fefcff8..fafa5f7 100644
--- a/src/lazy_buffer.rs
+++ b/src/lazy_buffer.rs
@@ -59,6 +59,10 @@ where
     pub fn get_at(&self, indices: &[usize]) -> Vec<I::Item> {
         indices.iter().map(|i| self.buffer[*i].clone()).collect()
     }
+
+    pub fn get_array<const K: usize>(&self, indices: [usize; K]) -> [I::Item; K] {
+        indices.map(|i| self.buffer[i].clone())
+    }
 }
 
 impl<I, J> Index<J> for LazyBuffer<I>
diff --git a/src/lib.rs b/src/lib.rs
index 9670390..dc316d8 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -97,7 +97,7 @@ pub mod structs {
         TakeWhileRef, TupleCombinations, Update, WhileSome,
     };
     #[cfg(feature = "use_alloc")]
-    pub use crate::combinations::Combinations;
+    pub use crate::combinations::{ArrayCombinations, Combinations};
     #[cfg(feature = "use_alloc")]
     pub use crate::combinations_with_replacement::CombinationsWithReplacement;
     pub use crate::cons_tuples_impl::ConsTuples;
@@ -1674,6 +1674,53 @@ 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 type is [Self::Item; K]. The iterator produces a new
+    /// array per iteration, and clones the iterator elements.
+    ///
+    /// # Guarantees
+    ///
+    /// If the adapted iterator is deterministic,
+    /// this iterator adapter yields items in a reliable order.
+    ///
+    /// ```
+    /// 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], [1, 4], [2, 3], [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::Range;
+    ///
+    /// let it: ArrayCombinations<Range<u32>, 3> = (1..5).array_combinations();
+    /// itertools::assert_equal(it, vec![[1, 2, 3], [1, 2, 4], [1, 3, 4], [2, 3, 4]]);
+    /// ```
+    #[cfg(feature = "use_alloc")]
+    fn array_combinations<const K: usize>(self) -> ArrayCombinations<Self, K>
+    where
+        Self: Sized + Clone,
+        Self::Item: Clone,
+    {
+        combinations::array_combinations(self)
+    }
+
     /// Return an iterator adaptor that iterates over the `k`-length combinations of
     /// the elements from an iterator.
     ///


Test patch:
diff --git a/tests/specializations.rs b/tests/specializations.rs
index 3e48310..e6694c8 100644
--- a/tests/specializations.rs
+++ b/tests/specializations.rs
@@ -273,6 +273,16 @@ quickcheck! {
         test_specializations(&v.into_iter().intersperse_with(|| 0));
     }
 
+    fn array_combinations(v: Vec<u8>) -> TestResult {
+        if v.len() > 10 {
+            return TestResult::discard();
+        }
+        test_specializations(&v.iter().array_combinations::<1>());
+        test_specializations(&v.iter().array_combinations::<2>());
+        test_specializations(&v.iter().array_combinations::<3>());
+        TestResult::passed()
+    }
+
     fn combinations(a: Vec<u8>, n: u8) -> TestResult {
         if n > 3 || a.len() > 8 {
             return TestResult::discard();