array_combinations using array::map

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 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,35 +122,15 @@
         &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,
             pool,
             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 @@
     ///
     /// 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,18 +175,18 @@
         }
 
         // 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
     }
 
     /// 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 @@
                 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 @@
             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 @@
 
     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 @@
     }
 }
 
-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();

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>

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.
     ///

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