Count 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 crate::combinations::checked_binomial;
 
 /// An iterator to iterate through all the `n`-length combinations in an iterator, with replacement.
 ///
@@ -107,3 +108,22 @@ where
     I: Iterator,
     I::Item: Clone,
 {}
+
+/// For a given size `n`, return the count of remaining combinations with replacement or None if it would overflow.
+fn remaining_for(n: usize, first: bool, indices: &[usize]) -> Option<usize> {
+    let count = |n: usize, k: usize| checked_binomial((n + k).saturating_sub(1), k);
+    let k = indices.len();
+    if first {
+        count(n, k)
+    } else {
+        indices
+            .iter()
+            .enumerate()
+            // TODO: Once the MSRV hits 1.37.0, we can sum options instead:
+            // .map(|(i, n0)| count(n - 1 - *n0, k - i))
+            // .sum()
+            .fold(Some(0), |sum, (i, n0)| {
+                sum.and_then(|s| s.checked_add(count(n - 1 - *n0, k - i)?))
+            })
+    }
+}