Repository: rust-itertools/itertools
Pull request #763 (2023-09-27 16:53:41)

`TupleCombinations`: `size_hint` and `count`
This was definitely simpler to implement than it was for `Combinations` (with vector items). The only difficulty is that it's recursive with a macro.

According to my notes, all iterators that should have a `size_hint` specialization now have. I like that.

The `count_ident` macro is moved to be available library-wide.

Philippe-Cholet: @jswrenn Ooops... I missed that `checked_binomial` in `combinations.rs` is behind the `use_alloc` feature. Where should I move it? In `adaptors/mod.rs` I guess as it seems to be the easiest move?
jswrenn: Yes, in adaptors/mod.rs is fine.

Patch:
diff --git a/benches/powerset.rs b/benches/powerset.rs
index f5862a4..018333d 100644
--- a/benches/powerset.rs
+++ b/benches/powerset.rs
@@ -20,6 +20,17 @@ fn powerset_n(c: &mut Criterion, n: usize) {
     });
 }
 
+fn powerset_n_fold(c: &mut Criterion, n: usize) {
+    let id = format!("powerset {} fold", n);
+    c.bench_function(id.as_str(), move |b| {
+        b.iter(|| {
+            for _ in 0..calc_iters(n) {
+                (0..n).powerset().fold(0, |s, elt| s + black_box(elt).len());
+            }
+        })
+    });
+}
+
 fn powerset_0(c: &mut Criterion) {
     powerset_n(c, 0);
 }
@@ -44,6 +55,30 @@ fn powerset_12(c: &mut Criterion) {
     powerset_n(c, 12);
 }
 
+fn powerset_0_fold(c: &mut Criterion) {
+    powerset_n_fold(c, 0);
+}
+
+fn powerset_1_fold(c: &mut Criterion) {
+    powerset_n_fold(c, 1);
+}
+
+fn powerset_2_fold(c: &mut Criterion) {
+    powerset_n_fold(c, 2);
+}
+
+fn powerset_4_fold(c: &mut Criterion) {
+    powerset_n_fold(c, 4);
+}
+
+fn powerset_8_fold(c: &mut Criterion) {
+    powerset_n_fold(c, 8);
+}
+
+fn powerset_12_fold(c: &mut Criterion) {
+    powerset_n_fold(c, 12);
+}
+
 criterion_group!(
     benches,
     powerset_0,
@@ -52,5 +87,11 @@ criterion_group!(
     powerset_4,
     powerset_8,
     powerset_12,
+    powerset_0_fold,
+    powerset_1_fold,
+    powerset_2_fold,
+    powerset_4_fold,
+    powerset_8_fold,
+    powerset_12_fold,
 );
 criterion_main!(benches);
diff --git a/src/adaptors/mod.rs b/src/adaptors/mod.rs
index 4bc50ea..90cf46e 100644
--- a/src/adaptors/mod.rs
+++ b/src/adaptors/mod.rs
@@ -14,7 +14,7 @@ pub use self::map::{map_into, map_ok, MapInto, MapOk};
 #[cfg(feature = "use_alloc")]
 pub use self::multi_product::*;
 
-use crate::size_hint;
+use crate::size_hint::{self, SizeHint};
 use std::fmt;
 use std::iter::{FromIterator, Fuse, FusedIterator};
 use std::marker::PhantomData;
@@ -626,6 +626,14 @@ where
     fn next(&mut self) -> Option<Self::Item> {
         self.iter.next()
     }
+
+    fn size_hint(&self) -> SizeHint {
+        self.iter.size_hint()
+    }
+
+    fn count(self) -> usize {
+        self.iter.count()
+    }
 }
 
 impl<I, T> FusedIterator for TupleCombinations<I, T>
@@ -652,6 +660,14 @@ impl<I: Iterator> Iterator for Tuple1Combination<I> {
     fn next(&mut self) -> Option<Self::Item> {
         self.iter.next().map(|x| (x,))
     }
+
+    fn size_hint(&self) -> SizeHint {
+        self.iter.size_hint()
+    }
+
+    fn count(self) -> usize {
+        self.iter.count()
+    }
 }
 
 impl<I: Iterator> HasCombination<I> for (I::Item,) {
@@ -701,6 +717,20 @@ macro_rules! impl_tuple_combination {
                     })
                 }
             }
+
+            fn size_hint(&self) -> SizeHint {
+                const K: usize = 1 + count_ident!($($X,)*);
+                let (mut n_min, mut n_max) = self.iter.size_hint();
+                n_min = checked_binomial(n_min, K).unwrap_or(usize::MAX);
+                n_max = n_max.and_then(|n| checked_binomial(n, K));
+                size_hint::add(self.c.size_hint(), (n_min, n_max))
+            }
+
+            fn count(self) -> usize {
+                const K: usize = 1 + count_ident!($($X,)*);
+                let n = self.iter.count();
+                checked_binomial(n, K).unwrap() + self.c.count()
+            }
         }
 
         impl<I, A> HasCombination<I> for (A, $(ignore_ident!($X, A)),*)
@@ -736,6 +766,42 @@ impl_tuple_combination!(Tuple10Combination Tuple9Combination; a b c d e f g h i)
 impl_tuple_combination!(Tuple11Combination Tuple10Combination; a b c d e f g h i j);
 impl_tuple_combination!(Tuple12Combination Tuple11Combination; a b c d e f g h i j k);
 
+// https://en.wikipedia.org/wiki/Binomial_coefficient#In_programming_languages
+pub(crate) fn checked_binomial(mut n: usize, mut k: usize) -> Option<usize> {
+    if n < k {
+        return Some(0);
+    }
+    // `factorial(n) / factorial(n - k) / factorial(k)` but trying to avoid it overflows:
+    k = (n - k).min(k); // symmetry
+    let mut c = 1;
+    for i in 1..=k {
+        c = (c / i)
+            .checked_mul(n)?
+            .checked_add((c % i).checked_mul(n)? / i)?;
+        n -= 1;
+    }
+    Some(c)
+}
+
+#[test]
+fn test_checked_binomial() {
+    // With the first row: [1, 0, 0, ...] and the first column full of 1s, we check
+    // row by row the recurrence relation of binomials (which is an equivalent definition).
+    // For n >= 1 and k >= 1 we have:
+    //   binomial(n, k) == binomial(n - 1, k - 1) + binomial(n - 1, k)
+    const LIMIT: usize = 500;
+    let mut row = vec![Some(0); LIMIT + 1];
+    row[0] = Some(1);
+    for n in 0..=LIMIT {
+        for k in 0..=LIMIT {
+            assert_eq!(row[k], checked_binomial(n, k));
+        }
+        row = std::iter::once(Some(1))
+            .chain((1..=LIMIT).map(|k| row[k - 1]?.checked_add(row[k]?)))
+            .collect();
+    }
+}
+
 /// An iterator adapter to filter values within a nested `Result::Ok`.
 ///
 /// See [`.filter_ok()`](crate::Itertools::filter_ok) for more information.
diff --git a/src/combinations.rs b/src/combinations.rs
index 332bcb3..1c834a7 100644
--- a/src/combinations.rs
+++ b/src/combinations.rs
@@ -4,6 +4,8 @@ use std::iter::FusedIterator;
 use super::lazy_buffer::LazyBuffer;
 use alloc::vec::Vec;
 
+use crate::adaptors::checked_binomial;
+
 /// An iterator to iterate through all the `k`-length combinations in an iterator.
 ///
 /// See [`.combinations()`](crate::Itertools::combinations) for more information.
@@ -160,42 +162,6 @@ where
 {
 }
 
-// https://en.wikipedia.org/wiki/Binomial_coefficient#In_programming_languages
-pub(crate) fn checked_binomial(mut n: usize, mut k: usize) -> Option<usize> {
-    if n < k {
-        return Some(0);
-    }
-    // `factorial(n) / factorial(n - k) / factorial(k)` but trying to avoid it overflows:
-    k = (n - k).min(k); // symmetry
-    let mut c = 1;
-    for i in 1..=k {
-        c = (c / i)
-            .checked_mul(n)?
-            .checked_add((c % i).checked_mul(n)? / i)?;
-        n -= 1;
-    }
-    Some(c)
-}
-
-#[test]
-fn test_checked_binomial() {
-    // With the first row: [1, 0, 0, ...] and the first column full of 1s, we check
-    // row by row the recurrence relation of binomials (which is an equivalent definition).
-    // For n >= 1 and k >= 1 we have:
-    //   binomial(n, k) == binomial(n - 1, k - 1) + binomial(n - 1, k)
-    const LIMIT: usize = 500;
-    let mut row = vec![Some(0); LIMIT + 1];
-    row[0] = Some(1);
-    for n in 0..=LIMIT {
-        for k in 0..=LIMIT {
-            assert_eq!(row[k], checked_binomial(n, k));
-        }
-        row = std::iter::once(Some(1))
-            .chain((1..=LIMIT).map(|k| row[k - 1]?.checked_add(row[k]?)))
-            .collect();
-    }
-}
-
 /// For a given size `n`, return the count of remaining combinations or None if it would overflow.
 fn remaining_for(n: usize, first: bool, indices: &[usize]) -> Option<usize> {
     let k = indices.len();
diff --git a/src/combinations_with_replacement.rs b/src/combinations_with_replacement.rs
index e4f1cda..f1a0bc0 100644
--- a/src/combinations_with_replacement.rs
+++ b/src/combinations_with_replacement.rs
@@ -3,7 +3,7 @@ use std::fmt;
 use std::iter::FusedIterator;
 
 use super::lazy_buffer::LazyBuffer;
-use crate::combinations::checked_binomial;
+use crate::adaptors::checked_binomial;
 
 /// An iterator to iterate through all the `n`-length combinations in an iterator, with replacement.
 ///
diff --git a/src/impl_macros.rs b/src/impl_macros.rs
index 5fd78cb..7c0b0c2 100644
--- a/src/impl_macros.rs
+++ b/src/impl_macros.rs
@@ -27,3 +27,8 @@ macro_rules! clone_fields {
 macro_rules! ignore_ident{
     ($id:ident, $($t:tt)*) => {$($t)*};
 }
+
+macro_rules! count_ident {
+    () => {0};
+    ($i0:ident, $($i:ident,)*) => {1 + count_ident!($($i,)*)};
+}
diff --git a/src/powerset.rs b/src/powerset.rs
index 811cfce..9a7131a 100644
--- a/src/powerset.rs
+++ b/src/powerset.rs
@@ -3,7 +3,8 @@ use std::fmt;
 use std::iter::FusedIterator;
 use std::usize;
 
-use super::combinations::{checked_binomial, combinations, Combinations};
+use super::combinations::{combinations, Combinations};
+use crate::adaptors::checked_binomial;
 use crate::size_hint::{self, SizeHint};
 
 /// An iterator to iterate through the powerset of the elements from an iterator.
@@ -74,6 +75,24 @@ where
         let (n, combs_count) = self.combs.n_and_count();
         combs_count + remaining_for(n, k).unwrap()
     }
+
+    fn fold<B, F>(self, mut init: B, mut f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B,
+    {
+        let mut it = self.combs;
+        if it.k() == 0 {
+            init = it.by_ref().fold(init, &mut f);
+            it.reset(1);
+        }
+        init = it.by_ref().fold(init, &mut f);
+        // n is now known for sure because k >= 1 and all k-combinations have been generated.
+        for k in it.k() + 1..=it.n() {
+            it.reset(k);
+            init = it.by_ref().fold(init, &mut f);
+        }
+        init
+    }
 }
 
 impl<I> FusedIterator for Powerset<I>
diff --git a/src/tuple_impl.rs b/src/tuple_impl.rs
index 1ce6a5d..c01190e 100644
--- a/src/tuple_impl.rs
+++ b/src/tuple_impl.rs
@@ -331,10 +331,6 @@ pub trait TupleCollect: Sized {
     fn left_shift_push(&mut self, item: Self::Item);
 }
 
-macro_rules! count_ident{
-    () => {0};
-    ($i0:ident, $($i:ident,)*) => {1 + count_ident!($($i,)*)};
-}
 macro_rules! rev_for_each_ident{
     ($m:ident, ) => {};
     ($m:ident, $i0:ident, $($i:ident,)*) => {


Test patch:
diff --git a/tests/quick.rs b/tests/quick.rs
index c779154..5729cfa 100644
--- a/tests/quick.rs
+++ b/tests/quick.rs
@@ -901,8 +901,31 @@ quickcheck! {
 }
 
 quickcheck! {
-    fn size_combinations(it: Iter<i16>) -> bool {
-        correct_size_hint(it.tuple_combinations::<(_, _)>())
+    fn size_combinations(a: Iter<i16>) -> bool {
+        let it = a.clone().tuple_combinations::<(_, _)>();
+        correct_size_hint(it.clone()) && it.count() == binomial(a.count(), 2)
+    }
+
+    fn exact_size_combinations_1(a: Vec<u8>) -> bool {
+        let it = a.iter().tuple_combinations::<(_,)>();
+        exact_size_for_this(it.clone()) && it.count() == binomial(a.len(), 1)
+    }
+    fn exact_size_combinations_2(a: Vec<u8>) -> bool {
+        let it = a.iter().tuple_combinations::<(_, _)>();
+        exact_size_for_this(it.clone()) && it.count() == binomial(a.len(), 2)
+    }
+    fn exact_size_combinations_3(mut a: Vec<u8>) -> bool {
+        a.truncate(15);
+        let it = a.iter().tuple_combinations::<(_, _, _)>();
+        exact_size_for_this(it.clone()) && it.count() == binomial(a.len(), 3)
+    }
+}
+
+fn binomial(n: usize, k: usize) -> usize {
+    if k > n {
+        0
+    } else {
+        (n - k + 1..=n).product::<usize>() / (1..=k).product::<usize>()
     }
 }
 
diff --git a/tests/specializations.rs b/tests/specializations.rs
index efc8ba6..7c21da6 100644
--- a/tests/specializations.rs
+++ b/tests/specializations.rs
@@ -72,9 +72,21 @@ where
 }
 
 quickcheck! {
+    fn tuple_combinations(v: Vec<u8>) -> () {
+        let mut v = v;
+        v.truncate(10);
+        test_specializations(&v.iter().tuple_combinations::<(_, _, _)>());
+    }
+
     fn intersperse(v: Vec<u8>) -> () {
         test_specializations(&v.into_iter().intersperse(0));
     }
+
+    fn powerset(a: Vec<u8>) -> () {
+        let mut a = a;
+        a.truncate(6);
+        test_specializations(&a.iter().powerset())
+    }
 }
 
 quickcheck! {