Skip to content
44 changes: 11 additions & 33 deletions src/permutations.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -123,12 +118,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;
Expand Down Expand Up @@ -156,8 +146,8 @@ where
PermutationState::StartUnknownLen { .. } |
PermutationState::OngoingUnknownLen { .. } => (0, None), // TODO can we improve this lower bound?
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))
}
Expand Down Expand Up @@ -238,39 +228,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)
}
}
}
Expand Down