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 [`chain`] or [`Iterator::chain`]. See their |
8 | /// documentation 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 | /// Converts the arguments to iterators and links them together, in a chain. |
42 | /// |
43 | /// See the documentation of [`Iterator::chain`] for more. |
44 | /// |
45 | /// # Examples |
46 | /// |
47 | /// ``` |
48 | /// #![feature(iter_chain)] |
49 | /// |
50 | /// use std::iter::chain; |
51 | /// |
52 | /// let a = [1, 2, 3]; |
53 | /// let b = [4, 5, 6]; |
54 | /// |
55 | /// let mut iter = chain(a, b); |
56 | /// |
57 | /// assert_eq!(iter.next(), Some(1)); |
58 | /// assert_eq!(iter.next(), Some(2)); |
59 | /// assert_eq!(iter.next(), Some(3)); |
60 | /// assert_eq!(iter.next(), Some(4)); |
61 | /// assert_eq!(iter.next(), Some(5)); |
62 | /// assert_eq!(iter.next(), Some(6)); |
63 | /// assert_eq!(iter.next(), None); |
64 | /// ``` |
65 | #[unstable (feature = "iter_chain" , reason = "recently added" , issue = "125964" )] |
66 | pub fn chain<A, B>(a: A, b: B) -> Chain<A::IntoIter, B::IntoIter> |
67 | where |
68 | A: IntoIterator, |
69 | B: IntoIterator<Item = A::Item>, |
70 | { |
71 | Chain::new(a.into_iter(), b.into_iter()) |
72 | } |
73 | |
74 | #[stable (feature = "rust1" , since = "1.0.0" )] |
75 | impl<A, B> Iterator for Chain<A, B> |
76 | where |
77 | A: Iterator, |
78 | B: Iterator<Item = A::Item>, |
79 | { |
80 | type Item = A::Item; |
81 | |
82 | #[inline ] |
83 | fn next(&mut self) -> Option<A::Item> { |
84 | and_then_or_clear(&mut self.a, Iterator::next).or_else(|| self.b.as_mut()?.next()) |
85 | } |
86 | |
87 | #[inline ] |
88 | #[rustc_inherit_overflow_checks ] |
89 | fn count(self) -> usize { |
90 | let a_count = match self.a { |
91 | Some(a) => a.count(), |
92 | None => 0, |
93 | }; |
94 | let b_count = match self.b { |
95 | Some(b) => b.count(), |
96 | None => 0, |
97 | }; |
98 | a_count + b_count |
99 | } |
100 | |
101 | fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R |
102 | where |
103 | Self: Sized, |
104 | F: FnMut(Acc, Self::Item) -> R, |
105 | R: Try<Output = Acc>, |
106 | { |
107 | if let Some(ref mut a) = self.a { |
108 | acc = a.try_fold(acc, &mut f)?; |
109 | self.a = None; |
110 | } |
111 | if let Some(ref mut b) = self.b { |
112 | acc = b.try_fold(acc, f)?; |
113 | // we don't fuse the second iterator |
114 | } |
115 | try { acc } |
116 | } |
117 | |
118 | fn fold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc |
119 | where |
120 | F: FnMut(Acc, Self::Item) -> Acc, |
121 | { |
122 | if let Some(a) = self.a { |
123 | acc = a.fold(acc, &mut f); |
124 | } |
125 | if let Some(b) = self.b { |
126 | acc = b.fold(acc, f); |
127 | } |
128 | acc |
129 | } |
130 | |
131 | #[inline ] |
132 | fn advance_by(&mut self, mut n: usize) -> Result<(), NonZero<usize>> { |
133 | if let Some(ref mut a) = self.a { |
134 | n = match a.advance_by(n) { |
135 | Ok(()) => return Ok(()), |
136 | Err(k) => k.get(), |
137 | }; |
138 | self.a = None; |
139 | } |
140 | |
141 | if let Some(ref mut b) = self.b { |
142 | return b.advance_by(n); |
143 | // we don't fuse the second iterator |
144 | } |
145 | |
146 | NonZero::new(n).map_or(Ok(()), Err) |
147 | } |
148 | |
149 | #[inline ] |
150 | fn nth(&mut self, mut n: usize) -> Option<Self::Item> { |
151 | if let Some(ref mut a) = self.a { |
152 | n = match a.advance_by(n) { |
153 | Ok(()) => match a.next() { |
154 | None => 0, |
155 | x => return x, |
156 | }, |
157 | Err(k) => k.get(), |
158 | }; |
159 | |
160 | self.a = None; |
161 | } |
162 | |
163 | self.b.as_mut()?.nth(n) |
164 | } |
165 | |
166 | #[inline ] |
167 | fn find<P>(&mut self, mut predicate: P) -> Option<Self::Item> |
168 | where |
169 | P: FnMut(&Self::Item) -> bool, |
170 | { |
171 | and_then_or_clear(&mut self.a, |a| a.find(&mut predicate)) |
172 | .or_else(|| self.b.as_mut()?.find(predicate)) |
173 | } |
174 | |
175 | #[inline ] |
176 | fn last(self) -> Option<A::Item> { |
177 | // Must exhaust a before b. |
178 | let a_last = self.a.and_then(Iterator::last); |
179 | let b_last = self.b.and_then(Iterator::last); |
180 | b_last.or(a_last) |
181 | } |
182 | |
183 | #[inline ] |
184 | fn size_hint(&self) -> (usize, Option<usize>) { |
185 | match self { |
186 | Chain { a: Some(a), b: Some(b) } => { |
187 | let (a_lower, a_upper) = a.size_hint(); |
188 | let (b_lower, b_upper) = b.size_hint(); |
189 | |
190 | let lower = a_lower.saturating_add(b_lower); |
191 | |
192 | let upper = match (a_upper, b_upper) { |
193 | (Some(x), Some(y)) => x.checked_add(y), |
194 | _ => None, |
195 | }; |
196 | |
197 | (lower, upper) |
198 | } |
199 | Chain { a: Some(a), b: None } => a.size_hint(), |
200 | Chain { a: None, b: Some(b) } => b.size_hint(), |
201 | Chain { a: None, b: None } => (0, Some(0)), |
202 | } |
203 | } |
204 | } |
205 | |
206 | #[stable (feature = "rust1" , since = "1.0.0" )] |
207 | impl<A, B> DoubleEndedIterator for Chain<A, B> |
208 | where |
209 | A: DoubleEndedIterator, |
210 | B: DoubleEndedIterator<Item = A::Item>, |
211 | { |
212 | #[inline ] |
213 | fn next_back(&mut self) -> Option<A::Item> { |
214 | and_then_or_clear(&mut self.b, |b| b.next_back()).or_else(|| self.a.as_mut()?.next_back()) |
215 | } |
216 | |
217 | #[inline ] |
218 | fn advance_back_by(&mut self, mut n: usize) -> Result<(), NonZero<usize>> { |
219 | if let Some(ref mut b) = self.b { |
220 | n = match b.advance_back_by(n) { |
221 | Ok(()) => return Ok(()), |
222 | Err(k) => k.get(), |
223 | }; |
224 | self.b = None; |
225 | } |
226 | |
227 | if let Some(ref mut a) = self.a { |
228 | return a.advance_back_by(n); |
229 | // we don't fuse the second iterator |
230 | } |
231 | |
232 | NonZero::new(n).map_or(Ok(()), Err) |
233 | } |
234 | |
235 | #[inline ] |
236 | fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> { |
237 | if let Some(ref mut b) = self.b { |
238 | n = match b.advance_back_by(n) { |
239 | Ok(()) => match b.next_back() { |
240 | None => 0, |
241 | x => return x, |
242 | }, |
243 | Err(k) => k.get(), |
244 | }; |
245 | |
246 | self.b = None; |
247 | } |
248 | |
249 | self.a.as_mut()?.nth_back(n) |
250 | } |
251 | |
252 | #[inline ] |
253 | fn rfind<P>(&mut self, mut predicate: P) -> Option<Self::Item> |
254 | where |
255 | P: FnMut(&Self::Item) -> bool, |
256 | { |
257 | and_then_or_clear(&mut self.b, |b| b.rfind(&mut predicate)) |
258 | .or_else(|| self.a.as_mut()?.rfind(predicate)) |
259 | } |
260 | |
261 | fn try_rfold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R |
262 | where |
263 | Self: Sized, |
264 | F: FnMut(Acc, Self::Item) -> R, |
265 | R: Try<Output = Acc>, |
266 | { |
267 | if let Some(ref mut b) = self.b { |
268 | acc = b.try_rfold(acc, &mut f)?; |
269 | self.b = None; |
270 | } |
271 | if let Some(ref mut a) = self.a { |
272 | acc = a.try_rfold(acc, f)?; |
273 | // we don't fuse the second iterator |
274 | } |
275 | try { acc } |
276 | } |
277 | |
278 | fn rfold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc |
279 | where |
280 | F: FnMut(Acc, Self::Item) -> Acc, |
281 | { |
282 | if let Some(b) = self.b { |
283 | acc = b.rfold(acc, &mut f); |
284 | } |
285 | if let Some(a) = self.a { |
286 | acc = a.rfold(acc, f); |
287 | } |
288 | acc |
289 | } |
290 | } |
291 | |
292 | // Note: *both* must be fused to handle double-ended iterators. |
293 | #[stable (feature = "fused" , since = "1.26.0" )] |
294 | impl<A, B> FusedIterator for Chain<A, B> |
295 | where |
296 | A: FusedIterator, |
297 | B: FusedIterator<Item = A::Item>, |
298 | { |
299 | } |
300 | |
301 | #[unstable (feature = "trusted_len" , issue = "37572" )] |
302 | unsafe impl<A, B> TrustedLen for Chain<A, B> |
303 | where |
304 | A: TrustedLen, |
305 | B: TrustedLen<Item = A::Item>, |
306 | { |
307 | } |
308 | |
309 | #[stable (feature = "default_iters" , since = "1.70.0" )] |
310 | impl<A: Default, B: Default> Default for Chain<A, B> { |
311 | /// Creates a `Chain` from the default values for `A` and `B`. |
312 | /// |
313 | /// ``` |
314 | /// # use core::iter::Chain; |
315 | /// # use core::slice; |
316 | /// # use std::collections::{btree_set, BTreeSet}; |
317 | /// # use std::mem; |
318 | /// struct Foo<'a>(Chain<slice::Iter<'a, u8>, btree_set::Iter<'a, u8>>); |
319 | /// |
320 | /// let set = BTreeSet::<u8>::new(); |
321 | /// let slice: &[u8] = &[]; |
322 | /// let mut foo = Foo(slice.iter().chain(set.iter())); |
323 | /// |
324 | /// // take requires `Default` |
325 | /// let _: Chain<_, _> = mem::take(&mut foo.0); |
326 | fn default() -> Self { |
327 | Chain::new(a:Default::default(), b:Default::default()) |
328 | } |
329 | } |
330 | |
331 | #[inline ] |
332 | fn and_then_or_clear<T, U>(opt: &mut Option<T>, f: impl FnOnce(&mut T) -> Option<U>) -> Option<U> { |
333 | let x: Option = f(opt.as_mut()?); |
334 | if x.is_none() { |
335 | *opt = None; |
336 | } |
337 | x |
338 | } |
339 | |