1 | use crate::iter::{FusedIterator, TrustedLen}; |
2 | use crate::num::NonZero; |
3 | use crate::ops::Try; |
4 | |
5 | /// An iterator that links two iterators together, in a chain. |
6 | /// |
7 | /// This `struct` is created by [`Iterator::chain`]. See its documentation |
8 | /// for more. |
9 | /// |
10 | /// # Examples |
11 | /// |
12 | /// ``` |
13 | /// use std::iter::Chain; |
14 | /// use std::slice::Iter; |
15 | /// |
16 | /// let a1 = [1, 2, 3]; |
17 | /// let a2 = [4, 5, 6]; |
18 | /// let iter: Chain<Iter<'_, _>, Iter<'_, _>> = a1.iter().chain(a2.iter()); |
19 | /// ``` |
20 | #[derive (Clone, Debug)] |
21 | #[must_use = "iterators are lazy and do nothing unless consumed" ] |
22 | #[stable (feature = "rust1" , since = "1.0.0" )] |
23 | pub struct Chain<A, B> { |
24 | // These are "fused" with `Option` so we don't need separate state to track which part is |
25 | // already exhausted, and we may also get niche layout for `None`. We don't use the real `Fuse` |
26 | // adapter because its specialization for `FusedIterator` unconditionally descends into the |
27 | // iterator, and that could be expensive to keep revisiting stuff like nested chains. It also |
28 | // hurts compiler performance to add more iterator layers to `Chain`. |
29 | // |
30 | // Only the "first" iterator is actually set `None` when exhausted, depending on whether you |
31 | // iterate forward or backward. If you mix directions, then both sides may be `None`. |
32 | a: Option<A>, |
33 | b: Option<B>, |
34 | } |
35 | impl<A, B> Chain<A, B> { |
36 | pub(in super::super) fn new(a: A, b: B) -> Chain<A, B> { |
37 | Chain { a: Some(a), b: Some(b) } |
38 | } |
39 | } |
40 | |
41 | #[stable (feature = "rust1" , since = "1.0.0" )] |
42 | impl<A, B> Iterator for Chain<A, B> |
43 | where |
44 | A: Iterator, |
45 | B: Iterator<Item = A::Item>, |
46 | { |
47 | type Item = A::Item; |
48 | |
49 | #[inline ] |
50 | fn next(&mut self) -> Option<A::Item> { |
51 | and_then_or_clear(&mut self.a, Iterator::next).or_else(|| self.b.as_mut()?.next()) |
52 | } |
53 | |
54 | #[inline ] |
55 | #[rustc_inherit_overflow_checks ] |
56 | fn count(self) -> usize { |
57 | let a_count = match self.a { |
58 | Some(a) => a.count(), |
59 | None => 0, |
60 | }; |
61 | let b_count = match self.b { |
62 | Some(b) => b.count(), |
63 | None => 0, |
64 | }; |
65 | a_count + b_count |
66 | } |
67 | |
68 | fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R |
69 | where |
70 | Self: Sized, |
71 | F: FnMut(Acc, Self::Item) -> R, |
72 | R: Try<Output = Acc>, |
73 | { |
74 | if let Some(ref mut a) = self.a { |
75 | acc = a.try_fold(acc, &mut f)?; |
76 | self.a = None; |
77 | } |
78 | if let Some(ref mut b) = self.b { |
79 | acc = b.try_fold(acc, f)?; |
80 | // we don't fuse the second iterator |
81 | } |
82 | try { acc } |
83 | } |
84 | |
85 | fn fold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc |
86 | where |
87 | F: FnMut(Acc, Self::Item) -> Acc, |
88 | { |
89 | if let Some(a) = self.a { |
90 | acc = a.fold(acc, &mut f); |
91 | } |
92 | if let Some(b) = self.b { |
93 | acc = b.fold(acc, f); |
94 | } |
95 | acc |
96 | } |
97 | |
98 | #[inline ] |
99 | fn advance_by(&mut self, mut n: usize) -> Result<(), NonZero<usize>> { |
100 | if let Some(ref mut a) = self.a { |
101 | n = match a.advance_by(n) { |
102 | Ok(()) => return Ok(()), |
103 | Err(k) => k.get(), |
104 | }; |
105 | self.a = None; |
106 | } |
107 | |
108 | if let Some(ref mut b) = self.b { |
109 | return b.advance_by(n); |
110 | // we don't fuse the second iterator |
111 | } |
112 | |
113 | NonZero::new(n).map_or(Ok(()), Err) |
114 | } |
115 | |
116 | #[inline ] |
117 | fn nth(&mut self, mut n: usize) -> Option<Self::Item> { |
118 | if let Some(ref mut a) = self.a { |
119 | n = match a.advance_by(n) { |
120 | Ok(()) => match a.next() { |
121 | None => 0, |
122 | x => return x, |
123 | }, |
124 | Err(k) => k.get(), |
125 | }; |
126 | |
127 | self.a = None; |
128 | } |
129 | |
130 | self.b.as_mut()?.nth(n) |
131 | } |
132 | |
133 | #[inline ] |
134 | fn find<P>(&mut self, mut predicate: P) -> Option<Self::Item> |
135 | where |
136 | P: FnMut(&Self::Item) -> bool, |
137 | { |
138 | and_then_or_clear(&mut self.a, |a| a.find(&mut predicate)) |
139 | .or_else(|| self.b.as_mut()?.find(predicate)) |
140 | } |
141 | |
142 | #[inline ] |
143 | fn last(self) -> Option<A::Item> { |
144 | // Must exhaust a before b. |
145 | let a_last = self.a.and_then(Iterator::last); |
146 | let b_last = self.b.and_then(Iterator::last); |
147 | b_last.or(a_last) |
148 | } |
149 | |
150 | #[inline ] |
151 | fn size_hint(&self) -> (usize, Option<usize>) { |
152 | match self { |
153 | Chain { a: Some(a), b: Some(b) } => { |
154 | let (a_lower, a_upper) = a.size_hint(); |
155 | let (b_lower, b_upper) = b.size_hint(); |
156 | |
157 | let lower = a_lower.saturating_add(b_lower); |
158 | |
159 | let upper = match (a_upper, b_upper) { |
160 | (Some(x), Some(y)) => x.checked_add(y), |
161 | _ => None, |
162 | }; |
163 | |
164 | (lower, upper) |
165 | } |
166 | Chain { a: Some(a), b: None } => a.size_hint(), |
167 | Chain { a: None, b: Some(b) } => b.size_hint(), |
168 | Chain { a: None, b: None } => (0, Some(0)), |
169 | } |
170 | } |
171 | } |
172 | |
173 | #[stable (feature = "rust1" , since = "1.0.0" )] |
174 | impl<A, B> DoubleEndedIterator for Chain<A, B> |
175 | where |
176 | A: DoubleEndedIterator, |
177 | B: DoubleEndedIterator<Item = A::Item>, |
178 | { |
179 | #[inline ] |
180 | fn next_back(&mut self) -> Option<A::Item> { |
181 | and_then_or_clear(&mut self.b, |b| b.next_back()).or_else(|| self.a.as_mut()?.next_back()) |
182 | } |
183 | |
184 | #[inline ] |
185 | fn advance_back_by(&mut self, mut n: usize) -> Result<(), NonZero<usize>> { |
186 | if let Some(ref mut b) = self.b { |
187 | n = match b.advance_back_by(n) { |
188 | Ok(()) => return Ok(()), |
189 | Err(k) => k.get(), |
190 | }; |
191 | self.b = None; |
192 | } |
193 | |
194 | if let Some(ref mut a) = self.a { |
195 | return a.advance_back_by(n); |
196 | // we don't fuse the second iterator |
197 | } |
198 | |
199 | NonZero::new(n).map_or(Ok(()), Err) |
200 | } |
201 | |
202 | #[inline ] |
203 | fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> { |
204 | if let Some(ref mut b) = self.b { |
205 | n = match b.advance_back_by(n) { |
206 | Ok(()) => match b.next_back() { |
207 | None => 0, |
208 | x => return x, |
209 | }, |
210 | Err(k) => k.get(), |
211 | }; |
212 | |
213 | self.b = None; |
214 | } |
215 | |
216 | self.a.as_mut()?.nth_back(n) |
217 | } |
218 | |
219 | #[inline ] |
220 | fn rfind<P>(&mut self, mut predicate: P) -> Option<Self::Item> |
221 | where |
222 | P: FnMut(&Self::Item) -> bool, |
223 | { |
224 | and_then_or_clear(&mut self.b, |b| b.rfind(&mut predicate)) |
225 | .or_else(|| self.a.as_mut()?.rfind(predicate)) |
226 | } |
227 | |
228 | fn try_rfold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R |
229 | where |
230 | Self: Sized, |
231 | F: FnMut(Acc, Self::Item) -> R, |
232 | R: Try<Output = Acc>, |
233 | { |
234 | if let Some(ref mut b) = self.b { |
235 | acc = b.try_rfold(acc, &mut f)?; |
236 | self.b = None; |
237 | } |
238 | if let Some(ref mut a) = self.a { |
239 | acc = a.try_rfold(acc, f)?; |
240 | // we don't fuse the second iterator |
241 | } |
242 | try { acc } |
243 | } |
244 | |
245 | fn rfold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc |
246 | where |
247 | F: FnMut(Acc, Self::Item) -> Acc, |
248 | { |
249 | if let Some(b) = self.b { |
250 | acc = b.rfold(acc, &mut f); |
251 | } |
252 | if let Some(a) = self.a { |
253 | acc = a.rfold(acc, f); |
254 | } |
255 | acc |
256 | } |
257 | } |
258 | |
259 | // Note: *both* must be fused to handle double-ended iterators. |
260 | #[stable (feature = "fused" , since = "1.26.0" )] |
261 | impl<A, B> FusedIterator for Chain<A, B> |
262 | where |
263 | A: FusedIterator, |
264 | B: FusedIterator<Item = A::Item>, |
265 | { |
266 | } |
267 | |
268 | #[unstable (feature = "trusted_len" , issue = "37572" )] |
269 | unsafe impl<A, B> TrustedLen for Chain<A, B> |
270 | where |
271 | A: TrustedLen, |
272 | B: TrustedLen<Item = A::Item>, |
273 | { |
274 | } |
275 | |
276 | #[stable (feature = "default_iters" , since = "1.70.0" )] |
277 | impl<A: Default, B: Default> Default for Chain<A, B> { |
278 | /// Creates a `Chain` from the default values for `A` and `B`. |
279 | /// |
280 | /// ``` |
281 | /// # use core::iter::Chain; |
282 | /// # use core::slice; |
283 | /// # use std::collections::{btree_set, BTreeSet}; |
284 | /// # use std::mem; |
285 | /// struct Foo<'a>(Chain<slice::Iter<'a, u8>, btree_set::Iter<'a, u8>>); |
286 | /// |
287 | /// let set = BTreeSet::<u8>::new(); |
288 | /// let slice: &[u8] = &[]; |
289 | /// let mut foo = Foo(slice.iter().chain(set.iter())); |
290 | /// |
291 | /// // take requires `Default` |
292 | /// let _: Chain<_, _> = mem::take(&mut foo.0); |
293 | fn default() -> Self { |
294 | Chain::new(a:Default::default(), b:Default::default()) |
295 | } |
296 | } |
297 | |
298 | #[inline ] |
299 | fn and_then_or_clear<T, U>(opt: &mut Option<T>, f: impl FnOnce(&mut T) -> Option<U>) -> Option<U> { |
300 | let x: Option = f(opt.as_mut()?); |
301 | if x.is_none() { |
302 | *opt = None; |
303 | } |
304 | x |
305 | } |
306 | |