use alloc::vec::Vec; use std::fmt; use std::iter::once; use super::lazy_buffer::LazyBuffer; /// An iterator adaptor that iterates through all the `k`-permutations of the /// elements from an iterator. /// /// See [`.permutations()`](crate::Itertools::permutations) for /// more information. #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] pub struct Permutations { vals: LazyBuffer, state: PermutationState, } impl Clone for Permutations where I: Clone + Iterator, I::Item: Clone, { clone_fields!(vals, state); } #[derive(Clone, Debug)] enum PermutationState { StartUnknownLen { k: usize, }, OngoingUnknownLen { k: usize, min_n: usize, }, Complete(CompleteState), Empty, } #[derive(Clone, Debug)] enum CompleteState { Start { n: usize, k: usize, }, Ongoing { indices: Vec, cycles: Vec, } } enum CompleteStateRemaining { Known(usize), Overflow, } impl fmt::Debug for Permutations where I: Iterator + fmt::Debug, I::Item: fmt::Debug, { debug_fmt_fields!(Permutations, vals, state); } pub fn permutations(iter: I, k: usize) -> Permutations { let mut vals = LazyBuffer::new(iter); if k == 0 { // Special case, yields single empty vec; `n` is irrelevant let state = PermutationState::Complete(CompleteState::Start { n: 0, k: 0 }); return Permutations { vals, state }; } let mut enough_vals = true; while vals.len() < k { if !vals.get_next() { enough_vals = false; break; } } let state = if enough_vals { PermutationState::StartUnknownLen { k } } else { PermutationState::Empty }; Permutations { vals, state } } impl Iterator for Permutations where I: Iterator, I::Item: Clone { type Item = Vec; fn next(&mut self) -> Option { self.advance(); let &mut Permutations { ref vals, ref state } = self; match *state { PermutationState::StartUnknownLen { .. } => panic!("unexpected iterator state"), PermutationState::OngoingUnknownLen { k, min_n } => { let latest_idx = min_n - 1; let indices = (0..(k - 1)).chain(once(latest_idx)); Some(indices.map(|i| vals[i].clone()).collect()) } PermutationState::Complete(CompleteState::Start { .. }) => None, PermutationState::Complete(CompleteState::Ongoing { ref indices, ref cycles }) => { let k = cycles.len(); Some(indices[0..k].iter().map(|&i| vals[i].clone()).collect()) }, PermutationState::Empty => None } } fn count(self) -> usize { let Permutations { vals, state } = self; fn from_complete(complete_state: CompleteState) -> usize { match complete_state.remaining() { CompleteStateRemaining::Known(count) => count, CompleteStateRemaining::Overflow => { panic!("Iterator count greater than usize::MAX"); } } } match state { PermutationState::StartUnknownLen { k } => { let n = vals.len() + vals.it.count(); let complete_state = CompleteState::Start { n, k }; from_complete(complete_state) } PermutationState::OngoingUnknownLen { k, min_n } => { let prev_iteration_count = min_n - k + 1; let n = vals.len() + vals.it.count(); let complete_state = CompleteState::Start { n, k }; from_complete(complete_state) - prev_iteration_count }, PermutationState::Complete(state) => from_complete(state), PermutationState::Empty => 0 } } fn size_hint(&self) -> (usize, Option) { match self.state { 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) } PermutationState::Empty => (0, Some(0)) } } } impl Permutations where I: Iterator, I::Item: Clone { fn advance(&mut self) { let &mut Permutations { ref mut vals, ref mut state } = self; *state = match *state { PermutationState::StartUnknownLen { k } => { PermutationState::OngoingUnknownLen { k, min_n: k } } PermutationState::OngoingUnknownLen { k, min_n } => { if vals.get_next() { PermutationState::OngoingUnknownLen { k, min_n: min_n + 1 } } else { let n = min_n; let prev_iteration_count = n - k + 1; let mut complete_state = CompleteState::Start { n, k }; // Advance the complete-state iterator to the correct point for _ in 0..(prev_iteration_count + 1) { complete_state.advance(); } PermutationState::Complete(complete_state) } } PermutationState::Complete(ref mut state) => { state.advance(); return; } PermutationState::Empty => { return; } }; } } impl CompleteState { fn advance(&mut self) { *self = match *self { CompleteState::Start { n, k } => { let indices = (0..n).collect(); let cycles = ((n - k)..n).rev().collect(); CompleteState::Ongoing { cycles, indices } }, CompleteState::Ongoing { ref mut indices, ref mut cycles } => { let n = indices.len(); let k = cycles.len(); for i in (0..k).rev() { if cycles[i] == 0 { cycles[i] = n - i - 1; let to_push = indices.remove(i); indices.push(to_push); } else { let swap_index = n - cycles[i]; indices.swap(i, swap_index); cycles[i] -= 1; return; } } CompleteState::Start { n, k } } } } fn remaining(&self) -> CompleteStateRemaining { use self::CompleteStateRemaining::{Known, Overflow}; match *self { CompleteState::Start { n, k } => { if n < k { return Known(0); } let count: Option = (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; } }; } Known(count) } } } }