1 files changed, 44 insertions, 9 deletions

diff --git a/src/combinations.rs b/src/combinations.rs
index 8759518..e6ba4ac 100644
--- a/src/combinations.rs
+++ b/src/combinations.rs
@@ -1,6 +1,7 @@
 use std::fmt;
 
 use super::lazy_buffer::LazyBuffer;
+use alloc::vec::Vec;
 
 /// An iterator to iterate through all the `k`-length combinations in an iterator.
 ///
@@ -30,13 +31,8 @@ impl<I> fmt::Debug for Combinations<I>
 pub fn combinations<I>(iter: I, k: usize) -> Combinations<I>
     where I: Iterator
 {
-    let mut pool: LazyBuffer<I> = LazyBuffer::new(iter);
-
-    for _ in 0..k {
-        if !pool.get_next() {
-            break;
-        }
-    }
+    let mut pool = LazyBuffer::new(iter);
+    pool.prefill(k);
 
     Combinations {
         indices: (0..k).collect(),
@@ -45,6 +41,45 @@ pub fn combinations<I>(iter: I, k: usize) -> Combinations<I>
     }
 }
 
+impl<I: Iterator> Combinations<I> {
+    /// Returns the length of a combination produced by this iterator.
+    #[inline]
+    pub fn k(&self) -> usize { self.indices.len() }
+
+    /// Returns the (current) length of the pool from which combination elements are
+    /// selected. This value can change between invocations of [`next`].
+    ///
+    /// [`next`]: #method.next
+    #[inline]
+    pub fn n(&self) -> usize { self.pool.len() }
+
+    /// Returns a reference to the source iterator.
+    #[inline]
+    pub(crate) fn src(&self) -> &I { &self.pool.it }
+
+    /// 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);
+        }
+    }
+}
+
 impl<I> Iterator for Combinations<I>
     where I: Iterator,
           I::Item: Clone
@@ -52,11 +87,11 @@ impl<I> Iterator for Combinations<I>
     type Item = Vec<I::Item>;
     fn next(&mut self) -> Option<Self::Item> {
         if self.first {
-            if self.pool.is_done() {
+            if self.k() > self.n() {
                 return None;
             }
             self.first = false;
-        } else if self.indices.len() == 0 {
+        } else if self.indices.is_empty() {
             return None;
         } else {
             // Scan from the end, looking for an index to increment