Complete Permutations::size_hint

src/permutations.rs

@@ -47,11 +47,6 @@ enum CompleteState {
     }
 }
 
-enum CompleteStateRemaining {
-    Known(usize),
-    Overflow,
-}
-
 impl<I> fmt::Debug for Permutations<I>
     where I: Iterator + fmt::Debug,
           I::Item: fmt::Debug,
@@ -72,14 +67,8 @@ pub fn permutations<I: Iterator>(iter: I, k: usize) -> Permutations<I> {
         };
     }
 
-    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 }
@@ -123,12 +112,7 @@ 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");
-                }
-            }
+            complete_state.remaining().expect("Iterator count greater than usize::MAX")
         }
 
         let Permutations { vals, state } = self;
@@ -153,11 +137,29 @@ where
 
     fn size_hint(&self) -> (usize, Option<usize>) {
         match self.state {
-            PermutationState::StartUnknownLen { .. } |
-            PermutationState::OngoingUnknownLen { .. } => (0, None), // TODO can we improve this lower bound?
+            PermutationState::StartUnknownLen { k } => {
+                let (mut low, mut upp) = self.vals.size_hint();
+                low = CompleteState::Start { n: low, k }.remaining().unwrap_or(usize::MAX);
+                upp = upp.and_then(|n| CompleteState::Start { n, k }.remaining());
+                (low, upp)
+            }
+            PermutationState::OngoingUnknownLen { k, min_n } => {
+                let prev_iteration_count = min_n - k + 1;
+                let (mut low, mut upp) = self.vals.size_hint();
+                low = CompleteState::Start { n: low, k }
+                    .remaining()
+                    .unwrap_or(usize::MAX)
+                    .saturating_sub(prev_iteration_count);
+                upp = upp.and_then(|n| {
+                    CompleteState::Start { n, k }
+                        .remaining()
+                        .map(|count| count.saturating_sub(prev_iteration_count))
+                });
+                (low, upp)
+            }
             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))
         }
@@ -238,39 +240,27 @@ impl CompleteState {
         }
     }
 
-    fn remaining(&self) -> CompleteStateRemaining {
-        use self::CompleteStateRemaining::{Known, Overflow};
-
+    /// Returns the count of remaining permutations, or None if it would overflow.
+    fn remaining(&self) -> Option<usize> {
         match *self {
             CompleteState::Start { n, k } => {
                 if n < k {
-                    return Known(0);
+                    return Some(0);
                 }
-
-                let count: Option<usize> = (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; }
-                    };
+                    count = count.checked_mul(radix)
+                        .and_then(|count| count.checked_add(c))?;
                 }
 
-                Known(count)
+                Some(count)
             }
         }
     }

tests/test_std.rs

@@ -986,6 +986,32 @@ fn permutations_zero() {
     it::assert_equal((0..0).permutations(0), vec![vec![]]);
 }
 
+#[test]
+fn permutations_range_count() {
+    for n in 0..=7 {
+        for k in 0..=7 {
+            let len = if k <= n {
+                (n - k + 1..=n).product()
+            } else {
+                0
+            };
+            let mut it = (0..n).permutations(k);
+            assert_eq!(len, it.clone().count());
+            assert_eq!(len, it.size_hint().0);
+            assert_eq!(Some(len), it.size_hint().1);
+            for count in (0..len).rev() {
+                let elem = it.next();
+                assert!(elem.is_some());
+                assert_eq!(count, it.clone().count());
+                assert_eq!(count, it.size_hint().0);
+                assert_eq!(Some(count), it.size_hint().1);
+            }
+            let should_be_none = it.next();
+            assert!(should_be_none.is_none());
+        }
+    }
+}
+
 #[test]
 fn combinations_with_replacement() {
     // Pool smaller than n