rustc_data_structures/
flat_map_in_place.rs

1use std::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    () => {
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            let mut read_i = 0;
23            let mut write_i = 0;
24            unsafe {
25                let mut old_len = self.len();
26                self.set_len(0); // make sure we just leak elements in case of panic
27
28                while read_i < old_len {
29                    // move the read_i'th item out of the vector and map it
30                    // to an iterator
31                    let e = ptr::read(self.as_ptr().add(read_i));
32                    let iter = f(e).into_iter();
33                    read_i += 1;
34
35                    for e in iter {
36                        if write_i < read_i {
37                            ptr::write(self.as_mut_ptr().add(write_i), e);
38                            write_i += 1;
39                        } else {
40                            // If this is reached we ran out of space
41                            // in the middle of the vector.
42                            // However, the vector is in a valid state here,
43                            // so we just do a somewhat inefficient insert.
44                            self.set_len(old_len);
45                            self.insert(write_i, e);
46
47                            old_len = self.len();
48                            self.set_len(0);
49
50                            read_i += 1;
51                            write_i += 1;
52                        }
53                    }
54                }
55
56                // write_i tracks the number of actually written new items.
57                self.set_len(write_i);
58            }
59        }
60    };
61}
62
63impl<T> FlatMapInPlace<T> for Vec<T> {
64    flat_map_in_place!();
65}
66
67impl<T, A: Array<Item = T>> FlatMapInPlace<T> for SmallVec<A> {
68    flat_map_in_place!();
69}
70
71impl<T> FlatMapInPlace<T> for ThinVec<T> {
72    flat_map_in_place!();
73}
rustc_data_structures/flat_map_in_place.rs

rustc_data_structures/
flat_map_in_place.rs
