solar_data_structures/
collect.rs

1//! Modified from [`rustc_type_id`](https://github.com/rust-lang/rust/blob/b389b0ab72cb0aa9acf4df0ae0c0e12090782da9/compiler/rustc_type_ir/src/interner.rs#L305).
2
3use smallvec::SmallVec;
4
5/// Implements `f(&iter.collect::<Vec<_>>())` more efficiently.
6///
7/// Imagine you have a function `F: FnOnce(&[T]) -> R`, plus an iterator `iter`
8/// that produces `T` items. You could combine them with
9/// `f(&iter.collect::<Vec<_>>())`, but this requires allocating memory for the
10/// `Vec`.
11///
12/// This trait allows for faster implementations, intended for cases where the
13/// number of items produced by the iterator is small. There is a blanket impl
14/// for `T` items, but there is also a fallible impl for `Result<T, E>` items.
15pub trait CollectAndApply<T, R>: Sized {
16    type Output;
17
18    /// Produce a result of type `Self::Output` from `iter`. The result will
19    /// typically be produced by applying `f` on the elements produced by
20    /// `iter`, though this may not happen in some impls, e.g. if an error
21    /// occurred during iteration.
22    fn collect_and_apply<I, F>(iter: I, f: F) -> Self::Output
23    where
24        I: Iterator<Item = Self>,
25        F: FnOnce(&[T]) -> R;
26}
27
28/// The blanket impl that always collects all elements and applies `f`.
29impl<T, R> CollectAndApply<T, R> for T {
30    type Output = R;
31
32    /// Equivalent to `f(&iter.collect::<Vec<_>>())`.
33    fn collect_and_apply<I, F>(mut iter: I, f: F) -> R
34    where
35        I: Iterator<Item = T>,
36        F: FnOnce(&[T]) -> R,
37    {
38        // This code is hot enough that it's worth specializing for the most
39        // common length lists, to avoid the overhead of `SmallVec` creation.
40        // Lengths 0, 1, and 2 typically account for ~95% of cases. If
41        // `size_hint` is incorrect a panic will occur via an `unwrap` or an
42        // `assert`.
43        match iter.size_hint() {
44            (0, Some(0)) => {
45                assert!(iter.next().is_none());
46                f(&[])
47            }
48            (1, Some(1)) => {
49                let t0 = iter.next().unwrap();
50                assert!(iter.next().is_none());
51                f(&[t0])
52            }
53            (2, Some(2)) => {
54                let t0 = iter.next().unwrap();
55                let t1 = iter.next().unwrap();
56                assert!(iter.next().is_none());
57                f(&[t0, t1])
58            }
59            _ => f(&iter.collect::<SmallVec<[_; 8]>>()),
60        }
61    }
62}
63
64/// A fallible impl that will fail, without calling `f`, if there are any
65/// errors during collection.
66impl<T, R, E> CollectAndApply<T, R> for Result<T, E> {
67    type Output = Result<R, E>;
68
69    /// Equivalent to `Ok(f(&iter.collect::<Result<Vec<_>>>()?))`.
70    fn collect_and_apply<I, F>(mut iter: I, f: F) -> Result<R, E>
71    where
72        I: Iterator<Item = Self>,
73        F: FnOnce(&[T]) -> R,
74    {
75        // This code is hot enough that it's worth specializing for the most
76        // common length lists, to avoid the overhead of `SmallVec` creation.
77        // Lengths 0, 1, and 2 typically account for ~95% of cases. If
78        // `size_hint` is incorrect a panic will occur via an `unwrap` or an
79        // `assert`, unless a failure happens first, in which case the result
80        // will be an error anyway.
81        Ok(match iter.size_hint() {
82            (0, Some(0)) => {
83                assert!(iter.next().is_none());
84                f(&[])
85            }
86            (1, Some(1)) => {
87                let t0 = iter.next().unwrap()?;
88                assert!(iter.next().is_none());
89                f(&[t0])
90            }
91            (2, Some(2)) => {
92                let t0 = iter.next().unwrap()?;
93                let t1 = iter.next().unwrap()?;
94                assert!(iter.next().is_none());
95                f(&[t0, t1])
96            }
97            _ => f(&iter.collect::<Result<SmallVec<[_; 8]>, _>>()?),
98        })
99    }
100}
solar_data_structures/collect.rs

solar_data_structures/
collect.rs
