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

core/iter/adapters/
chain.rs
