diff --git a/src/combinations.rs b/src/combinations.rs index 7e12c3803..2ebdaa29d 100644 --- a/src/combinations.rs +++ b/src/combinations.rs @@ -30,21 +30,32 @@ impl fmt::Debug for Combinations pub fn combinations(iter: I, k: usize) -> Combinations where I: Iterator { + let mut pool = LazyBuffer::new(iter); + pool.prefill(k); + Combinations { indices: (0..k).collect(), - pool: LazyBuffer::from_parts(iter, Vec::new(), k), - first: true + pool, + first: true, } } impl Combinations { - #[inline] /// 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. + #[inline] + pub(crate) 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 input data, avoiding (re)allocations if possible. - pub(crate) fn reset(&mut self, k: usize) { + pub(crate) fn init(&mut self, k: usize) { let old_k = self.indices.len(); let reset_limit: usize; @@ -53,10 +64,10 @@ impl Combinations { self.indices.truncate(k); reset_limit = k; } else { - self.indices.reserve(k - old_k); for i in old_k..k { self.indices.push(i); } + reset_limit = old_k; } @@ -77,7 +88,7 @@ impl Iterator for Combinations type Item = Vec; fn next(&mut self) -> Option { if self.first { - if self.pool.len() < self.indices.len() { + if self.pool.is_done() && self.pool.len() < self.indices.len() { return None } self.first = false; diff --git a/src/lazy_buffer.rs b/src/lazy_buffer.rs index 7bec54f55..945a400d1 100644 --- a/src/lazy_buffer.rs +++ b/src/lazy_buffer.rs @@ -1,5 +1,4 @@ use std::ops::Index; -use std::cmp; #[derive(Debug, Clone)] pub struct LazyBuffer { @@ -13,41 +12,10 @@ where I: Iterator, { pub fn new(it: I) -> LazyBuffer { - Self::from_parts(it, Vec::new(), 0) - } - - pub(crate) fn from_parts(it: I, buf: Vec, prefill: usize) -> LazyBuffer { - let mut result = LazyBuffer { + LazyBuffer { it, done: false, - buffer: buf, - }; - - result.prefill(prefill); - result - } - - /// Attempt to prefill the buffer with the requested number of elements `n` from its source - /// iterator. If the buffer already has `n` or more elements filled, this method has no effect. - pub(crate) fn prefill(&mut self, n: usize) { - let buf_len = self.buffer.len(); - if n > buf_len { - let extra = n - buf_len; - - // Allocate capacity in one go, limiting extra allocation to the source iterator - // length's lower bound. - let cap_alloc = cmp::min(n - buf_len, self.it.size_hint().0); - self.buffer.reserve(cap_alloc); - - // Try to push requested elements. - for _ in 0..extra { - if let Some(elt) = self.it.next() { - self.buffer.push(elt); - } else { - self.done = true; - break; - } - } + buffer: Vec::new(), } } @@ -75,6 +43,22 @@ where } } } + + pub fn prefill(&mut self, k: usize) { + if k > self.buffer.len() { + let extra = k - self.buffer.len(); + self.buffer.reserve(extra); + + for _ in 0..extra { + if let Some(elt) = self.it.next() { + self.buffer.push(elt); + } else { + self.done = true; + break; + } + } + } + } } impl Index for LazyBuffer diff --git a/src/powerset.rs b/src/powerset.rs index 46a4fb346..407c74dfd 100644 --- a/src/powerset.rs +++ b/src/powerset.rs @@ -1,8 +1,7 @@ use std::fmt; use std::usize; -use std::vec::IntoIter; -use super::combinations::{self, Combinations}; +use super::combinations::Combinations; use super::size_hint; use crate::Itertools; use std::fmt::Debug; @@ -51,43 +50,40 @@ impl Iterator for Powerset type Item = Vec; fn next(&mut self) -> Option> { - let result = if self.done { - None + if let Some(elt) = self.combs.next() { + self.pos += 1; + Some(elt) } else { - let result = match self.combs.next() { - Some(elt) => { + if self.done { + None + } else { + self.combs.init(self.combs.k() + 1); + if let Some(elt) = self.combs.next() { self.pos += 1; Some(elt) - }, - None => { - self.combs.reset(self.combs.k() + 1); - self.combs.next() + } else { + // None returned from a new Combinations indicates we're finished. + self.done = true; + None } - }; - self.done = result.is_none(); - result - }; - - result + } + } } fn size_hint(&self) -> (usize, Option) { - // TODO try and improve on this - // Might be screwed here unless we keep hold of the src iterator etc. - (0, None) + // Aggregate bounds for source iterator + let src = self.combs.src(); + let src_total = size_hint::add_scalar(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); - // // Aggregate bounds for 'unwalked' source iterator - // let src_total = size_hint::add_scalar(self.src.size_hint(), self.buf.len()); - // - // // Aggregate bounds for self ( length(powerset(set) == 2 ^ length(set) ) - // let self_total = size_hint::two_raised_pow(src_total); - // - // if self.pos < usize::MAX { - // // Subtract count of walked elements from total - // size_hint::sub_scalar(self_total, self.pos) - // } else { - // // self.pos is saturated, no longer reliable - // (0, self_total.1) - // } + if self.pos < usize::MAX { + // Subtract count of walked elements from total + size_hint::sub_scalar(self_total, self.pos) + } else { + // 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 89ed50f14..ca59dc96d 100644 --- a/src/size_hint.rs +++ b/src/size_hint.rs @@ -76,12 +76,12 @@ pub fn mul_scalar(sh: SizeHint, x: usize) -> SizeHint { /// Raise two correctly by a **SizeHint** exponent. #[inline] -pub fn two_raised_pow(exp: SizeHint) -> SizeHint { +pub fn two_exp(exp: SizeHint) -> SizeHint { let (mut low, mut hi) = exp; - let shl_range = 1_usize.leading_zeros() as usize; // saturating 'raise two by exponent' + let shl_range = 1_usize.leading_zeros() as usize; low = if low <= shl_range { 1 << low } else {