rustc_data_structures/
flat_map_in_place.rs

1use std::{mem, ptr};
2
3use smallvec::{Array, SmallVec};
4use thin_vec::ThinVec;
5
6pub trait FlatMapInPlace<T>: Sized {
7    fn flat_map_in_place<F, I>(&mut self, f: F)
8    where
9        F: FnMut(T) -> I,
10        I: IntoIterator<Item = T>;
11}
12
13// The implementation of this method is syntactically identical for all the
14// different vector types.
15macro_rules! flat_map_in_place {
16    ($vec:ident $( where T: $bound:path)?) => {
17        fn flat_map_in_place<F, I>(&mut self, mut f: F)
18        where
19            F: FnMut(T) -> I,
20            I: IntoIterator<Item = T>,
21        {
22            struct LeakGuard<'a, T $(: $bound)?>(&'a mut $vec<T>);
23
24            impl<'a, T $(: $bound)?> Drop for LeakGuard<'a, T> {
25                fn drop(&mut self) {
26                    unsafe {
27                        self.0.set_len(0); // make sure we just leak elements in case of panic
28                    }
29                }
30            }
31
32            let this = LeakGuard(self);
33
34            let mut read_i = 0;
35            let mut write_i = 0;
36            unsafe {
37                while read_i < this.0.len() {
38                    // move the read_i'th item out of the vector and map it
39                    // to an iterator
40                    let e = ptr::read(this.0.as_ptr().add(read_i));
41                    let iter = f(e).into_iter();
42                    read_i += 1;
43
44                    for e in iter {
45                        if write_i < read_i {
46                            ptr::write(this.0.as_mut_ptr().add(write_i), e);
47                            write_i += 1;
48                        } else {
49                            // If this is reached we ran out of space
50                            // in the middle of the vector.
51                            // However, the vector is in a valid state here,
52                            // so we just do a somewhat inefficient insert.
53                            this.0.insert(write_i, e);
54
55                            read_i += 1;
56                            write_i += 1;
57                        }
58                    }
59                }
60
61                // write_i tracks the number of actually written new items.
62                this.0.set_len(write_i);
63
64                // The ThinVec is in a sane state again. Prevent the LeakGuard from leaking the data.
65                mem::forget(this);
66            }
67        }
68    };
69}
70
71impl<T> FlatMapInPlace<T> for Vec<T> {
72    flat_map_in_place!(Vec);
73}
74
75impl<T, A: Array<Item = T>> FlatMapInPlace<T> for SmallVec<A> {
76    flat_map_in_place!(SmallVec where T: Array);
77}
78
79impl<T> FlatMapInPlace<T> for ThinVec<T> {
80    flat_map_in_place!(ThinVec);
81}
rustc_data_structures/flat_map_in_place.rs

rustc_data_structures/
flat_map_in_place.rs
