Skip to content

Commit

Permalink
Add Iterator::size_hint() impl for Powerset
Browse files Browse the repository at this point in the history
  • Loading branch information
willcrozi committed Sep 26, 2020
1 parent 69357be commit 2f41e1d
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 4 deletions.
4 changes: 4 additions & 0 deletions src/combinations.rs
Expand Up @@ -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.
Expand Down
36 changes: 32 additions & 4 deletions src/powerset.rs
Expand Up @@ -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.
///
Expand All @@ -10,28 +11,32 @@ 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.
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>
Expand All @@ -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)
}
}
}
14 changes: 14 additions & 0 deletions src/size_hint.rs
Expand Up @@ -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 {
Expand Down

0 comments on commit 2f41e1d

Please sign in to comment.