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")]
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) const 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    #[inline]
100    fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
101    where
102        Self: Sized,
103        F: FnMut(Acc, Self::Item) -> R,
104        R: Try<Output = Acc>,
105    {
106        if let Some(ref mut a) = self.a {
107            acc = a.try_fold(acc, &mut f)?;
108            self.a = None;
109        }
110        if let Some(ref mut b) = self.b {
111            acc = b.try_fold(acc, f)?;
112            // we don't fuse the second iterator
113        }
114        try { acc }
115    }
116
117    #[inline]
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    #[inline]
262    fn try_rfold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
263    where
264        Self: Sized,
265        F: FnMut(Acc, Self::Item) -> R,
266        R: Try<Output = Acc>,
267    {
268        if let Some(ref mut b) = self.b {
269            acc = b.try_rfold(acc, &mut f)?;
270            self.b = None;
271        }
272        if let Some(ref mut a) = self.a {
273            acc = a.try_rfold(acc, f)?;
274            // we don't fuse the second iterator
275        }
276        try { acc }
277    }
278
279    #[inline]
280    fn rfold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
281    where
282        F: FnMut(Acc, Self::Item) -> Acc,
283    {
284        if let Some(b) = self.b {
285            acc = b.rfold(acc, &mut f);
286        }
287        if let Some(a) = self.a {
288            acc = a.rfold(acc, f);
289        }
290        acc
291    }
292}
293
294// Note: *both* must be fused to handle double-ended iterators.
295#[stable(feature = "fused", since = "1.26.0")]
296impl<A, B> FusedIterator for Chain<A, B>
297where
298    A: FusedIterator,
299    B: FusedIterator<Item = A::Item>,
300{
301}
302
303#[unstable(feature = "trusted_len", issue = "37572")]
304unsafe impl<A, B> TrustedLen for Chain<A, B>
305where
306    A: TrustedLen,
307    B: TrustedLen<Item = A::Item>,
308{
309}
310
311#[stable(feature = "default_iters", since = "1.70.0")]
312impl<A: Default, B: Default> Default for Chain<A, B> {
313    /// Creates a `Chain` from the default values for `A` and `B`.
314    ///
315    /// ```
316    /// # use core::iter::Chain;
317    /// # use core::slice;
318    /// # use std::collections::{btree_set, BTreeSet};
319    /// # use std::mem;
320    /// struct Foo<'a>(Chain<slice::Iter<'a, u8>, btree_set::Iter<'a, u8>>);
321    ///
322    /// let set = BTreeSet::<u8>::new();
323    /// let slice: &[u8] = &[];
324    /// let mut foo = Foo(slice.iter().chain(set.iter()));
325    ///
326    /// // take requires `Default`
327    /// let _: Chain<_, _> = mem::take(&mut foo.0);
328    /// ```
329    fn default() -> Self {
330        Chain::new(Default::default(), Default::default())
331    }
332}
333
334#[inline]
335fn and_then_or_clear<T, U>(opt: &mut Option<T>, f: impl FnOnce(&mut T) -> Option<U>) -> Option<U> {
336    let x = f(opt.as_mut()?);
337    if x.is_none() {
338        *opt = None;
339    }
340    x
341}
core/iter/adapters/chain.rs

core/iter/adapters/
chain.rs
