1use crate::iter::{FusedIterator, TrustedLen};
2use crate::num::NonZero;
3use 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")]
23pub 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}
35impl<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")]
42impl<A, B> Iterator for Chain<A, B>
43where
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")]
174impl<A, B> DoubleEndedIterator for Chain<A, B>
175where
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")]
261impl<A, B> FusedIterator for Chain<A, B>
262where
263 A: FusedIterator,
264 B: FusedIterator<Item = A::Item>,
265{
266}
267
268#[unstable(feature = "trusted_len", issue = "37572")]
269unsafe impl<A, B> TrustedLen for Chain<A, B>
270where
271 A: TrustedLen,
272 B: TrustedLen<Item = A::Item>,
273{
274}
275
276#[stable(feature = "default_iters", since = "1.70.0")]
277impl<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]
299fn 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