From 83c0f046c077a71185042e020a04262f388a5157 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() method impl. for Powerset

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

diff --git a/src/combinations.rs b/src/combinations.rs
index 31fb5e597..e6ba4ac29 100644
--- a/src/combinations.rs
+++ b/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`
diff --git a/src/powerset.rs b/src/powerset.rs
index df42ff514..ef17752b3 100644
--- a/src/powerset.rs
+++ b/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,20 +12,22 @@ 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.
@@ -31,7 +35,10 @@ 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)
+        }
+    }
 }
diff --git a/src/size_hint.rs b/src/size_hint.rs
index 9d9b8b8a1..1168ecaa3 100644
--- a/src/size_hint.rs
+++ b/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 {