From 2f41e1d31441a34733ae4d619118c9641c03d54d Mon Sep 17 00:00:00 2001
From: Will Crozier <willcrozi@gmail.com>
Date: Sat, 26 Sep 2020 15:55:36 +0100
Subject: [PATCH] Add Iterator::size_hint() impl for Powerset

---
 src/combinations.rs |  4 ++++
 src/powerset.rs     | 36 ++++++++++++++++++++++++++++++++----
 src/size_hint.rs    | 14 ++++++++++++++
 3 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/src/combinations.rs b/src/combinations.rs
index 22ad205da..d2cf070a4 100644
--- a/src/combinations.rs
+++ b/src/combinations.rs
@@ -56,6 +56,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 previous value
     /// an attempt is made to prefill the pool so that it holds `k` elements.
diff --git a/src/powerset.rs b/src/powerset.rs
index 6cc317654..a524d484a 100644
--- a/src/powerset.rs
+++ b/src/powerset.rs
@@ -2,6 +2,7 @@ use std::fmt;
 use alloc::vec::Vec;
 
 use super::combinations::Combinations;
+use super::size_hint;
 
 /// An iterator over the powerset of all elements from its source iterator.
 ///
@@ -10,20 +11,21 @@ use super::combinations::Combinations;
 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
 pub struct Powerset<I: Iterator> {
     combs: Combinations<I>,
+    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 cloneable iterator.
@@ -31,7 +33,10 @@ pub fn powerset<I>(src: I) -> Powerset<I>
     where I: Iterator,
           I::Item: Clone,
 {
-    Powerset { combs: Combinations::new(src, 0) }
+    Powerset {
+        combs: Combinations::new(src, 0),
+        pos: 0,
+    }
 }
 
 impl<I> Iterator for Powerset<I>
@@ -43,14 +48,37 @@ impl<I> Iterator for Powerset<I>
 
     fn next(&mut self) -> Option<Vec<I::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()
+
+                let result = self.combs.next();
+                if result.is_some() {
+                    self.pos = self.pos.saturating_add(1);
+                }
+
+                result
             } else {
                 None
             }
         }
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        // Total bounds for source iterator (or return 'empty' hint if we're done).
+        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::two_exp(src_total);
+
+        if self.pos < usize::MAX {
+            // Subtract count of already yielded elements from total.
+            size_hint::sub_scalar(self_total, self.pos)
+        } else {
+            // Fallback: self.pos is saturated and no longer reliable.
+            (0, self_total.1)
+        }
+    }
 }
diff --git a/src/size_hint.rs b/src/size_hint.rs
index 9d9b8b8a1..6151ce921 100644
--- a/src/size_hint.rs
+++ b/src/size_hint.rs
@@ -74,6 +74,20 @@ pub fn mul_scalar(sh: SizeHint, x: usize) -> SizeHint {
     (low, hi)
 }
 
+/// Raise two correctly by a **SizeHint** exponent.
+#[inline]
+pub fn two_exp(exp: SizeHint) -> SizeHint {
+    let exp_low = cmp::min(exp.0, u32::MAX as usize) as u32;
+    let low = 2usize.saturating_pow(exp_low);
+
+    let hi = exp.1.and_then(|exp| {
+        let exp_hi = cmp::min(exp, u32::MAX as usize) as u32;
+        2usize.checked_pow(exp_hi)
+    });
+
+    (low, hi)
+}
+
 /// Return the maximum
 #[inline]
 pub fn max(a: SizeHint, b: SizeHint) -> SizeHint {