core/iter/adapters/
chain.rs

1#[cfg(not(feature = "ferrocene_subset"))]
2use crate::iter::{FusedIterator, TrustedLen};
3use crate::num::NonZero;
4use crate::ops::Try;
5
6/// An iterator that links two iterators together, in a chain.
7///
8/// This `struct` is created by [`chain`] or [`Iterator::chain`]. See their
9/// documentation for more.
10///
11/// # Examples
12///
13/// ```
14/// use std::iter::Chain;
15/// use std::slice::Iter;
16///
17/// let a1 = [1, 2, 3];
18/// let a2 = [4, 5, 6];
19/// let iter: Chain<Iter<'_, _>, Iter<'_, _>> = a1.iter().chain(a2.iter());
20/// ```
21#[cfg_attr(not(feature = "ferrocene_subset"), derive(Clone, Debug))]
22#[must_use = "iterators are lazy and do nothing unless consumed"]
23#[stable(feature = "rust1", since = "1.0.0")]
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    pub(in super::super) fn new(a: A, b: B) -> Chain<A, B> {
38        Chain { a: Some(a), b: Some(b) }
39    }
40}
41
42/// Converts the arguments to iterators and links them together, in a chain.
43///
44/// See the documentation of [`Iterator::chain`] for more.
45///
46/// # Examples
47///
48/// ```
49/// use std::iter::chain;
50///
51/// let a = [1, 2, 3];
52/// let b = [4, 5, 6];
53///
54/// let mut iter = chain(a, b);
55///
56/// assert_eq!(iter.next(), Some(1));
57/// assert_eq!(iter.next(), Some(2));
58/// assert_eq!(iter.next(), Some(3));
59/// assert_eq!(iter.next(), Some(4));
60/// assert_eq!(iter.next(), Some(5));
61/// assert_eq!(iter.next(), Some(6));
62/// assert_eq!(iter.next(), None);
63/// ```
64#[cfg(not(feature = "ferrocene_subset"))]
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    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#[cfg(not(feature = "ferrocene_subset"))]
207#[stable(feature = "rust1", since = "1.0.0")]
208impl<A, B> DoubleEndedIterator for Chain<A, B>
209where
210    A: DoubleEndedIterator,
211    B: DoubleEndedIterator<Item = A::Item>,
212{
213    #[inline]
214    fn next_back(&mut self) -> Option<A::Item> {
215        and_then_or_clear(&mut self.b, |b| b.next_back()).or_else(|| self.a.as_mut()?.next_back())
216    }
217
218    #[inline]
219    fn advance_back_by(&mut self, mut n: usize) -> Result<(), NonZero<usize>> {
220        if let Some(ref mut b) = self.b {
221            n = match b.advance_back_by(n) {
222                Ok(()) => return Ok(()),
223                Err(k) => k.get(),
224            };
225            self.b = None;
226        }
227
228        if let Some(ref mut a) = self.a {
229            return a.advance_back_by(n);
230            // we don't fuse the second iterator
231        }
232
233        NonZero::new(n).map_or(Ok(()), Err)
234    }
235
236    #[inline]
237    fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> {
238        if let Some(ref mut b) = self.b {
239            n = match b.advance_back_by(n) {
240                Ok(()) => match b.next_back() {
241                    None => 0,
242                    x => return x,
243                },
244                Err(k) => k.get(),
245            };
246
247            self.b = None;
248        }
249
250        self.a.as_mut()?.nth_back(n)
251    }
252
253    #[inline]
254    fn rfind<P>(&mut self, mut predicate: P) -> Option<Self::Item>
255    where
256        P: FnMut(&Self::Item) -> bool,
257    {
258        and_then_or_clear(&mut self.b, |b| b.rfind(&mut predicate))
259            .or_else(|| self.a.as_mut()?.rfind(predicate))
260    }
261
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    fn rfold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
280    where
281        F: FnMut(Acc, Self::Item) -> Acc,
282    {
283        if let Some(b) = self.b {
284            acc = b.rfold(acc, &mut f);
285        }
286        if let Some(a) = self.a {
287            acc = a.rfold(acc, f);
288        }
289        acc
290    }
291}
292
293// Note: *both* must be fused to handle double-ended iterators.
294#[cfg(not(feature = "ferrocene_subset"))]
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#[cfg(not(feature = "ferrocene_subset"))]
304#[unstable(feature = "trusted_len", issue = "37572")]
305unsafe impl<A, B> TrustedLen for Chain<A, B>
306where
307    A: TrustedLen,
308    B: TrustedLen<Item = A::Item>,
309{
310}
311
312#[cfg(not(feature = "ferrocene_subset"))]
313#[stable(feature = "default_iters", since = "1.70.0")]
314impl<A: Default, B: Default> Default for Chain<A, B> {
315    /// Creates a `Chain` from the default values for `A` and `B`.
316    ///
317    /// ```
318    /// # use core::iter::Chain;
319    /// # use core::slice;
320    /// # use std::collections::{btree_set, BTreeSet};
321    /// # use std::mem;
322    /// struct Foo<'a>(Chain<slice::Iter<'a, u8>, btree_set::Iter<'a, u8>>);
323    ///
324    /// let set = BTreeSet::<u8>::new();
325    /// let slice: &[u8] = &[];
326    /// let mut foo = Foo(slice.iter().chain(set.iter()));
327    ///
328    /// // take requires `Default`
329    /// let _: Chain<_, _> = mem::take(&mut foo.0);
330    /// ```
331    fn default() -> Self {
332        Chain::new(Default::default(), Default::default())
333    }
334}
335
336#[inline]
337fn and_then_or_clear<T, U>(opt: &mut Option<T>, f: impl FnOnce(&mut T) -> Option<U>) -> Option<U> {
338    let x = f(opt.as_mut()?);
339    if x.is_none() {
340        *opt = None;
341    }
342    x
343}
core/iter/adapters/chain.rs

core/iter/adapters/
chain.rs
