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/// #![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")]
66pub fn chain<A, B>(a: A, b: B) -> Chain<A::IntoIter, B::IntoIter>
67where
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")]
75impl<A, B> Iterator for Chain<A, B>
76where
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")]
207impl<A, B> DoubleEndedIterator for Chain<A, B>
208where
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")]
294impl<A, B> FusedIterator for Chain<A, B>
295where
296 A: FusedIterator,
297 B: FusedIterator<Item = A::Item>,
298{
299}
300
301#[unstable(feature = "trusted_len", issue = "37572")]
302unsafe impl<A, B> TrustedLen for Chain<A, B>
303where
304 A: TrustedLen,
305 B: TrustedLen<Item = A::Item>,
306{
307}
308
309#[stable(feature = "default_iters", since = "1.70.0")]
310impl<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]
332fn 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