/
array_combinations.rs
119 lines (107 loc) · 3.32 KB
/
array_combinations.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
use core::fmt;
use core::fmt::Debug;
use core::iter::FusedIterator;
use crate::adaptors::generic_combinations::GenericCombinations;
/// An extension trait that provides the [`array_combinations`] method for
/// iterators.
///
/// [`array_combinations`]: IterArrayCombinations::array_combinations
#[cfg_attr(docsrs, doc(cfg(feature = "array_combinations")))]
pub trait IterArrayCombinations: Iterator {
/// Returns an iterator adaptor that iterates over `K` length combinations
/// of all the elements in the underlying iterator.
///
/// The iterator is consumed as elements are required. In the first
/// iteration `K` elements will be consumed by the iterator.
///
/// # Panics
///
/// If called with `K = 0`.
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// use itermore::IterArrayCombinations;
///
/// let mut iter = "abcd".chars().array_combinations();
/// assert_eq!(iter.next(), Some(['a', 'b', 'c']));
/// assert_eq!(iter.next(), Some(['a', 'b', 'd']));
/// assert_eq!(iter.next(), Some(['a', 'c', 'd']));
/// assert_eq!(iter.next(), Some(['b', 'c', 'd']));
/// assert_eq!(iter.next(), None);
/// ```
#[inline]
fn array_combinations<const K: usize>(self) -> ArrayCombinations<Self, K>
where
Self: Sized,
Self::Item: Clone,
{
ArrayCombinations::new(self)
}
}
impl<I: ?Sized> IterArrayCombinations for I where I: Iterator {}
/// An iterator that iterates over `K` length combinations of all the elements
/// in the underlying iterator.
///
/// This struct is created by the [`array_combinations`] method on iterators.
/// See its documentation for more.
///
/// [`array_combinations`]: IterArrayCombinations::array_combinations
#[cfg_attr(docsrs, doc(cfg(feature = "array_combinations")))]
#[must_use = "iterators are lazy and do nothing unless consumed"]
pub struct ArrayCombinations<I, const K: usize>(GenericCombinations<I, [usize; K]>)
where
I: Iterator;
impl<I, const K: usize> ArrayCombinations<I, K>
where
I: Iterator,
I::Item: Clone,
{
#[track_caller]
pub(crate) fn new(iter: I) -> Self {
assert!(K != 0, "combination size must be non-zero");
// SAFETY: The range 0..K yields at least K elements.
let comb = unsafe { arrays::from_iter_unchecked(0..K) };
Self(GenericCombinations::new(iter, comb))
}
}
impl<I, const K: usize> Clone for ArrayCombinations<I, K>
where
I: Iterator + Clone,
I::Item: Clone,
{
fn clone(&self) -> Self {
Self(self.0.clone())
}
}
impl<I, const K: usize> Debug for ArrayCombinations<I, K>
where
I: Iterator + Debug,
I::Item: Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt_with(f, "ArrayCombinations")
}
}
impl<I, const K: usize> Iterator for ArrayCombinations<I, K>
where
I: Iterator,
I::Item: Clone,
{
type Item = [I::Item; K];
fn next(&mut self) -> Option<Self::Item> {
self.0.fill_next().map(|it| {
// SAFETY: The iterator is guaranteed to yield K elements because
// it is derived from `self.0.comb` which is an array of length K.
unsafe { arrays::from_iter_unchecked(it) }
})
}
}
impl<I, const K: usize> FusedIterator for ArrayCombinations<I, K>
where
I: Iterator,
I::Item: Clone,
{
}