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
63 #[test]
64 fn test_array_flat_map() {
65 static TEST_CASES: &[(&[u16], fn(u16) -> [u8; 2], &[u8])] = &[
66 // Empty input
67 (&[], u16::to_be_bytes, &[]),
68 // Non-empty input.
69 (
70 &[0x0102, 0x0304, 0x0506],
71 u16::to_be_bytes,
72 &[1, 2, 3, 4, 5, 6],
73 ),
74 // Test with a different mapping function.
75 (
76 &[0x0102, 0x0304, 0x0506],
77 u16::to_le_bytes,
78 &[2, 1, 4, 3, 6, 5],
79 ),
80 ];
81 TEST_CASES.iter().copied().for_each(|(input, f, expected)| {
82 let mapped = ArrayFlatMap::new(input.iter().copied(), f).unwrap();
83 super::super::test::assert_iterator(mapped, expected);
84 });
85 }
86
87 // Does ArrayFlatMap::new() handle overflow correctly?
88 #[test]
89 fn test_array_flat_map_len_overflow() {
90 struct DownwardCounter {
91 remaining: usize,
92 }
93 impl Iterator for DownwardCounter {
94 type Item = usize;
95
96 fn next(&mut self) -> Option<Self::Item> {
97 if self.remaining > 0 {
98 let result = self.remaining;
99 self.remaining -= 1;
100 Some(result)
101 } else {
102 None
103 }
104 }
105
106 fn size_hint(&self) -> (usize, Option<usize>) {
107 (self.remaining, Some(self.remaining))
108 }
109 }
110 impl ExactSizeIterator for DownwardCounter {}
111
112 const MAX: usize = usize::MAX / core::mem::size_of::<usize>();
113
114 static TEST_CASES: &[(usize, bool)] = &[(MAX, true), (MAX + 1, false)];
115 TEST_CASES.iter().copied().for_each(|(input_len, is_some)| {
116 let inner = DownwardCounter {
117 remaining: input_len,
118 };
119 let mapped = ArrayFlatMap::new(inner, usize::to_be_bytes);
120 assert_eq!(mapped.is_some(), is_some);
121 if let Some(mapped) = mapped {
122 assert_eq!(mapped.len(), input_len * core::mem::size_of::<usize>());
123 }
124 });
125 }
126}
127