diff --git a/src/combinations.rs b/src/combinations.rs index 54a027551..b7c5bb275 100644 --- a/src/combinations.rs +++ b/src/combinations.rs @@ -1,9 +1,7 @@ -use core::array; -use core::borrow::BorrowMut; use std::fmt; use std::iter::FusedIterator; -use super::lazy_buffer::LazyBuffer; +use super::lazy_buffer::{ArrayOrVecHelper, ConstUsize, LazyBuffer, MaybeConstUsize}; use alloc::vec::Vec; use crate::adaptors::checked_binomial; @@ -18,7 +16,7 @@ pub fn combinations(iter: I, k: usize) -> Combinations where I::Item: Clone, { - Combinations::new(iter, (0..k).collect()) + Combinations::new(iter, ArrayOrVecHelper::start(k)) } /// Create a new `ArrayCombinations` from a clonable iterator. @@ -26,7 +24,7 @@ pub fn array_combinations(iter: I) -> ArrayCombinat where I::Item: Clone, { - ArrayCombinations::new(iter, array::from_fn(|i| i)) + ArrayCombinations::new(iter, ArrayOrVecHelper::start(ConstUsize::)) } /// An iterator to iterate through all the `k`-length combinations in an iterator. @@ -39,42 +37,6 @@ pub struct CombinationsGeneric { first: bool, } -/// 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: Iterator + Clone, @@ -93,7 +55,7 @@ where debug_fmt_fields!(Combinations, indices, pool, first); } -impl> CombinationsGeneric { +impl CombinationsGeneric { /// Constructor with arguments the inner iterator and the initial state for the indices. fn new(iter: I, indices: Idx) -> Self { Self { @@ -105,7 +67,7 @@ impl> CombinationsGeneric { /// Returns the length of a combination produced by this iterator. #[inline] - pub fn k(&self) -> usize { + pub fn k(&self) -> Idx::Length { self.indices.len() } @@ -136,8 +98,8 @@ impl> CombinationsGeneric { /// Initialises the iterator by filling a buffer with elements from the /// iterator. Returns true if there are no combinations, false otherwise. fn init(&mut self) -> bool { - self.pool.prefill(self.k()); - let done = self.k() > self.n(); + self.pool.prefill(self.k().value()); + let done = self.k().value() > self.n(); if !done { self.first = false; } @@ -210,9 +172,9 @@ impl Iterator for CombinationsGeneric where I: Iterator, I::Item: Clone, - Idx: PoolIndex, + Idx: ArrayOrVecHelper, { - type Item = Idx::Item; + type Item = Idx::Item; fn next(&mut self) -> Option { let done = if self.first { self.init() @@ -248,7 +210,7 @@ impl FusedIterator for CombinationsGeneric where I: Iterator, I::Item: Clone, - Idx: PoolIndex, + Idx: ArrayOrVecHelper, { } diff --git a/src/lazy_buffer.rs b/src/lazy_buffer.rs index fafa5f726..bb064f23b 100644 --- a/src/lazy_buffer.rs +++ b/src/lazy_buffer.rs @@ -1,4 +1,7 @@ +use alloc::boxed::Box; use alloc::vec::Vec; +use core::borrow::BorrowMut; +use core::ops::Deref; use std::iter::Fuse; use std::ops::Index; @@ -77,3 +80,121 @@ where self.buffer.index(index) } } + +pub trait MaybeConstUsize: Clone + Copy + std::fmt::Debug { + /*TODO const*/ + fn value(self) -> usize; +} + +#[derive(Clone, Copy, Debug)] +pub struct ConstUsize; +impl MaybeConstUsize for ConstUsize { + fn value(self) -> usize { + N + } +} + +impl MaybeConstUsize for usize { + fn value(self) -> usize { + self + } +} + +/// A type holding indices, mostly used to pick out different combinations of elements from +/// a pool or buffer of items from an inner iterator in a generic way. +pub trait ArrayOrVecHelper: BorrowMut<[usize]> { + type Item; + type Length: MaybeConstUsize; + + fn extract_item(&self, pool: &LazyBuffer) -> Self::Item + where + I::Item: Clone; + + // TODO if only ever used to index into LazyBuffer, specialize to + // extract_from_fn(Self::Length, Fn(usize) -> usize, &LazyBuffer) -> Item? + fn item_from_fn T>(len: Self::Length, f: F) -> Self::Item; + + fn len(&self) -> Self::Length; + + fn from_fn usize>(k: Self::Length, f: F) -> Self; + + /// Create an array/vec/... of indices from 0 to `len - 1`. + fn start(len: Self::Length) -> Self + where + Self: Sized, + { + Self::from_fn(len, |i| i) + } +} + +impl ArrayOrVecHelper for Vec { + type Item = Vec; + type Length = usize; + + fn extract_item(&self, pool: &LazyBuffer) -> Self::Item + where + I::Item: Clone, + { + pool.get_at(self) + } + + fn item_from_fn T>(len: Self::Length, f: F) -> Self::Item { + (0..len).map(f).collect() + } + + fn len(&self) -> Self::Length { + self.len() + } + + fn from_fn usize>(k: Self::Length, f: F) -> Self { + (0..k).map(f).collect() + } +} + +impl ArrayOrVecHelper for Box<[usize]> { + type Item = Vec; + type Length = usize; + + fn extract_item(&self, pool: &LazyBuffer) -> Self::Item + where + I::Item: Clone, + { + pool.get_at(self) + } + + fn item_from_fn T>(len: Self::Length, f: F) -> Self::Item { + (0..len).map(f).collect() + } + + fn len(&self) -> Self::Length { + self.deref().len() + } + + fn from_fn usize>(k: Self::Length, f: F) -> Self { + (0..k).map(f).collect() + } +} + +impl ArrayOrVecHelper for [usize; K] { + type Item = [T; K]; + type Length = ConstUsize; + + fn extract_item(&self, pool: &LazyBuffer) -> Self::Item + where + I::Item: Clone, + { + pool.get_array(*self) + } + + fn item_from_fn T>(_len: Self::Length, f: F) -> Self::Item { + std::array::from_fn(f) + } + + fn len(&self) -> Self::Length { + ConstUsize:: + } + + fn from_fn usize>(_len: Self::Length, f: F) -> Self { + std::array::from_fn(f) + } +} diff --git a/src/permutations.rs b/src/permutations.rs index 91389a73a..682d97f70 100644 --- a/src/permutations.rs +++ b/src/permutations.rs @@ -1,52 +1,51 @@ use alloc::boxed::Box; -use alloc::vec::Vec; use std::fmt; -use std::iter::once; use std::iter::FusedIterator; -use super::lazy_buffer::LazyBuffer; +use super::lazy_buffer::{ArrayOrVecHelper, LazyBuffer, MaybeConstUsize}; use crate::size_hint::{self, SizeHint}; +#[must_use = "iterator adaptors are lazy and do nothing unless consumed"] +pub struct PermutationsGeneric { + vals: LazyBuffer, + state: PermutationState, +} + /// 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, -} +pub type Permutations = PermutationsGeneric>; -impl Clone for Permutations +impl Clone for PermutationsGeneric where I: Clone + Iterator, I::Item: Clone, + Idx: Clone + ArrayOrVecHelper, { clone_fields!(vals, state); } #[derive(Clone, Debug)] -enum PermutationState { +enum PermutationState { /// No permutation generated yet. - Start { k: usize }, + Start { k: Idx::Length }, /// Values from the iterator are not fully loaded yet so `n` is still unknown. - Buffered { k: usize, min_n: usize }, + Buffered { indices: Idx, min_n: usize }, /// All values from the iterator are known so `n` is known. - Loaded { - indices: Box<[usize]>, - cycles: Box<[usize]>, - }, + Loaded { indices: Box<[usize]>, cycles: Idx }, /// No permutation left to generate. End, } -impl fmt::Debug for Permutations +impl fmt::Debug for PermutationsGeneric where I: Iterator + fmt::Debug, I::Item: fmt::Debug, + Idx: fmt::Debug, { - debug_fmt_fields!(Permutations, vals, state); + debug_fmt_fields!(PermutationsGeneric, vals, state); } pub fn permutations(iter: I, k: usize) -> Permutations { @@ -56,61 +55,70 @@ pub fn permutations(iter: I, k: usize) -> Permutations { } } -impl Iterator for Permutations +impl Iterator for PermutationsGeneric where I: Iterator, I::Item: Clone, { - type Item = Vec; + type Item = Idx::Item; fn next(&mut self) -> Option { let Self { vals, state } = self; match state { - PermutationState::Start { k: 0 } => { - *state = PermutationState::End; - Some(Vec::new()) - } &mut PermutationState::Start { k } => { - vals.prefill(k); - if vals.len() != k { + if k.value() == 0 { *state = PermutationState::End; - return None; + } else { + vals.prefill(k.value()); + if vals.len() != k.value() { + *state = PermutationState::End; + return None; + } + *state = PermutationState::Buffered { + indices: Idx::start(k), + min_n: k.value(), + }; } - *state = PermutationState::Buffered { k, min_n: k }; - Some(vals[0..k].to_vec()) + Some(Idx::item_from_fn(k, |i| vals[i].clone())) } - PermutationState::Buffered { ref k, min_n } => { + PermutationState::Buffered { indices, min_n } => { + let k = indices.len(); if vals.get_next() { - let item = (0..*k - 1) - .chain(once(*min_n)) - .map(|i| vals[i].clone()) - .collect(); + // This could be an unwrap instead of ?, indices.len() should be + // k.value() > 0 to be in this case + *indices.borrow_mut().last_mut()? += 1; + debug_assert!(indices + .borrow() + .iter() + .copied() + .eq((0..k.value() - 1).chain([*min_n]))); *min_n += 1; - Some(item) + Some(indices.extract_item(vals)) } else { let n = *min_n; - let prev_iteration_count = n - *k + 1; + let prev_iteration_count = n - k.value() + 1; let mut indices: Box<[_]> = (0..n).collect(); - let mut cycles: Box<[_]> = (n - k..n).rev().collect(); + let mut cycles = Idx::from_fn(k, |i| n - 1 - i); // Advance the state to the correct point. for _ in 0..prev_iteration_count { - if advance(&mut indices, &mut cycles) { + if advance(&mut indices, cycles.borrow_mut()) { *state = PermutationState::End; return None; } } - let item = vals.get_at(&indices[0..*k]); + let item = Idx::item_from_fn(k, |i| vals[indices[i]].clone()); *state = PermutationState::Loaded { indices, cycles }; Some(item) } } PermutationState::Loaded { indices, cycles } => { - if advance(indices, cycles) { + if advance(indices, cycles.borrow_mut()) { *state = PermutationState::End; return None; } - let k = cycles.len(); - Some(vals.get_at(&indices[0..k])) + Some(Idx::item_from_fn(cycles.len(), |i| { + vals[indices[i]].clone() + })) } PermutationState::End => None, } @@ -130,7 +138,7 @@ where } } -impl FusedIterator for Permutations +impl FusedIterator for PermutationsGeneric where I: Iterator, I::Item: Clone, @@ -155,29 +163,34 @@ fn advance(indices: &mut [usize], cycles: &mut [usize]) -> bool { true } -impl PermutationState { +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)); + let at_start = |n, k: Idx::Length| { + debug_assert!(n >= k.value()); + let total = (n - k.value() + 1..=n).try_fold(1usize, |acc, i| acc.checked_mul(i)); (total.unwrap_or(usize::MAX), total) }; - match *self { - Self::Start { k } if n < k => (0, Some(0)), - Self::Start { k } => at_start(n, k), - Self::Buffered { k, min_n } => { + match self { + Self::Start { k } if n < k.value() => (0, Some(0)), + Self::Start { k } => at_start(n, *k), + Self::Buffered { indices, min_n } => { + let k = indices.len(); // Same as `Start` minus the previously generated items. - size_hint::sub_scalar(at_start(n, k), min_n - k + 1) + size_hint::sub_scalar(at_start(n, k), min_n - k.value() + 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)) - }); + let count = cycles + .borrow() + .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)),