array_combinations using array::from_fn

Author: ronnodas

src/array_combinations.rs

@@ -0,0 +1,166 @@
+use core::iter::FusedIterator;
+use core::{array, fmt};
+
+use crate::combinations::{n_and_count, remaining_for};
+use crate::lazy_buffer::LazyBuffer;
+
+/// An iterator to iterate through all combinations in an iterator of `Clone`-able items that
+/// produces arrays of a specific size.
+///
+/// See [`.array_combinations()`](crate::Itertools::array_combinations) for more
+/// information.
+#[derive(Clone)]
+#[must_use = "this iterator adaptor is not lazy but does nearly nothing unless consumed"]
+pub struct ArrayCombinations<I: Iterator, const K: usize>
+where
+    I::Item: Clone,
+{
+    indices: [usize; K],
+    pool: LazyBuffer<I>,
+    first: bool,
+}
+
+/// 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,
+{
+    let indices = array::from_fn(|i| i);
+    let pool = LazyBuffer::new(iter);
+
+    ArrayCombinations {
+        indices,
+        pool,
+        first: true,
+    }
+}
+
+impl<I: Iterator, const K: usize> ArrayCombinations<I, K>
+where
+    I::Item: Clone,
+{
+    /// Returns the (current) length of the pool from which combination elements are
+    /// selected. This value can change between invocations of [`next`](Combinations::next).
+    #[inline]
+    pub fn n(&self) -> usize {
+        self.pool.len()
+    }
+
+    /// Initialises the iterator by filling a buffer with elements from the
+    /// iterator. Returns true if there are no combinations, false otherwise.
+    fn init(&mut self) -> bool {
+        self.pool.prefill(K);
+        let done = K > self.n();
+        if !done {
+            self.first = false;
+        }
+
+        done
+    }
+
+    /// Increments indices representing the combination to advance to the next
+    /// (in lexicographic order by increasing sequence) combination. For example
+    /// if we have n=4 & k=2 then `[0, 1] -> [0, 2] -> [0, 3] -> [1, 2] -> ...`
+    ///
+    /// Returns true if we've run out of combinations, false otherwise.
+    fn increment_indices(&mut self) -> bool {
+        if K == 0 {
+            return true; // Done
+        }
+
+        // Scan from the end, looking for an index to increment
+        let mut i: usize = K - 1;
+
+        // Check if we need to consume more from the iterator
+        if self.indices[i] == self.pool.len() - 1 {
+            _ = self.pool.get_next(); // may change pool size
+        }
+
+        while self.indices[i] == i + self.pool.len() - K {
+            if i > 0 {
+                i -= 1;
+            } else {
+                // Reached the last combination
+                return true;
+            }
+        }
+
+        // Increment index, and reset the ones to its right
+        self.indices[i] += 1;
+        for j in i + 1..K {
+            self.indices[j] = self.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::Item: Clone,
+    {
+        let done = if self.first {
+            self.init()
+        } else {
+            self.increment_indices()
+        };
+        if done {
+            return Err(0);
+        }
+        for i in 0..n {
+            if self.increment_indices() {
+                return Err(i + 1);
+            }
+        }
+        Ok(self.pool.get_array(self.indices))
+    }
+}
+
+impl<I: Iterator, const K: usize> Iterator for ArrayCombinations<I, K>
+where
+    I::Item: Clone,
+{
+    type Item = [I::Item; K];
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let done = if self.first {
+            self.init()
+        } else {
+            self.increment_indices()
+        };
+
+        (!done).then(|| self.pool.get_array(self.indices))
+    }
+
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        self.try_nth(n).ok()
+    }
+
+    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, upp)
+    }
+
+    #[inline]
+    fn count(self) -> usize {
+        n_and_count(self.pool, self.first, &self.indices).1
+    }
+}
+
+impl<I, const K: usize> fmt::Debug for ArrayCombinations<I, K>
+where
+    I: Iterator + fmt::Debug,
+    I::Item: Clone + fmt::Debug,
+{
+    debug_fmt_fields!(ArrayCombinations, indices, pool, first);
+}
+
+impl<I, const K: usize> FusedIterator for ArrayCombinations<I, K>
+where
+    I: FusedIterator,
+    I::Item: Clone,
+{
+}

src/combinations.rs

@@ -91,8 +91,7 @@ impl<I: Iterator> Combinations<I> {
             pool,
             first,
         } = self;
-        let n = pool.count();
-        (n, remaining_for(n, first, &indices).unwrap())
+        n_and_count(pool, first, &indices)
     }
 
     /// Initialises the iterator by filling a buffer with elements from the
@@ -210,8 +209,17 @@ where
 {
 }
 
+pub(crate) fn n_and_count<I: Iterator>(
+    pool: LazyBuffer<I>,
+    first: bool,
+    indices: &[usize],
+) -> (usize, usize) {
+    let n = pool.count();
+    (n, remaining_for(n, first, indices).unwrap())
+}
+
 /// 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> {
+pub(crate) fn remaining_for(n: usize, first: bool, indices: &[usize]) -> Option<usize> {
     let k = indices.len();
     if n < k {
         Some(0)

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

@@ -96,6 +96,7 @@ pub mod structs {
         FilterOk, Interleave, InterleaveShortest, MapInto, MapOk, Positions, Product, PutBack,
         TakeWhileRef, TupleCombinations, Update, WhileSome,
     };
+    pub use crate::array_combinations::ArrayCombinations;
     #[cfg(feature = "use_alloc")]
     pub use crate::combinations::Combinations;
     #[cfg(feature = "use_alloc")]
@@ -177,6 +178,7 @@ pub use crate::either_or_both::EitherOrBoth;
 pub mod free;
 #[doc(inline)]
 pub use crate::free::*;
+mod array_combinations;
 #[cfg(feature = "use_alloc")]
 mod combinations;
 #[cfg(feature = "use_alloc")]
@@ -1674,6 +1676,51 @@ 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`.
+    ///
+    /// # 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]]);
+    /// ```
+    fn array_combinations<const K: usize>(self) -> ArrayCombinations<Self, K>
+    where
+        Self: Sized + Clone,
+        Self::Item: Clone,
+    {
+        array_combinations::array_combinations(self)
+    }
+
     /// Return an iterator adaptor that iterates over the `k`-length combinations of
     /// the elements from an iterator.
     ///