From 3405f4b0c1723689bc261c9895fbb5c1ee4d9acb Mon Sep 17 00:00:00 2001 From: Philippe Cholet <44676486+Philippe-Cholet@users.noreply.github.com> Date: Sat, 12 Aug 2023 14:57:44 +0200 Subject: [PATCH 1/6] size hint: lazy_buffer / permutations As I intend to improve the `permutations` algorithm, I first wanted to simplify a few things: - `CompleteStateRemaining` basically stands for `Option` without its usuability. A doc comment seems enough. - `prefill` is way clearer to see if there are enough values. - `CompleteState::count` method rather than an one-off function. - `0..=x` instead of `0..(x+1)` ; one more fold. Then I encountered a TODO that I fixed: It makes sense to have a `LazyBuffer::size_hint` method which greatly helps in fixing `Permutations::size_hint` and therefore correctly allocate the memory for codes like `(0..n).permutations(k).map(...).collect_vec()`. I think we need to be able to apply any function to both bounds of a `SizeHint`. Here, it avoid repetition and therefore readability. Then permutations' size hints are tested. --- src/lazy_buffer.rs | 11 +++++ src/permutations.rs | 102 ++++++++++++++++++-------------------------- src/size_hint.rs | 12 ++++++ tests/test_std.rs | 14 ++++++ 4 files changed, 78 insertions(+), 61 deletions(-) diff --git a/src/lazy_buffer.rs b/src/lazy_buffer.rs index ca24062aa..a952449ae 100644 --- a/src/lazy_buffer.rs +++ b/src/lazy_buffer.rs @@ -1,6 +1,8 @@ use std::ops::Index; use alloc::vec::Vec; +use crate::size_hint::{self, SizeHint}; + #[derive(Debug, Clone)] pub struct LazyBuffer { pub it: I, @@ -24,6 +26,15 @@ where self.buffer.len() } + pub fn size_hint(&self) -> SizeHint { + let len = self.len(); + if self.done { + (len, Some(len)) + } else { + size_hint::add_scalar(self.it.size_hint(), len) + } + } + pub fn get_next(&mut self) -> bool { if self.done { return false; diff --git a/src/permutations.rs b/src/permutations.rs index d03b85262..9c29248bf 100644 --- a/src/permutations.rs +++ b/src/permutations.rs @@ -3,6 +3,7 @@ use std::fmt; use std::iter::once; use super::lazy_buffer::LazyBuffer; +use super::size_hint::{self, SizeHint}; /// An iterator adaptor that iterates through all the `k`-permutations of the /// elements from an iterator. @@ -47,11 +48,6 @@ enum CompleteState { } } -enum CompleteStateRemaining { - Known(usize), - Overflow, -} - impl fmt::Debug for Permutations where I: Iterator + fmt::Debug, I::Item: fmt::Debug, @@ -72,14 +68,8 @@ pub fn permutations(iter: I, k: usize) -> Permutations { }; } - let mut enough_vals = true; - - while vals.len() < k { - if !vals.get_next() { - enough_vals = false; - break; - } - } + vals.prefill(k); + let enough_vals = vals.len() == k; let state = if enough_vals { PermutationState::StartUnknownLen { k } @@ -122,42 +112,42 @@ where } fn count(self) -> usize { - fn from_complete(complete_state: CompleteState) -> usize { - match complete_state.remaining() { - CompleteStateRemaining::Known(count) => count, - CompleteStateRemaining::Overflow => { - panic!("Iterator count greater than usize::MAX"); - } - } - } - let Permutations { vals, state } = self; match state { PermutationState::StartUnknownLen { k } => { let n = vals.len() + vals.it.count(); - let complete_state = CompleteState::Start { n, k }; - - from_complete(complete_state) + CompleteState::Start { n, k }.count() } PermutationState::OngoingUnknownLen { k, min_n } => { let prev_iteration_count = min_n - k + 1; let n = vals.len() + vals.it.count(); - let complete_state = CompleteState::Start { n, k }; - - from_complete(complete_state) - prev_iteration_count + CompleteState::Start { n, k }.count() - prev_iteration_count }, - PermutationState::Complete(state) => from_complete(state), + PermutationState::Complete(state) => state.count(), PermutationState::Empty => 0 } } - fn size_hint(&self) -> (usize, Option) { + fn size_hint(&self) -> SizeHint { match self.state { - PermutationState::StartUnknownLen { .. } | - PermutationState::OngoingUnknownLen { .. } => (0, None), // TODO can we improve this lower bound? + // Note: the product for `CompleteState::Start` in `remaining` increases with `n`. + PermutationState::StartUnknownLen { k } => { + size_hint::try_map( + self.vals.size_hint(), + |n| CompleteState::Start { n, k }.remaining(), + ) + } + PermutationState::OngoingUnknownLen { k, min_n } => { + let prev_iteration_count = min_n - k + 1; + size_hint::try_map(self.vals.size_hint(), |n| { + CompleteState::Start { n, k } + .remaining() + .and_then(|count| count.checked_sub(prev_iteration_count)) + }) + } PermutationState::Complete(ref state) => match state.remaining() { - CompleteStateRemaining::Known(count) => (count, Some(count)), - CompleteStateRemaining::Overflow => (::std::usize::MAX, None) + Some(count) => (count, Some(count)), + None => (::std::usize::MAX, None) } PermutationState::Empty => (0, Some(0)) } @@ -185,7 +175,7 @@ where let mut complete_state = CompleteState::Start { n, k }; // Advance the complete-state iterator to the correct point - for _ in 0..(prev_iteration_count + 1) { + for _ in 0..=prev_iteration_count { complete_state.advance(); } @@ -238,40 +228,30 @@ impl CompleteState { } } - fn remaining(&self) -> CompleteStateRemaining { - use self::CompleteStateRemaining::{Known, Overflow}; - + /// The remaining count of elements, if it does not overflow. + fn remaining(&self) -> Option { match *self { CompleteState::Start { n, k } => { if n < k { - return Known(0); + return Some(0); } - - let count: Option = (n - k + 1..n + 1).fold(Some(1), |acc, i| { + (n - k + 1..n + 1).fold(Some(1), |acc, i| { acc.and_then(|acc| acc.checked_mul(i)) - }); - - match count { - Some(count) => Known(count), - None => Overflow - } + }) } CompleteState::Ongoing { ref indices, ref cycles } => { - let mut count: usize = 0; - - for (i, &c) in cycles.iter().enumerate() { - let radix = indices.len() - i; - let next_count = count.checked_mul(radix) - .and_then(|count| count.checked_add(c)); - - count = match next_count { - Some(count) => count, - None => { return Overflow; } - }; - } - - Known(count) + cycles.iter().enumerate().fold(Some(0), |acc, (i, c)| { + acc.and_then(|count| { + let radix = indices.len() - i; + count.checked_mul(radix)?.checked_add(*c) + }) + }) } } } + + /// The remaining count of elements, panics if it overflows. + fn count(&self) -> usize { + self.remaining().expect("Iterator count greater than usize::MAX") + } } diff --git a/src/size_hint.rs b/src/size_hint.rs index 71ea1412b..008a0a289 100644 --- a/src/size_hint.rs +++ b/src/size_hint.rs @@ -117,3 +117,15 @@ pub fn min(a: SizeHint, b: SizeHint) -> SizeHint { }; (lower, upper) } + +/// Try to apply a function `f` on both bounds of a `SizeHint`, failure means overflow. +#[inline] +pub fn try_map(sh: SizeHint, mut f: F) -> SizeHint +where + F: FnMut(usize) -> Option, +{ + let (mut low, mut hi) = sh; + low = f(low).unwrap_or(usize::MAX); + hi = hi.and_then(f); + (low, hi) +} diff --git a/tests/test_std.rs b/tests/test_std.rs index 77207d87e..2f3848a02 100644 --- a/tests/test_std.rs +++ b/tests/test_std.rs @@ -915,6 +915,20 @@ fn permutations_zero() { it::assert_equal((0..0).permutations(0), vec![vec![]]); } +#[test] +fn permutations_range_size_hint() { + for n in 0..6 { + for k in 0..=n { + let len: usize = (n - k + 1..=n).product(); + let mut it = (0..n).permutations(k); + for count in (0..=len).rev() { + assert_eq!(it.size_hint(), (count, Some(count))); + it.next(); + } + } + } +} + #[test] fn combinations_with_replacement() { // Pool smaller than n From 2524a9c37613811a60499f45fcedad24e08e8b76 Mon Sep 17 00:00:00 2001 From: Philippe Cholet <44676486+Philippe-Cholet@users.noreply.github.com> Date: Sun, 13 Aug 2023 09:14:19 +0200 Subject: [PATCH 2/6] size_hint for combinations with replacement --- src/combinations_with_replacement.rs | 26 ++++++++++++++++++++++++++ tests/test_std.rs | 16 +++++++++++++++- 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/src/combinations_with_replacement.rs b/src/combinations_with_replacement.rs index 0fec9671a..9b62f8ce8 100644 --- a/src/combinations_with_replacement.rs +++ b/src/combinations_with_replacement.rs @@ -3,6 +3,7 @@ use std::fmt; use std::iter::FusedIterator; use super::lazy_buffer::LazyBuffer; +use super::size_hint::{self, SizeHint}; /// An iterator to iterate through all the `n`-length combinations in an iterator, with replacement. /// @@ -100,6 +101,31 @@ where None => None, } } + + fn size_hint(&self) -> SizeHint { + fn binomial(n: usize, k: usize) -> Option { + if n < k { + return Some(0); + } + // n! / (n - k)! / k! but trying to avoid it overflows: + let k = (n - k).min(k); + (1..=k).fold(Some(1), |res, i| res.and_then(|x| x.checked_mul(n - i + 1).map(|x| x / i))) + } + let k_perms = |n: usize, k: usize| binomial((n + k).saturating_sub(1), k); + let k = self.indices.len(); + size_hint::try_map(self.pool.size_hint(), |n| { + if self.first { + k_perms(n, k) + } else { + self.indices + .iter() + .enumerate() + .fold(Some(0), |sum, (k0, n0)| { + sum.and_then(|s| s.checked_add(k_perms(n - 1 - *n0, k - k0)?)) + }) + } + }) + } } impl FusedIterator for CombinationsWithReplacement diff --git a/tests/test_std.rs b/tests/test_std.rs index 2f3848a02..500175f9d 100644 --- a/tests/test_std.rs +++ b/tests/test_std.rs @@ -923,7 +923,7 @@ fn permutations_range_size_hint() { let mut it = (0..n).permutations(k); for count in (0..=len).rev() { assert_eq!(it.size_hint(), (count, Some(count))); - it.next(); + assert_eq!(it.next().is_none(), count == 0); } } } @@ -962,6 +962,20 @@ fn combinations_with_replacement() { ); } +#[test] +fn combinations_with_replacement_range_size_hint() { + for n in 0..6 { + for k in 0..=n { + let len = (n..n + k).product::() / (1..=k).product::(); + let mut it = (0..n).combinations_with_replacement(k); + for count in (0..=len).rev() { + assert_eq!(it.size_hint(), (count, Some(count))); + assert_eq!(it.next().is_none(), count == 0); + } + } + } +} + #[test] fn powerset() { it::assert_equal((0..0).powerset(), vec![vec![]]); From 06d50952ab6beb90be5820bd44f930a2e3ef3bc6 Mon Sep 17 00:00:00 2001 From: Philippe Cholet <44676486+Philippe-Cholet@users.noreply.github.com> Date: Sun, 13 Aug 2023 10:58:31 +0200 Subject: [PATCH 3/6] size_hint for combinations --- src/combinations.rs | 26 ++++++++++++++++++++++++++ src/combinations_with_replacement.rs | 9 +-------- src/size_hint.rs | 2 ++ tests/test_std.rs | 14 ++++++++++++++ 4 files changed, 43 insertions(+), 8 deletions(-) diff --git a/src/combinations.rs b/src/combinations.rs index 68a59c5e4..59654a1c0 100644 --- a/src/combinations.rs +++ b/src/combinations.rs @@ -2,6 +2,7 @@ use std::fmt; use std::iter::FusedIterator; use super::lazy_buffer::LazyBuffer; +use super::size_hint::{self, SizeHint}; use alloc::vec::Vec; /// An iterator to iterate through all the `k`-length combinations in an iterator. @@ -120,9 +121,34 @@ impl Iterator for Combinations // Create result vector based on the indices Some(self.indices.iter().map(|i| self.pool[*i].clone()).collect()) } + + fn size_hint(&self) -> SizeHint { + let k = self.k(); + size_hint::try_map(self.pool.size_hint(), |n| { + if self.first { + binomial(n, k) + } else { + self.indices + .iter() + .enumerate() + .fold(Some(0), |sum, (k0, n0)| { + sum.and_then(|s| s.checked_add(binomial(n - 1 - *n0, k - k0)?)) + }) + } + }) + } } impl FusedIterator for Combinations where I: Iterator, I::Item: Clone {} + +pub(crate) fn binomial(n: usize, k: usize) -> Option { + if n < k { + return Some(0); + } + // n! / (n - k)! / k! but trying to avoid it overflows: + let k = (n - k).min(k); + (1..=k).fold(Some(1), |res, i| res.and_then(|x| x.checked_mul(n - i + 1).map(|x| x / i))) +} diff --git a/src/combinations_with_replacement.rs b/src/combinations_with_replacement.rs index 9b62f8ce8..610960ffd 100644 --- a/src/combinations_with_replacement.rs +++ b/src/combinations_with_replacement.rs @@ -2,6 +2,7 @@ use alloc::vec::Vec; use std::fmt; use std::iter::FusedIterator; +use super::combinations::binomial; use super::lazy_buffer::LazyBuffer; use super::size_hint::{self, SizeHint}; @@ -103,14 +104,6 @@ where } fn size_hint(&self) -> SizeHint { - fn binomial(n: usize, k: usize) -> Option { - if n < k { - return Some(0); - } - // n! / (n - k)! / k! but trying to avoid it overflows: - let k = (n - k).min(k); - (1..=k).fold(Some(1), |res, i| res.and_then(|x| x.checked_mul(n - i + 1).map(|x| x / i))) - } let k_perms = |n: usize, k: usize| binomial((n + k).saturating_sub(1), k); let k = self.indices.len(); size_hint::try_map(self.pool.size_hint(), |n| { diff --git a/src/size_hint.rs b/src/size_hint.rs index 008a0a289..920964cac 100644 --- a/src/size_hint.rs +++ b/src/size_hint.rs @@ -119,6 +119,8 @@ pub fn min(a: SizeHint, b: SizeHint) -> SizeHint { } /// Try to apply a function `f` on both bounds of a `SizeHint`, failure means overflow. +/// +/// For the resulting size hint to be correct, `f` must be increasing. #[inline] pub fn try_map(sh: SizeHint, mut f: F) -> SizeHint where diff --git a/tests/test_std.rs b/tests/test_std.rs index 500175f9d..f76ee45b6 100644 --- a/tests/test_std.rs +++ b/tests/test_std.rs @@ -909,6 +909,20 @@ fn combinations_zero() { it::assert_equal((0..0).combinations(0), vec![vec![]]); } +#[test] +fn combinations_range_size_hint() { + for n in 0..6 { + for k in 0..=n { + let len = (n - k + 1..=n).product::() / (1..=k).product::(); + let mut it = (0..n).combinations(k); + for count in (0..=len).rev() { + assert_eq!(it.size_hint(), (count, Some(count))); + assert_eq!(it.next().is_none(), count == 0); + } + } + } +} + #[test] fn permutations_zero() { it::assert_equal((1..3).permutations(0), vec![vec![]]); From 5750bc35199e5f99acb208bcea6c71f6a0c2ab9a Mon Sep 17 00:00:00 2001 From: Philippe Cholet <44676486+Philippe-Cholet@users.noreply.github.com> Date: Sun, 13 Aug 2023 11:44:23 +0200 Subject: [PATCH 4/6] binomial: better avoid overflow --- src/combinations.rs | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/combinations.rs b/src/combinations.rs index 59654a1c0..911375abb 100644 --- a/src/combinations.rs +++ b/src/combinations.rs @@ -144,11 +144,16 @@ impl FusedIterator for Combinations I::Item: Clone {} -pub(crate) fn binomial(n: usize, k: usize) -> Option { +pub(crate) fn binomial(mut n: usize, mut k: usize) -> Option { if n < k { return Some(0); } // n! / (n - k)! / k! but trying to avoid it overflows: - let k = (n - k).min(k); - (1..=k).fold(Some(1), |res, i| res.and_then(|x| x.checked_mul(n - i + 1).map(|x| x / i))) + k = (n - k).min(k); + let mut c = 1; + for i in 1..=k { + c = (c / i).checked_mul(n)? + c % i * n / i; + n -= 1; + } + Some(c) } From a3c62c9c1d9e67c69096e19f4bedc1f592c17757 Mon Sep 17 00:00:00 2001 From: Philippe Cholet <44676486+Philippe-Cholet@users.noreply.github.com> Date: Sun, 13 Aug 2023 12:57:30 +0200 Subject: [PATCH 5/6] count: `combinations[_with_replacement]`, `powerset` --- src/combinations.rs | 39 ++++++++++++++++++---------- src/combinations_with_replacement.rs | 36 +++++++++++++++---------- src/powerset.rs | 9 ++++++- tests/test_std.rs | 14 +++++++--- 4 files changed, 67 insertions(+), 31 deletions(-) diff --git a/src/combinations.rs b/src/combinations.rs index 911375abb..dc45a4438 100644 --- a/src/combinations.rs +++ b/src/combinations.rs @@ -53,6 +53,12 @@ impl Combinations { #[inline] pub fn n(&self) -> usize { self.pool.len() } + /// Fill the pool to get its length. + pub(crate) fn real_n(&mut self) -> usize { + while self.pool.get_next() {} + self.pool.len() + } + /// Returns a reference to the source iterator. #[inline] pub(crate) fn src(&self) -> &I { &self.pool.it } @@ -78,6 +84,20 @@ impl Combinations { self.pool.prefill(k); } } + + fn remaining_for(&self, n: usize) -> Option { + let k = self.k(); + if self.first { + binomial(n, k) + } else { + self.indices + .iter() + .enumerate() + .fold(Some(0), |sum, (k0, n0)| { + sum.and_then(|s| s.checked_add(binomial(n - 1 - *n0, k - k0)?)) + }) + } + } } impl Iterator for Combinations @@ -123,19 +143,12 @@ impl Iterator for Combinations } fn size_hint(&self) -> SizeHint { - let k = self.k(); - size_hint::try_map(self.pool.size_hint(), |n| { - if self.first { - binomial(n, k) - } else { - self.indices - .iter() - .enumerate() - .fold(Some(0), |sum, (k0, n0)| { - sum.and_then(|s| s.checked_add(binomial(n - 1 - *n0, k - k0)?)) - }) - } - }) + size_hint::try_map(self.pool.size_hint(), |n| self.remaining_for(n)) + } + + fn count(mut self) -> usize { + let n = self.real_n(); + self.remaining_for(n).expect("Iterator count greater than usize::MAX") } } diff --git a/src/combinations_with_replacement.rs b/src/combinations_with_replacement.rs index 610960ffd..1e5ffb573 100644 --- a/src/combinations_with_replacement.rs +++ b/src/combinations_with_replacement.rs @@ -38,6 +38,21 @@ where fn current(&self) -> Vec { self.indices.iter().map(|i| self.pool[*i].clone()).collect() } + + fn remaining_for(&self, n: usize) -> Option { + let k_perms = |n: usize, k: usize| binomial((n + k).saturating_sub(1), k); + let k = self.indices.len(); + if self.first { + k_perms(n, k) + } else { + self.indices + .iter() + .enumerate() + .fold(Some(0), |sum, (k0, n0)| { + sum.and_then(|s| s.checked_add(k_perms(n - 1 - *n0, k - k0)?)) + }) + } + } } /// Create a new `CombinationsWithReplacement` from a clonable iterator. @@ -104,20 +119,13 @@ where } fn size_hint(&self) -> SizeHint { - let k_perms = |n: usize, k: usize| binomial((n + k).saturating_sub(1), k); - let k = self.indices.len(); - size_hint::try_map(self.pool.size_hint(), |n| { - if self.first { - k_perms(n, k) - } else { - self.indices - .iter() - .enumerate() - .fold(Some(0), |sum, (k0, n0)| { - sum.and_then(|s| s.checked_add(k_perms(n - 1 - *n0, k - k0)?)) - }) - } - }) + size_hint::try_map(self.pool.size_hint(), |n| self.remaining_for(n)) + } + + fn count(mut self) -> usize { + while self.pool.get_next() {} + let n = self.pool.len(); + self.remaining_for(n).expect("Iterator count greater than usize::MAX") } } diff --git a/src/powerset.rs b/src/powerset.rs index 4d7685b12..b43223462 100644 --- a/src/powerset.rs +++ b/src/powerset.rs @@ -3,7 +3,7 @@ use std::iter::FusedIterator; use std::usize; use alloc::vec::Vec; -use super::combinations::{Combinations, combinations}; +use super::combinations::{Combinations, binomial, combinations}; use super::size_hint; /// An iterator to iterate through the powerset of the elements from an iterator. @@ -81,6 +81,13 @@ impl Iterator for Powerset (0, self_total.1) } } + + fn count(mut self) -> usize { + let k = self.combs.k(); + let n = self.combs.real_n(); + // It could be `(1 << n) - self.pos` but `1 << n` might overflow. + self.combs.count() + (k + 1..=n).map(|k| binomial(n, k).unwrap()).sum::() + } } impl FusedIterator for Powerset diff --git a/tests/test_std.rs b/tests/test_std.rs index f76ee45b6..49b137c24 100644 --- a/tests/test_std.rs +++ b/tests/test_std.rs @@ -910,13 +910,14 @@ fn combinations_zero() { } #[test] -fn combinations_range_size_hint() { +fn combinations_range_count() { for n in 0..6 { for k in 0..=n { let len = (n - k + 1..=n).product::() / (1..=k).product::(); let mut it = (0..n).combinations(k); for count in (0..=len).rev() { assert_eq!(it.size_hint(), (count, Some(count))); + assert_eq!(it.clone().count(), count); assert_eq!(it.next().is_none(), count == 0); } } @@ -930,13 +931,14 @@ fn permutations_zero() { } #[test] -fn permutations_range_size_hint() { +fn permutations_range_count() { for n in 0..6 { for k in 0..=n { let len: usize = (n - k + 1..=n).product(); let mut it = (0..n).permutations(k); for count in (0..=len).rev() { assert_eq!(it.size_hint(), (count, Some(count))); + assert_eq!(it.clone().count(), count); assert_eq!(it.next().is_none(), count == 0); } } @@ -977,13 +979,14 @@ fn combinations_with_replacement() { } #[test] -fn combinations_with_replacement_range_size_hint() { +fn combinations_with_replacement_range_count() { for n in 0..6 { for k in 0..=n { let len = (n..n + k).product::() / (1..=k).product::(); let mut it = (0..n).combinations_with_replacement(k); for count in (0..=len).rev() { assert_eq!(it.size_hint(), (count, Some(count))); + assert_eq!(it.clone().count(), count); assert_eq!(it.next().is_none(), count == 0); } } @@ -1005,6 +1008,11 @@ fn powerset() { assert_eq!((0..4).powerset().count(), 1 << 4); assert_eq!((0..8).powerset().count(), 1 << 8); assert_eq!((0..16).powerset().count(), 1 << 16); + let mut it = (0..8).powerset(); + for count in (0..=(1 << 8)).rev() { + assert_eq!(it.clone().count(), count); + assert_eq!(it.next().is_none(), count == 0); + } } #[test] From 99fc19a091387cd3c263f00226077a9fcbcd3b37 Mon Sep 17 00:00:00 2001 From: Philippe-Cholet Date: Mon, 14 Aug 2023 18:18:03 +0200 Subject: [PATCH 6/6] Fuse the iterator in `LazyBuffer` --- src/combinations.rs | 4 ++-- src/lazy_buffer.rs | 23 +++++------------------ 2 files changed, 7 insertions(+), 20 deletions(-) diff --git a/src/combinations.rs b/src/combinations.rs index dc45a4438..57c440c5c 100644 --- a/src/combinations.rs +++ b/src/combinations.rs @@ -1,5 +1,5 @@ use std::fmt; -use std::iter::FusedIterator; +use std::iter::{Fuse, FusedIterator}; use super::lazy_buffer::LazyBuffer; use super::size_hint::{self, SizeHint}; @@ -61,7 +61,7 @@ impl Combinations { /// Returns a reference to the source iterator. #[inline] - pub(crate) fn src(&self) -> &I { &self.pool.it } + pub(crate) fn src(&self) -> &Fuse { &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 diff --git a/src/lazy_buffer.rs b/src/lazy_buffer.rs index a952449ae..88ee06c7c 100644 --- a/src/lazy_buffer.rs +++ b/src/lazy_buffer.rs @@ -1,3 +1,4 @@ +use std::iter::Fuse; use std::ops::Index; use alloc::vec::Vec; @@ -5,8 +6,7 @@ use crate::size_hint::{self, SizeHint}; #[derive(Debug, Clone)] pub struct LazyBuffer { - pub it: I, - done: bool, + pub it: Fuse, buffer: Vec, } @@ -16,8 +16,7 @@ where { pub fn new(it: I) -> LazyBuffer { LazyBuffer { - it, - done: false, + it: it.fuse(), buffer: Vec::new(), } } @@ -27,35 +26,23 @@ where } pub fn size_hint(&self) -> SizeHint { - let len = self.len(); - if self.done { - (len, Some(len)) - } else { - size_hint::add_scalar(self.it.size_hint(), len) - } + size_hint::add_scalar(self.it.size_hint(), self.len()) } pub fn get_next(&mut self) -> bool { - if self.done { - return false; - } if let Some(x) = self.it.next() { self.buffer.push(x); true } else { - self.done = true; false } } pub fn prefill(&mut self, len: usize) { let buffer_len = self.buffer.len(); - - if !self.done && len > buffer_len { + if len > buffer_len { let delta = len - buffer_len; - self.buffer.extend(self.it.by_ref().take(delta)); - self.done = self.buffer.len() < len; } } }