From cb2449faf8b60e3e09c02f190e7d02b00285409e Mon Sep 17 00:00:00 2001
From: Ronno Das <ronnodas@gmail.com>
Date: Tue, 10 Sep 2024 22:43:58 +0200
Subject: [PATCH] refactored to share code

---
 src/array_combinations.rs | 166 -----------------------------------
 src/combinations.rs       | 180 ++++++++++++++++++++++++++------------
 src/lib.rs                |   8 +-
 3 files changed, 126 insertions(+), 228 deletions(-)
 delete mode 100644 src/array_combinations.rs

diff --git a/src/array_combinations.rs b/src/array_combinations.rs
deleted file mode 100644
index bd2a067d3..000000000
--- a/src/array_combinations.rs
+++ /dev/null
@@ -1,166 +0,0 @@
-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`](ArrayCombinations::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,
-{
-}
diff --git a/src/combinations.rs b/src/combinations.rs
index 9e7492e95..0b2806f68 100644
--- a/src/combinations.rs
+++ b/src/combinations.rs
@@ -1,3 +1,5 @@
+use core::array;
+use core::ops::IndexMut;
 use std::fmt;
 use std::iter::FusedIterator;
 
@@ -6,45 +8,112 @@ 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>
+pub trait PoolIndex<T>: IndexMut<usize, Output = usize> {
+    type Item;
+
+    fn len(&self) -> usize;
+    fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> Self::Item
+    where
+        T: Clone;
+    fn as_slice(&self) -> &[usize];
+}
+
+impl<T> PoolIndex<T> for Vec<usize> {
+    type Item = Vec<T>;
+
+    fn len(&self) -> usize {
+        self.len()
+    }
+    fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> Vec<T>
+    where
+        T: Clone,
+    {
+        pool.get_at(self)
+    }
+
+    fn as_slice(&self) -> &[usize] {
+        self
+    }
+}
+
+impl<T, const K: usize> PoolIndex<T> for [usize; K] {
+    type Item = [T; K];
+
+    fn len(&self) -> usize {
+        K
+    }
+
+    fn extract_item<I: Iterator<Item = T>>(&self, pool: &LazyBuffer<I>) -> [T; K]
+    where
+        T: Clone,
+    {
+        pool.get_array(*self)
+    }
+
+    fn as_slice(&self) -> &[usize] {
+        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,34 +133,17 @@ 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,
             pool,
             first,
         } = self;
-        n_and_count(pool, first, &indices)
+        {
+            let n = pool.count();
+            (n, remaining_for(n, first, indices.as_slice()).unwrap())
+        }
     }
 
     /// Initialises the iterator by filling a buffer with elements from the
@@ -112,7 +164,7 @@ 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() {
+        if self.indices.len() == 0 {
             return true; // Done
         }
 
@@ -144,8 +196,9 @@ 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>
+    pub(crate) fn try_nth(&mut self, n: usize) -> Result<Idx::Item, usize>
     where
+        I: Iterator,
         I::Item: Clone,
     {
         let done = if self.first {
@@ -161,16 +214,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()
@@ -182,7 +236,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> {
@@ -191,8 +245,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.as_slice()).unwrap_or(usize::MAX);
+        upp = upp.and_then(|upp| remaining_for(upp, self.first, self.indices.as_slice()));
         (low, upp)
     }
 
@@ -202,21 +256,35 @@ 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>,
 {
 }
 
-/// Return the length of the inner iterator and the count of remaining combinations.
-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())
+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.
diff --git a/src/lib.rs b/src/lib.rs
index 551dbddd6..cd99d77ca 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -97,9 +97,7 @@ pub mod structs {
         TakeWhileRef, TupleCombinations, Update, WhileSome,
     };
     #[cfg(feature = "use_alloc")]
-    pub use crate::array_combinations::ArrayCombinations;
-    #[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;
@@ -180,8 +178,6 @@ pub mod free;
 #[doc(inline)]
 pub use crate::free::*;
 #[cfg(feature = "use_alloc")]
-mod array_combinations;
-#[cfg(feature = "use_alloc")]
 mod combinations;
 #[cfg(feature = "use_alloc")]
 mod combinations_with_replacement;
@@ -1718,7 +1714,7 @@ pub trait Itertools: Iterator {
         Self: Sized + Clone,
         Self::Item: Clone,
     {
-        array_combinations::array_combinations(self)
+        combinations::array_combinations(self)
     }
 
     /// Return an iterator adaptor that iterates over the `k`-length combinations of