core/iter/adapters/
chain.rs

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")]
23#[ferrocene::prevalidated]
24pub struct Chain<A, B> {
25    // These are "fused" with `Option` so we don't need separate state to track which part is
26    // already exhausted, and we may also get niche layout for `None`. We don't use the real `Fuse`
27    // adapter because its specialization for `FusedIterator` unconditionally descends into the
28    // iterator, and that could be expensive to keep revisiting stuff like nested chains. It also
29    // hurts compiler performance to add more iterator layers to `Chain`.
30    //
31    // Only the "first" iterator is actually set `None` when exhausted, depending on whether you
32    // iterate forward or backward. If you mix directions, then both sides may be `None`.
33    a: Option<A>,
34    b: Option<B>,
35}
36impl<A, B> Chain<A, B> {
37    #[ferrocene::prevalidated]
38    pub(in super::super) fn new(a: A, b: B) -> Chain<A, B> {
39        Chain { a: Some(a), b: Some(b) }
40    }
41}
42
43/// Converts the arguments to iterators and links them together, in a chain.
44///
45/// See the documentation of [`Iterator::chain`] for more.
46///
47/// # Examples
48///
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#[stable(feature = "iter_chain", since = "1.91.0")]
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    #[ferrocene::prevalidated]
84    fn next(&mut self) -> Option<A::Item> {
85        and_then_or_clear(&mut self.a, Iterator::next).or_else(|| self.b.as_mut()?.next())
86    }
87
88    #[inline]
89    #[rustc_inherit_overflow_checks]
90    #[ferrocene::prevalidated]
91    fn count(self) -> usize {
92        let a_count = match self.a {
93            Some(a) => a.count(),
94            None => 0,
95        };
96        let b_count = match self.b {
97            Some(b) => b.count(),
98            None => 0,
99        };
100        a_count + b_count
101    }
102
103    #[ferrocene::prevalidated]
104    fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
105    where
106        Self: Sized,
107        F: FnMut(Acc, Self::Item) -> R,
108        R: Try<Output = Acc>,
109    {
110        if let Some(ref mut a) = self.a {
111            acc = a.try_fold(acc, &mut f)?;
112            self.a = None;
113        }
114        if let Some(ref mut b) = self.b {
115            acc = b.try_fold(acc, f)?;
116            // we don't fuse the second iterator
117        }
118        try { acc }
119    }
120
121    #[ferrocene::prevalidated]
122    fn fold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
123    where
124        F: FnMut(Acc, Self::Item) -> Acc,
125    {
126        if let Some(a) = self.a {
127            acc = a.fold(acc, &mut f);
128        }
129        if let Some(b) = self.b {
130            acc = b.fold(acc, f);
131        }
132        acc
133    }
134
135    #[inline]
136    #[ferrocene::prevalidated]
137    fn advance_by(&mut self, mut n: usize) -> Result<(), NonZero<usize>> {
138        if let Some(ref mut a) = self.a {
139            n = match a.advance_by(n) {
140                Ok(()) => return Ok(()),
141                Err(k) => k.get(),
142            };
143            self.a = None;
144        }
145
146        if let Some(ref mut b) = self.b {
147            return b.advance_by(n);
148            // we don't fuse the second iterator
149        }
150
151        NonZero::new(n).map_or(Ok(()), Err)
152    }
153
154    #[inline]
155    #[ferrocene::prevalidated]
156    fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
157        if let Some(ref mut a) = self.a {
158            n = match a.advance_by(n) {
159                Ok(()) => match a.next() {
160                    None => 0,
161                    x => return x,
162                },
163                Err(k) => k.get(),
164            };
165
166            self.a = None;
167        }
168
169        self.b.as_mut()?.nth(n)
170    }
171
172    #[inline]
173    #[ferrocene::prevalidated]
174    fn find<P>(&mut self, mut predicate: P) -> Option<Self::Item>
175    where
176        P: FnMut(&Self::Item) -> bool,
177    {
178        and_then_or_clear(&mut self.a, |a| a.find(&mut predicate))
179            .or_else(|| self.b.as_mut()?.find(predicate))
180    }
181
182    #[inline]
183    #[ferrocene::prevalidated]
184    fn last(self) -> Option<A::Item> {
185        // Must exhaust a before b.
186        let a_last = self.a.and_then(Iterator::last);
187        let b_last = self.b.and_then(Iterator::last);
188        b_last.or(a_last)
189    }
190
191    #[inline]
192    #[ferrocene::prevalidated]
193    fn size_hint(&self) -> (usize, Option<usize>) {
194        match self {
195            Chain { a: Some(a), b: Some(b) } => {
196                let (a_lower, a_upper) = a.size_hint();
197                let (b_lower, b_upper) = b.size_hint();
198
199                let lower = a_lower.saturating_add(b_lower);
200
201                let upper = match (a_upper, b_upper) {
202                    (Some(x), Some(y)) => x.checked_add(y),
203                    _ => None,
204                };
205
206                (lower, upper)
207            }
208            Chain { a: Some(a), b: None } => a.size_hint(),
209            Chain { a: None, b: Some(b) } => b.size_hint(),
210            Chain { a: None, b: None } => (0, Some(0)),
211        }
212    }
213}
214
215#[stable(feature = "rust1", since = "1.0.0")]
216impl<A, B> DoubleEndedIterator for Chain<A, B>
217where
218    A: DoubleEndedIterator,
219    B: DoubleEndedIterator<Item = A::Item>,
220{
221    #[inline]
222    fn next_back(&mut self) -> Option<A::Item> {
223        and_then_or_clear(&mut self.b, |b| b.next_back()).or_else(|| self.a.as_mut()?.next_back())
224    }
225
226    #[inline]
227    fn advance_back_by(&mut self, mut n: usize) -> Result<(), NonZero<usize>> {
228        if let Some(ref mut b) = self.b {
229            n = match b.advance_back_by(n) {
230                Ok(()) => return Ok(()),
231                Err(k) => k.get(),
232            };
233            self.b = None;
234        }
235
236        if let Some(ref mut a) = self.a {
237            return a.advance_back_by(n);
238            // we don't fuse the second iterator
239        }
240
241        NonZero::new(n).map_or(Ok(()), Err)
242    }
243
244    #[inline]
245    fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> {
246        if let Some(ref mut b) = self.b {
247            n = match b.advance_back_by(n) {
248                Ok(()) => match b.next_back() {
249                    None => 0,
250                    x => return x,
251                },
252                Err(k) => k.get(),
253            };
254
255            self.b = None;
256        }
257
258        self.a.as_mut()?.nth_back(n)
259    }
260
261    #[inline]
262    fn rfind<P>(&mut self, mut predicate: P) -> Option<Self::Item>
263    where
264        P: FnMut(&Self::Item) -> bool,
265    {
266        and_then_or_clear(&mut self.b, |b| b.rfind(&mut predicate))
267            .or_else(|| self.a.as_mut()?.rfind(predicate))
268    }
269
270    fn try_rfold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
271    where
272        Self: Sized,
273        F: FnMut(Acc, Self::Item) -> R,
274        R: Try<Output = Acc>,
275    {
276        if let Some(ref mut b) = self.b {
277            acc = b.try_rfold(acc, &mut f)?;
278            self.b = None;
279        }
280        if let Some(ref mut a) = self.a {
281            acc = a.try_rfold(acc, f)?;
282            // we don't fuse the second iterator
283        }
284        try { acc }
285    }
286
287    fn rfold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
288    where
289        F: FnMut(Acc, Self::Item) -> Acc,
290    {
291        if let Some(b) = self.b {
292            acc = b.rfold(acc, &mut f);
293        }
294        if let Some(a) = self.a {
295            acc = a.rfold(acc, f);
296        }
297        acc
298    }
299}
300
301// Note: *both* must be fused to handle double-ended iterators.
302#[stable(feature = "fused", since = "1.26.0")]
303impl<A, B> FusedIterator for Chain<A, B>
304where
305    A: FusedIterator,
306    B: FusedIterator<Item = A::Item>,
307{
308}
309
310#[unstable(feature = "trusted_len", issue = "37572")]
311unsafe impl<A, B> TrustedLen for Chain<A, B>
312where
313    A: TrustedLen,
314    B: TrustedLen<Item = A::Item>,
315{
316}
317
318#[stable(feature = "default_iters", since = "1.70.0")]
319impl<A: Default, B: Default> Default for Chain<A, B> {
320    /// Creates a `Chain` from the default values for `A` and `B`.
321    ///
322    /// ```
323    /// # use core::iter::Chain;
324    /// # use core::slice;
325    /// # use std::collections::{btree_set, BTreeSet};
326    /// # use std::mem;
327    /// struct Foo<'a>(Chain<slice::Iter<'a, u8>, btree_set::Iter<'a, u8>>);
328    ///
329    /// let set = BTreeSet::<u8>::new();
330    /// let slice: &[u8] = &[];
331    /// let mut foo = Foo(slice.iter().chain(set.iter()));
332    ///
333    /// // take requires `Default`
334    /// let _: Chain<_, _> = mem::take(&mut foo.0);
335    /// ```
336    fn default() -> Self {
337        Chain::new(Default::default(), Default::default())
338    }
339}
340
341#[inline]
342#[ferrocene::prevalidated]
343fn and_then_or_clear<T, U>(opt: &mut Option<T>, f: impl FnOnce(&mut T) -> Option<U>) -> Option<U> {
344    let x = f(opt.as_mut()?);
345    if x.is_none() {
346        *opt = None;
347    }
348    x
349}
core/iter/adapters/chain.rs

core/iter/adapters/
chain.rs
