diff --git a/src/combinations.rs b/src/combinations.rs index 6bb2f3ec6..54a027551 100644 --- a/src/combinations.rs +++ b/src/combinations.rs @@ -1,3 +1,5 @@ +use core::array; +use core::borrow::BorrowMut; use std::fmt; use std::iter::FusedIterator; @@ -6,45 +8,101 @@ use alloc::vec::Vec; use crate::adaptors::checked_binomial; +/// Iterator for `Vec` valued combinations returned by [`.combinations()`](crate::Itertools::combinations) +pub type Combinations = CombinationsGeneric>; +/// Iterator for const generic combinations returned by [`.array_combinations()`](crate::Itertools::array_combinations) +pub type ArrayCombinations = CombinationsGeneric; + +/// Create a new `Combinations` from a clonable iterator. +pub fn combinations(iter: I, k: usize) -> Combinations +where + I::Item: Clone, +{ + Combinations::new(iter, (0..k).collect()) +} + +/// Create a new `ArrayCombinations` from a clonable iterator. +pub fn array_combinations(iter: I) -> ArrayCombinations +where + I::Item: Clone, +{ + ArrayCombinations::new(iter, array::from_fn(|i| i)) +} + /// An iterator to iterate through all the `k`-length combinations in an iterator. /// -/// See [`.combinations()`](crate::Itertools::combinations) for more information. +/// See [`.combinations()`](crate::Itertools::combinations) and [`.array_combinations()`](crate::Itertools::array_combinations) for more information. #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] -pub struct Combinations { - indices: Vec, +pub struct CombinationsGeneric { + indices: Idx, pool: LazyBuffer, first: bool, } -impl Clone for Combinations +/// A type holding indices of elements in a pool or buffer of items from an inner iterator +/// and used to pick out different combinations in a generic way. +pub trait PoolIndex: BorrowMut<[usize]> { + type Item; + + fn extract_item>(&self, pool: &LazyBuffer) -> Self::Item + where + T: Clone; + + fn len(&self) -> usize { + self.borrow().len() + } +} + +impl PoolIndex for Vec { + type Item = Vec; + + fn extract_item>(&self, pool: &LazyBuffer) -> Vec + where + T: Clone, + { + pool.get_at(self) + } +} + +impl PoolIndex for [usize; K] { + type Item = [T; K]; + + fn extract_item>(&self, pool: &LazyBuffer) -> [T; K] + where + T: Clone, + { + pool.get_array(*self) + } +} + +impl Clone for CombinationsGeneric where - I: Clone + Iterator, + I: Iterator + Clone, I::Item: Clone, + Idx: Clone, { clone_fields!(indices, pool, first); } -impl fmt::Debug for Combinations +impl fmt::Debug for CombinationsGeneric where I: Iterator + fmt::Debug, I::Item: fmt::Debug, + Idx: fmt::Debug, { debug_fmt_fields!(Combinations, indices, pool, first); } -/// Create a new `Combinations` from a clonable iterator. -pub fn combinations(iter: I, k: usize) -> Combinations -where - I: Iterator, -{ - Combinations { - indices: (0..k).collect(), - pool: LazyBuffer::new(iter), - first: true, +impl> CombinationsGeneric { + /// Constructor with arguments the inner iterator and the initial state for the indices. + fn new(iter: I, indices: Idx) -> Self { + Self { + indices, + pool: LazyBuffer::new(iter), + first: true, + } } -} -impl Combinations { /// Returns the length of a combination produced by this iterator. #[inline] pub fn k(&self) -> usize { @@ -64,27 +122,7 @@ impl Combinations { &self.pool } - /// Resets this `Combinations` back to an initial state for combinations of length - /// `k` over the same pool data source. If `k` is larger than the current length - /// of the data pool an attempt is made to prefill the pool so that it holds `k` - /// elements. - pub(crate) fn reset(&mut self, k: usize) { - self.first = true; - - if k < self.indices.len() { - self.indices.truncate(k); - for i in 0..k { - self.indices[i] = i; - } - } else { - for i in 0..self.indices.len() { - self.indices[i] = i; - } - self.indices.extend(self.indices.len()..k); - self.pool.prefill(k); - } - } - + /// Return the length of the inner iterator and the count of remaining combinations. pub(crate) fn n_and_count(self) -> (usize, usize) { let Self { indices, @@ -92,7 +130,7 @@ impl Combinations { first, } = self; let n = pool.count(); - (n, remaining_for(n, first, &indices).unwrap()) + (n, remaining_for(n, first, indices.borrow()).unwrap()) } /// Initialises the iterator by filling a buffer with elements from the @@ -113,19 +151,21 @@ impl Combinations { /// /// Returns true if we've run out of combinations, false otherwise. fn increment_indices(&mut self) -> bool { - if self.indices.is_empty() { + // Borrow once instead of noise each time it's indexed + let indices = self.indices.borrow_mut(); + + if indices.is_empty() { return true; // Done } - // Scan from the end, looking for an index to increment - let mut i: usize = self.indices.len() - 1; + let mut i: usize = indices.len() - 1; // Check if we need to consume more from the iterator - if self.indices[i] == self.pool.len() - 1 { + if indices[i] == self.pool.len() - 1 { self.pool.get_next(); // may change pool size } - while self.indices[i] == i + self.pool.len() - self.indices.len() { + while indices[i] == i + self.pool.len() - indices.len() { if i > 0 { i -= 1; } else { @@ -135,11 +175,10 @@ impl Combinations { } // Increment index, and reset the ones to its right - self.indices[i] += 1; - for j in i + 1..self.indices.len() { - self.indices[j] = self.indices[j - 1] + 1; + indices[i] += 1; + for j in i + 1..indices.len() { + indices[j] = indices[j - 1] + 1; } - // If we've made it this far, we haven't run out of combos false } @@ -147,6 +186,7 @@ impl Combinations { /// Returns the n-th item or the number of successful steps. pub(crate) fn try_nth(&mut self, n: usize) -> Result<::Item, usize> where + I: Iterator, I::Item: Clone, { let done = if self.first { @@ -162,16 +202,17 @@ impl Combinations { return Err(i + 1); } } - Ok(self.pool.get_at(&self.indices)) + Ok(self.indices.extract_item(&self.pool)) } } -impl Iterator for Combinations +impl Iterator for CombinationsGeneric where I: Iterator, I::Item: Clone, + Idx: PoolIndex, { - type Item = Vec; + type Item = Idx::Item; fn next(&mut self) -> Option { let done = if self.first { self.init() @@ -183,7 +224,7 @@ where return None; } - Some(self.pool.get_at(&self.indices)) + Some(self.indices.extract_item(&self.pool)) } fn nth(&mut self, n: usize) -> Option { @@ -192,8 +233,8 @@ where fn size_hint(&self) -> (usize, Option) { let (mut low, mut upp) = self.pool.size_hint(); - low = remaining_for(low, self.first, &self.indices).unwrap_or(usize::MAX); - upp = upp.and_then(|upp| remaining_for(upp, self.first, &self.indices)); + low = remaining_for(low, self.first, self.indices.borrow()).unwrap_or(usize::MAX); + upp = upp.and_then(|upp| remaining_for(upp, self.first, self.indices.borrow())); (low, upp) } @@ -203,13 +244,37 @@ where } } -impl FusedIterator for Combinations +impl FusedIterator for CombinationsGeneric where I: Iterator, I::Item: Clone, + Idx: PoolIndex, { } +impl Combinations { + /// Resets this `Combinations` back to an initial state for combinations of length + /// `k` over the same pool data source. If `k` is larger than the current length + /// of the data pool an attempt is made to prefill the pool so that it holds `k` + /// elements. + pub(crate) fn reset(&mut self, k: usize) { + self.first = true; + + if k < self.indices.len() { + self.indices.truncate(k); + for i in 0..k { + self.indices[i] = i; + } + } else { + for i in 0..self.indices.len() { + self.indices[i] = i; + } + self.indices.extend(self.indices.len()..k); + self.pool.prefill(k); + } + } +} + /// For a given size `n`, return the count of remaining combinations or None if it would overflow. fn remaining_for(n: usize, first: bool, indices: &[usize]) -> Option { let k = indices.len(); diff --git a/src/lazy_buffer.rs b/src/lazy_buffer.rs index fefcff8f5..fafa5f726 100644 --- a/src/lazy_buffer.rs +++ b/src/lazy_buffer.rs @@ -59,6 +59,10 @@ where pub fn get_at(&self, indices: &[usize]) -> Vec { indices.iter().map(|i| self.buffer[*i].clone()).collect() } + + pub fn get_array(&self, indices: [usize; K]) -> [I::Item; K] { + indices.map(|i| self.buffer[i].clone()) + } } impl Index for LazyBuffer diff --git a/src/lib.rs b/src/lib.rs index 9670390a9..dc316d847 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -97,7 +97,7 @@ pub mod structs { TakeWhileRef, TupleCombinations, Update, WhileSome, }; #[cfg(feature = "use_alloc")] - pub use crate::combinations::Combinations; + pub use crate::combinations::{ArrayCombinations, Combinations}; #[cfg(feature = "use_alloc")] pub use crate::combinations_with_replacement::CombinationsWithReplacement; pub use crate::cons_tuples_impl::ConsTuples; @@ -1674,6 +1674,53 @@ pub trait Itertools: Iterator { adaptors::tuple_combinations(self) } + /// Return an iterator adaptor that iterates over the combinations 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. + /// + /// # Guarantees + /// + /// If the adapted iterator is deterministic, + /// this iterator adapter yields items in a reliable order. + /// + /// ``` + /// use itertools::Itertools; + /// + /// let mut v = Vec::new(); + /// for [a, b] in (1..5).array_combinations() { + /// v.push([a, b]); + /// } + /// assert_eq!(v, vec![[1, 2], [1, 3], [1, 4], [2, 3], [2, 4], [3, 4]]); + /// + /// let mut it = (1..5).array_combinations(); + /// assert_eq!(Some([1, 2, 3]), it.next()); + /// assert_eq!(Some([1, 2, 4]), it.next()); + /// assert_eq!(Some([1, 3, 4]), it.next()); + /// assert_eq!(Some([2, 3, 4]), it.next()); + /// assert_eq!(None, it.next()); + /// + /// // this requires a type hint + /// let it = (1..5).array_combinations::<3>(); + /// itertools::assert_equal(it, vec![[1, 2, 3], [1, 2, 4], [1, 3, 4], [2, 3, 4]]); + /// + /// // you can also specify the complete type + /// use itertools::ArrayCombinations; + /// use std::ops::Range; + /// + /// let it: ArrayCombinations, 3> = (1..5).array_combinations(); + /// itertools::assert_equal(it, vec![[1, 2, 3], [1, 2, 4], [1, 3, 4], [2, 3, 4]]); + /// ``` + #[cfg(feature = "use_alloc")] + fn array_combinations(self) -> ArrayCombinations + where + Self: Sized + Clone, + Self::Item: Clone, + { + combinations::array_combinations(self) + } + /// Return an iterator adaptor that iterates over the `k`-length combinations of /// the elements from an iterator. /// diff --git a/tests/specializations.rs b/tests/specializations.rs index 3e4831024..e6694c8e7 100644 --- a/tests/specializations.rs +++ b/tests/specializations.rs @@ -273,6 +273,16 @@ quickcheck! { test_specializations(&v.into_iter().intersperse_with(|| 0)); } + fn array_combinations(v: Vec) -> TestResult { + if v.len() > 10 { + return TestResult::discard(); + } + test_specializations(&v.iter().array_combinations::<1>()); + test_specializations(&v.iter().array_combinations::<2>()); + test_specializations(&v.iter().array_combinations::<3>()); + TestResult::passed() + } + fn combinations(a: Vec, n: u8) -> TestResult { if n > 3 || a.len() > 8 { return TestResult::discard();