diff --git a/src/lazy_buffer.rs b/src/lazy_buffer.rs index fafa5f726..69ece7b19 100644 --- a/src/lazy_buffer.rs +++ b/src/lazy_buffer.rs @@ -1,4 +1,5 @@ use alloc::vec::Vec; +use core::array::from_fn; use std::iter::Fuse; use std::ops::Index; @@ -63,6 +64,13 @@ where pub fn get_array(&self, indices: [usize; K]) -> [I::Item; K] { indices.map(|i| self.buffer[i].clone()) } + + pub fn get_array_from_fn( + &self, + mut f: impl FnMut(usize) -> usize, + ) -> [I::Item; K] { + from_fn(|i| self.buffer[f(i)].clone()) + } } impl Index for LazyBuffer diff --git a/src/lib.rs b/src/lib.rs index cba3ad570..a2e17fd48 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -124,7 +124,7 @@ pub mod structs { pub use crate::peek_nth::PeekNth; pub use crate::peeking_take_while::PeekingTakeWhile; #[cfg(feature = "use_alloc")] - pub use crate::permutations::Permutations; + pub use crate::permutations::{ArrayPermutations, Permutations}; #[cfg(feature = "use_alloc")] pub use crate::powerset::Powerset; pub use crate::process_results_impl::ProcessResults; @@ -1792,6 +1792,51 @@ pub trait Itertools: Iterator { combinations_with_replacement::combinations_with_replacement(self, k) } + /// Return an iterator adaptor that iterates over all `K`-permutations of the + /// elements from an iterator. + /// + /// Iterator element type is `[Self::Item; K]`. The iterator + /// produces a new array per iteration, and clones the iterator elements. + /// + /// If `K` is greater than the length of the input iterator, the resultant + /// iterator adaptor will be empty. + /// ``` + /// use itertools::Itertools; + /// + /// let perms = (5..8).array_permutations(); + /// itertools::assert_equal(perms, vec![ + /// [5, 6], + /// [5, 7], + /// [6, 5], + /// [6, 7], + /// [7, 5], + /// [7, 6], + /// ]); + /// ``` + /// + /// Note: Permutations does not take into account the equality of the iterated values. + /// + /// ``` + /// use itertools::Itertools; + /// + /// let it = vec![-2, -2].into_iter().array_permutations(); + /// itertools::assert_equal(it, vec![ + /// [-2, -2], // Note: these are the same + /// [-2, -2], // Note: these are the same + /// ]); + /// ``` + /// + /// Note: The source iterator is collected lazily, and will not be + /// re-iterated if the permutations adaptor is completed and re-iterated. + #[cfg(feature = "use_alloc")] + fn array_permutations(self) -> ArrayPermutations + where + Self: Sized, + Self::Item: Clone, + { + permutations::array_permutations(self) + } + /// Return an iterator adaptor that iterates over all k-permutations of the /// elements from an iterator. /// @@ -1823,10 +1868,10 @@ pub trait Itertools: Iterator { /// ``` /// use itertools::Itertools; /// - /// let it = vec![2, 2].into_iter().permutations(2); + /// let it = vec![-2, -2].into_iter().permutations(2); /// itertools::assert_equal(it, vec![ - /// vec![2, 2], // Note: these are the same - /// vec![2, 2], // Note: these are the same + /// vec![-2, -2], // Note: these are the same + /// vec![-2, -2], // Note: these are the same /// ]); /// ``` /// diff --git a/src/permutations.rs b/src/permutations.rs index 91389a73a..29e97ac1a 100644 --- a/src/permutations.rs +++ b/src/permutations.rs @@ -1,5 +1,6 @@ use alloc::boxed::Box; use alloc::vec::Vec; +use core::marker::PhantomData; use std::fmt; use std::iter::once; use std::iter::FusedIterator; @@ -7,23 +8,32 @@ use std::iter::FusedIterator; use super::lazy_buffer::LazyBuffer; use crate::size_hint::{self, SizeHint}; +/// Iterator for `Vec` valued permutations returned by +/// [`.permutations()`](crate::Itertools::permutations) +pub type ArrayPermutations = PermutationsGeneric::Item; K]>; +/// Iterator for const generic permutations returned by +/// [`.array_permutations()`](crate::Itertools::array_permutations) +pub type Permutations = PermutationsGeneric::Item>>; + /// An iterator adaptor that iterates through all the `k`-permutations of the /// elements from an iterator. /// -/// See [`.permutations()`](crate::Itertools::permutations) for +/// See [`.permutations()`](crate::Itertools::permutations) and +/// [`.array_permutations()`](crate::Itertools::array_permutations) for /// more information. #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] -pub struct Permutations { +pub struct PermutationsGeneric { vals: LazyBuffer, state: PermutationState, + _item: PhantomData, } -impl Clone for Permutations +impl Clone for PermutationsGeneric where I: Clone + Iterator, I::Item: Clone, { - clone_fields!(vals, state); + clone_fields!(vals, state, _item); } #[derive(Clone, Debug)] @@ -41,7 +51,7 @@ enum PermutationState { End, } -impl fmt::Debug for Permutations +impl fmt::Debug for PermutationsGeneric where I: Iterator + fmt::Debug, I::Item: fmt::Debug, @@ -49,26 +59,36 @@ where debug_fmt_fields!(Permutations, vals, state); } +pub fn array_permutations(iter: I) -> ArrayPermutations { + PermutationsGeneric { + vals: LazyBuffer::new(iter), + state: PermutationState::Start { k: K }, + _item: PhantomData, + } +} + pub fn permutations(iter: I, k: usize) -> Permutations { - Permutations { + PermutationsGeneric { vals: LazyBuffer::new(iter), state: PermutationState::Start { k }, + _item: PhantomData, } } -impl Iterator for Permutations +impl Iterator for PermutationsGeneric where I: Iterator, I::Item: Clone, + Item: PermItem, { - type Item = Vec; + type Item = Item; fn next(&mut self) -> Option { - let Self { vals, state } = self; + let Self { vals, state, _item } = self; match state { PermutationState::Start { k: 0 } => { *state = PermutationState::End; - Some(Vec::new()) + Some(Item::extract_start(vals, 0, 0)) } &mut PermutationState::Start { k } => { vals.prefill(k); @@ -77,14 +97,11 @@ where return None; } *state = PermutationState::Buffered { k, min_n: k }; - Some(vals[0..k].to_vec()) + Some(Item::extract_start(vals, k, k - 1)) } PermutationState::Buffered { ref k, min_n } => { if vals.get_next() { - let item = (0..*k - 1) - .chain(once(*min_n)) - .map(|i| vals[i].clone()) - .collect(); + let item = Item::extract_start(vals, *k, *min_n); *min_n += 1; Some(item) } else { @@ -99,7 +116,7 @@ where return None; } } - let item = vals.get_at(&indices[0..*k]); + let item = Item::extract_from_prefix(vals, &indices[0..*k]); *state = PermutationState::Loaded { indices, cycles }; Some(item) } @@ -110,14 +127,14 @@ where return None; } let k = cycles.len(); - Some(vals.get_at(&indices[0..k])) + Some(Item::extract_from_prefix(vals, &indices[0..k])) } PermutationState::End => None, } } fn count(self) -> usize { - let Self { vals, state } = self; + let Self { vals, state, _item } = self; let n = vals.count(); state.size_hint_for(n).1.unwrap() } @@ -130,10 +147,11 @@ where } } -impl FusedIterator for Permutations +impl FusedIterator for PermutationsGeneric where I: Iterator, I::Item: Clone, + Item: PermItem, { } @@ -184,3 +202,39 @@ impl PermutationState { } } } + +/// A type that can be picked out from a pool or buffer of items from an inner iterator +/// and in a generic way given their indices. +pub trait PermItem { + fn extract_start>(buf: &LazyBuffer, len: usize, last: usize) -> Self; + + fn extract_from_prefix>(buf: &LazyBuffer, indices: &[usize]) -> Self; +} + +impl PermItem for [T; K] { + fn extract_start>(buf: &LazyBuffer, len: usize, last: usize) -> Self { + assert_eq!(len, K); + buf.get_array_from_fn(|i| if i + 1 < len { i } else { last }) + } + + fn extract_from_prefix>(buf: &LazyBuffer, indices: &[usize]) -> Self { + buf.get_array_from_fn(|i| indices[i]) + } +} + +impl PermItem for Vec { + fn extract_start>(buf: &LazyBuffer, len: usize, last: usize) -> Self { + if len == 0 { + Vec::new() + } else { + (0..len - 1) + .chain(once(last)) + .map(|i| buf[i].clone()) + .collect() + } + } + + fn extract_from_prefix>(buf: &LazyBuffer, indices: &[usize]) -> Self { + buf.get_at(indices) + } +} diff --git a/tests/laziness.rs b/tests/laziness.rs index c559d33ad..c6d53fbd0 100644 --- a/tests/laziness.rs +++ b/tests/laziness.rs @@ -207,6 +207,11 @@ must_use_tests! { tuple_combinations3 { let _ = Panicking.tuple_combinations::<(_, _, _)>(); } + array_combinations { + let _ = Panicking.array_combinations::<0>(); + let _ = Panicking.array_combinations::<1>(); + let _ = Panicking.array_combinations::<2>(); + } combinations { let _ = Panicking.combinations(0); let _ = Panicking.combinations(1); @@ -217,6 +222,11 @@ must_use_tests! { let _ = Panicking.combinations_with_replacement(1); let _ = Panicking.combinations_with_replacement(2); } + array_permutations { + let _ = Panicking.array_permutations::<0>(); + let _ = Panicking.array_permutations::<1>(); + let _ = Panicking.array_permutations::<2>(); + } permutations { let _ = Panicking.permutations(0); let _ = Panicking.permutations(1);