diff --git a/src/combinations.rs b/src/combinations.rs index 1c834a74e..d8b5351ec 100644 --- a/src/combinations.rs +++ b/src/combinations.rs @@ -37,12 +37,9 @@ 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, + pool: LazyBuffer::new(iter), first: true, } } @@ -107,6 +104,7 @@ where type Item = Vec; fn next(&mut self) -> Option { if self.first { + self.pool.prefill(self.k()); if self.k() > self.n() { return None; } diff --git a/tests/test_std.rs b/tests/test_std.rs index 732790ef4..732be7b7d 100644 --- a/tests/test_std.rs +++ b/tests/test_std.rs @@ -1044,7 +1044,7 @@ fn combinations_inexact_size_hints() { let len = binomial(real_n, k); assert_eq!(len, it.clone().count()); - let mut nb_loaded = numbers.by_ref().take(k).count(); // because of `LazyBuffer::prefill(k)` + let mut nb_loaded = 0; let sh = numbers.size_hint(); assert_eq!(binomial(sh.0 + nb_loaded, k), it.size_hint().0); assert_eq!(sh.1.map(|n| binomial(n + nb_loaded, k)), it.size_hint().1); @@ -1053,8 +1053,10 @@ fn combinations_inexact_size_hints() { let elem = it.next(); assert!(elem.is_some()); assert_eq!(len - next_count, it.clone().count()); - // It does not load anything more the very first time (it's prefilled). - if next_count > 1 { + if next_count == 1 { + // The very first time, the lazy buffer is prefilled. + nb_loaded = numbers.by_ref().take(k).count(); + } else { // Then it loads one item each time until exhausted. let nb = numbers.next(); if nb.is_some() {