1use core::iter::FlatMap;
2
3/// A specialized version of `core::iter::FlatMap` for mapping over exact-sized
4/// iterators with a function that returns an array.
5///
6/// `ArrayFlatMap` differs from `FlatMap` in that `ArrayFlatMap` implements
7/// `ExactSizeIterator`. Since the result of `F` always has `LEN` elements, if
8/// `I` is an exact-sized iterator of length `inner_len` then we know the
9/// length of the flat-mapped result is `inner_len * LEN`. (The constructor
10/// verifies that this multiplication doesn't overflow `usize`.)
11#[derive(Clone)]
12pub struct ArrayFlatMap<I, Item, F, const LEN: usize> {
13 inner: FlatMap<I, [Item; LEN], F>,
14 remaining: usize,
15}
16
17impl<I, Item, F, const LEN: usize> ArrayFlatMap<I, Item, F, LEN>
18where
19 I: ExactSizeIterator,
20 F: FnMut(I::Item) -> [Item; LEN],
21{
22 /// Constructs an `ArrayFlatMap` wrapping the given iterator, using the
23 /// given function
24 pub fn new(inner: I, f: F) -> Option<Self> {
25 let remaining: usize = inner.len().checked_mul(LEN)?;
26 let inner: impl Iterator = inner.flat_map(f);
27 Some(Self { inner, remaining })
28 }
29}
30
31impl<I, Item, F, const LEN: usize> Iterator for ArrayFlatMap<I, Item, F, LEN>
32where
33 I: Iterator,
34 F: FnMut(I::Item) -> [Item; LEN],
35{
36 type Item = Item;
37
38 fn next(&mut self) -> Option<Self::Item> {
39 let result: Option = self.inner.next();
40 if result.is_some() {
41 self.remaining -= 1;
42 }
43 result
44 }
45
46 /// Required for implementing `ExactSizeIterator`.
47 fn size_hint(&self) -> (usize, Option<usize>) {
48 (self.remaining, Some(self.remaining))
49 }
50}
51
52impl<I, Item, F, const LEN: usize> ExactSizeIterator for ArrayFlatMap<I, Item, F, LEN>
53where
54 I: Iterator,
55 F: FnMut(I::Item) -> [Item; LEN],
56{
57}
58
59#[cfg(test)]
60mod tests {
61 use super::*;
62 use core::mem::size_of;
63
64 #[test]
65 fn test_array_flat_map() {
66 static TEST_CASES: &[(&[u16], fn(u16) -> [u8; 2], &[u8])] = &[
67 // Empty input
68 (&[], u16::to_be_bytes, &[]),
69 // Non-empty input.
70 (
71 &[0x0102, 0x0304, 0x0506],
72 u16::to_be_bytes,
73 &[1, 2, 3, 4, 5, 6],
74 ),
75 // Test with a different mapping function.
76 (
77 &[0x0102, 0x0304, 0x0506],
78 u16::to_le_bytes,
79 &[2, 1, 4, 3, 6, 5],
80 ),
81 ];
82 TEST_CASES.iter().copied().for_each(|(input, f, expected)| {
83 let mapped = ArrayFlatMap::new(input.iter().copied(), f).unwrap();
84 super::super::test::assert_iterator(mapped, expected);
85 });
86 }
87
88 // Does ArrayFlatMap::new() handle overflow correctly?
89 #[test]
90 fn test_array_flat_map_len_overflow() {
91 struct DownwardCounter {
92 remaining: usize,
93 }
94 impl Iterator for DownwardCounter {
95 type Item = usize;
96
97 fn next(&mut self) -> Option<Self::Item> {
98 if self.remaining > 0 {
99 let result = self.remaining;
100 self.remaining -= 1;
101 Some(result)
102 } else {
103 None
104 }
105 }
106
107 fn size_hint(&self) -> (usize, Option<usize>) {
108 (self.remaining, Some(self.remaining))
109 }
110 }
111 impl ExactSizeIterator for DownwardCounter {}
112
113 const MAX: usize = usize::MAX / size_of::<usize>();
114
115 static TEST_CASES: &[(usize, bool)] = &[(MAX, true), (MAX + 1, false)];
116 TEST_CASES.iter().copied().for_each(|(input_len, is_some)| {
117 let inner = DownwardCounter {
118 remaining: input_len,
119 };
120 let mapped = ArrayFlatMap::new(inner, usize::to_be_bytes);
121 assert_eq!(mapped.is_some(), is_some);
122 if let Some(mapped) = mapped {
123 assert_eq!(mapped.len(), input_len * size_of::<usize>());
124 }
125 });
126 }
127}
128