diff --git a/src/lib.rs b/src/lib.rs index 018f877e5..59b838ffe 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -174,6 +174,8 @@ pub use crate::either_or_both::EitherOrBoth; pub mod free; #[doc(inline)] pub use crate::free::*; +use crate::permutations_const::PermutationsConst; + #[cfg(feature = "use_alloc")] mod combinations; #[cfg(feature = "use_alloc")] @@ -233,6 +235,7 @@ mod with_position; mod zip_eq_impl; mod zip_longest; mod ziptuple; +mod permutations_const; #[macro_export] /// Create an iterator over the “cartesian product” of iterators. @@ -1786,6 +1789,15 @@ pub trait Itertools: Iterator { permutations::permutations(self, k) } + #[cfg(feature = "use_alloc")] + fn permutations_const(self) -> PermutationsConst + where + Self: Sized, + Self::Item: Clone, + { + permutations_const::permutations_const(self) + } + /// Return an iterator that iterates through the powerset of the elements from an /// iterator. /// diff --git a/src/permutations_const.rs b/src/permutations_const.rs new file mode 100644 index 000000000..1dce2f565 --- /dev/null +++ b/src/permutations_const.rs @@ -0,0 +1,188 @@ +use alloc::boxed::Box; +use alloc::vec::Vec; +use std::fmt; +use std::iter::once; +use std::iter::FusedIterator; +use std::convert::TryInto; + +use super::lazy_buffer::LazyBuffer; +use crate::size_hint::{self, SizeHint}; + +/// 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 PermutationsConst { + vals: LazyBuffer, + state: PermutationState, +} + +impl Clone for PermutationsConst +where + I: Clone + Iterator, + I::Item: Clone, +{ + clone_fields!(vals, state); +} + +#[derive(Clone, Debug)] +enum PermutationState { + /// No permutation generated yet. + Start, + /// Values from the iterator are not fully loaded yet so `n` is still unknown. + Buffered { min_n: usize }, + /// All values from the iterator are known so `n` is known. + Loaded { + indices: Box<[usize]>, + cycles: Box<[usize]>, + }, + /// No permutation left to generate. + End, +} + +impl fmt::Debug for PermutationsConst +where + I: Iterator + fmt::Debug, + I::Item: fmt::Debug, +{ + debug_fmt_fields!(Permutations, vals, state); +} + +pub fn permutations_const(iter: I) -> PermutationsConst { + PermutationsConst { + vals: LazyBuffer::new(iter), + state: PermutationState::Start, + } +} + +impl Iterator for PermutationsConst +where + I: Iterator, + I::Item: Clone + Default, +{ + type Item = [I::Item; K]; + + fn next(&mut self) -> Option { + let Self { vals, state } = self; + match state { + PermutationState::Start => { + *state = PermutationState::End; + // TODO: Consider this case and handle it somehow, currently just using default + Some(std::array::from_fn(|_|I::Item::default())) + } + &mut PermutationState::Start => { + vals.prefill(K); + if vals.len() != K { + *state = PermutationState::End; + return None; + } + *state = PermutationState::Buffered { min_n: K }; + let mut iter = vals[0..K].into_iter().cloned(); + Some(std::array::from_fn(|_|iter.next().unwrap())) // TODO: Handle error case, maybe make this better + } + PermutationState::Buffered { min_n } => { + if vals.get_next() { + let mut item = (0..K - 1) + .chain(once(*min_n)) + .map(|i| vals[i].clone()); + *min_n += 1; + Some(std::array::from_fn(|_|item.next().unwrap())) + } else { + let n = *min_n; + let prev_iteration_count = n - K + 1; + let mut indices: Box<[_]> = (0..n).collect(); + let mut cycles: Box<[_]> = (n - K..n).rev().collect(); + // Advance the state to the correct point. + for _ in 0..prev_iteration_count { + if advance(&mut indices, &mut cycles) { + *state = PermutationState::End; + return None; + } + } + let item = vals.get_at(&indices[0..K]); // TODO: Impl const sized variant otherwise this is pointless + *state = PermutationState::Loaded { indices, cycles }; + Some(item.try_into().ok()?) // TODO: Handle error case + } + } + PermutationState::Loaded { indices, cycles } => { + if advance(indices, cycles) { + *state = PermutationState::End; + return None; + } + let k = cycles.len(); + Some(vals.get_at(&indices[0..k]).try_into().ok()?) // TODO: Handle error case and const size indexing + } + PermutationState::End => None, + } + } + + fn count(self) -> usize { + let Self { vals, state } = self; + let n = vals.count(); + state.size_hint_for(n).1.unwrap() + } + + fn size_hint(&self) -> SizeHint { + let (mut low, mut upp) = self.vals.size_hint(); + low = self.state.size_hint_for(low).0; + upp = upp.and_then(|n| self.state.size_hint_for(n).1); + (low, upp) + } +} + +impl FusedIterator for PermutationsConst +where + I: Iterator, + I::Item: Clone + Default, +{ +} + +fn advance(indices: &mut [usize], cycles: &mut [usize]) -> bool { + let n = indices.len(); + let k = cycles.len(); + // NOTE: if `cycles` are only zeros, then we reached the last permutation. + for i in (0..k).rev() { + if cycles[i] == 0 { + cycles[i] = n - i - 1; + indices[i..].rotate_left(1); + } else { + let swap_index = n - cycles[i]; + indices.swap(i, swap_index); + cycles[i] -= 1; + return false; + } + } + true +} + +impl PermutationState { + fn size_hint_for(&self, n: usize) -> SizeHint { + // At the beginning, there are `n!/(n-k)!` items to come. + let at_start = |n, k| { + debug_assert!(n >= k); + let total = (n - k + 1..=n).try_fold(1usize, |acc, i| acc.checked_mul(i)); + (total.unwrap_or(usize::MAX), total) + }; + match *self { + Self::Start if n < K => (0, Some(0)), + Self::Start => at_start(n, K), + Self::Buffered { min_n } => { + // Same as `Start` minus the previously generated items. + size_hint::sub_scalar(at_start(n, K), min_n - K + 1) + } + Self::Loaded { + ref indices, + ref cycles, + } => { + let count = cycles.iter().enumerate().try_fold(0usize, |acc, (i, &c)| { + acc.checked_mul(indices.len() - i) + .and_then(|count| count.checked_add(c)) + }); + (count.unwrap_or(usize::MAX), count) + } + Self::End => (0, Some(0)), + } + } +}