diff options
| author | mo khan <mo@mokhan.ca> | 2025-07-02 18:36:06 -0600 |
|---|---|---|
| committer | mo khan <mo@mokhan.ca> | 2025-07-02 18:36:06 -0600 |
| commit | 8cdfa445d6629ffef4cb84967ff7017654045bc2 (patch) | |
| tree | 22f0b0907c024c78d26a731e2e1f5219407d8102 /vendor/itertools/src/permutations.rs | |
| parent | 4351c74c7c5f97156bc94d3a8549b9940ac80e3f (diff) | |
chore: add vendor directory
Diffstat (limited to 'vendor/itertools/src/permutations.rs')
| -rw-r--r-- | vendor/itertools/src/permutations.rs | 186 |
1 files changed, 186 insertions, 0 deletions
diff --git a/vendor/itertools/src/permutations.rs b/vendor/itertools/src/permutations.rs new file mode 100644 index 00000000..91389a73 --- /dev/null +++ b/vendor/itertools/src/permutations.rs @@ -0,0 +1,186 @@ +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 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 Permutations<I: Iterator> { + vals: LazyBuffer<I>, + state: PermutationState, +} + +impl<I> Clone for Permutations<I> +where + I: Clone + Iterator, + I::Item: Clone, +{ + clone_fields!(vals, state); +} + +#[derive(Clone, Debug)] +enum PermutationState { + /// No permutation generated yet. + Start { k: usize }, + /// Values from the iterator are not fully loaded yet so `n` is still unknown. + Buffered { k: usize, 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<I> fmt::Debug for Permutations<I> +where + I: Iterator + fmt::Debug, + I::Item: fmt::Debug, +{ + debug_fmt_fields!(Permutations, vals, state); +} + +pub fn permutations<I: Iterator>(iter: I, k: usize) -> Permutations<I> { + Permutations { + vals: LazyBuffer::new(iter), + state: PermutationState::Start { k }, + } +} + +impl<I> Iterator for Permutations<I> +where + I: Iterator, + I::Item: Clone, +{ + type Item = Vec<I::Item>; + + fn next(&mut self) -> Option<Self::Item> { + 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 { + *state = PermutationState::End; + return None; + } + *state = PermutationState::Buffered { k, min_n: k }; + Some(vals[0..k].to_vec()) + } + 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(); + *min_n += 1; + Some(item) + } 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]); + *state = PermutationState::Loaded { indices, cycles }; + Some(item) + } + } + PermutationState::Loaded { indices, cycles } => { + if advance(indices, cycles) { + *state = PermutationState::End; + return None; + } + let k = cycles.len(); + Some(vals.get_at(&indices[0..k])) + } + 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<I> FusedIterator for Permutations<I> +where + I: Iterator, + I::Item: Clone, +{ +} + +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 { k } if n < k => (0, Some(0)), + Self::Start { k } => at_start(n, k), + Self::Buffered { k, 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)), + } + } +} |
