size_hint for combinations with replacement

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<usize> {
+            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<I> FusedIterator for CombinationsWithReplacement<I>

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::<usize>() / (1..=k).product::<usize>();
+            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![]]);