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