Add Iterator::size_hint() method impl. for Powerset

Author: willcrozi

src/combinations.rs

@@ -53,6 +53,10 @@ impl<I: Iterator> Combinations<I> {
     #[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`

src/powerset.rs

@@ -1,7 +1,9 @@
 use std::fmt;
+use std::usize;
 use alloc::vec::Vec;
 
 use super::combinations::{Combinations, combinations};
+use super::size_hint;
 
 /// An iterator to iterate through the powerset of the elements from an iterator.
 ///
@@ -10,28 +12,33 @@ use super::combinations::{Combinations, combinations};
 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
 pub struct Powerset<I: Iterator> {
     combs: Combinations<I>,
+    // Iterator `position` (equal to count of yielded elements).
+    pos: usize,
 }
 
 impl<I> Clone for Powerset<I>
     where I: Clone + Iterator,
           I::Item: Clone,
 {
-    clone_fields!(combs);
+    clone_fields!(combs, pos);
 }
 
 impl<I> fmt::Debug for Powerset<I>
     where I: Iterator + fmt::Debug,
           I::Item: fmt::Debug,
 {
-    debug_fmt_fields!(Powerset, combs);
+    debug_fmt_fields!(Powerset, combs, pos);
 }
 
 /// Create a new `Powerset` from a clonable iterator.
 pub fn powerset<I>(src: I) -> Powerset<I>
     where I: Iterator,
           I::Item: Clone,
 {
-    Powerset { combs: combinations(src, 0) }
+    Powerset {
+        combs: combinations(src, 0),
+        pos: 0,
+    }
 }
 
 impl<I> Iterator for Powerset<I>
@@ -43,14 +50,34 @@ impl<I> Iterator for Powerset<I>
 
     fn next(&mut self) -> Option<Self::Item> {
         if let Some(elt) = self.combs.next() {
+            self.pos = self.pos.saturating_add(1);
             Some(elt)
         } else if self.combs.k() < self.combs.n()
             || self.combs.k() == 0
         {
             self.combs.reset(self.combs.k() + 1);
-            self.combs.next()
+            self.combs.next().map(|elt| {
+                self.pos = self.pos.saturating_add(1);
+                elt
+            })
         } else {
             None
         }
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        // Total bounds for source iterator.
+        let src_total = size_hint::add_scalar(self.combs.src().size_hint(), self.combs.n());
+
+        // Total bounds for self ( length(powerset(set) == 2 ^ length(set) )
+        let self_total = size_hint::pow_scalar_base(2, src_total);
+
+        if self.pos < usize::MAX {
+            // Subtract count of elements already yielded from total.
+            size_hint::sub_scalar(self_total, self.pos)
+        } else {
+            // Fallback: self.pos is saturated and no longer reliable.
+            (0, self_total.1)
+        }
+    }
 }

src/size_hint.rs

@@ -3,6 +3,7 @@
 
 use std::usize;
 use std::cmp;
+use std::u32;
 
 /// **SizeHint** is the return type of **Iterator::size_hint()**.
 pub type SizeHint = (usize, Option<usize>);
@@ -74,6 +75,20 @@ pub fn mul_scalar(sh: SizeHint, x: usize) -> SizeHint {
     (low, hi)
 }
 
+/// Raise `base` correctly by a **`SizeHint`** exponent.
+#[inline]
+pub fn pow_scalar_base(base: usize, exp: SizeHint) -> SizeHint {
+    let exp_low = cmp::min(exp.0, u32::MAX as usize) as u32;
+    let low = base.saturating_pow(exp_low);
+
+    let hi = exp.1.and_then(|exp| {
+        let exp_hi = cmp::min(exp, u32::MAX as usize) as u32;
+        base.checked_pow(exp_hi)
+    });
+
+    (low, hi)
+}
+
 /// Return the maximum
 #[inline]
 pub fn max(a: SizeHint, b: SizeHint) -> SizeHint {