rustc_data_structures/
flat_map_in_place.rs

1use std::{mem, ptr};
2
3use smallvec::SmallVec;
4use thin_vec::ThinVec;
5
6pub trait FlatMapInPlace<T> {
7    /// `f` turns each element into 0..many elements. This function will consume the existing
8    /// elements in a vec-like structure and replace them with any number of new elements — fewer,
9    /// more, or the same number — as efficiently as possible.
10    fn flat_map_in_place<F, I>(&mut self, f: F)
11    where
12        F: FnMut(T) -> I,
13        I: IntoIterator<Item = T>;
14}
15
16// Blanket impl for all vec-like types that impl `FlatMapInPlaceVec`.
17impl<V: FlatMapInPlaceVec> FlatMapInPlace<V::Elem> for V {
18    fn flat_map_in_place<F, I>(&mut self, mut f: F)
19    where
20        F: FnMut(V::Elem) -> I,
21        I: IntoIterator<Item = V::Elem>,
22    {
23        struct LeakGuard<'a, V: FlatMapInPlaceVec>(&'a mut V);
24
25        impl<'a, V: FlatMapInPlaceVec> Drop for LeakGuard<'a, V> {
26            fn drop(&mut self) {
27                unsafe {
28                    // Leak all elements in case of panic.
29                    self.0.set_len(0);
30                }
31            }
32        }
33
34        let guard = LeakGuard(self);
35
36        let mut read_i = 0;
37        let mut write_i = 0;
38        unsafe {
39            while read_i < guard.0.len() {
40                // Move the read_i'th item out of the vector and map it to an iterator.
41                let e = ptr::read(guard.0.as_ptr().add(read_i));
42                let iter = f(e).into_iter();
43                read_i += 1;
44
45                for e in iter {
46                    if write_i < read_i {
47                        ptr::write(guard.0.as_mut_ptr().add(write_i), e);
48                        write_i += 1;
49                    } else {
50                        // If this is reached we ran out of space in the middle of the vector.
51                        // However, the vector is in a valid state here, so we just do a somewhat
52                        // inefficient insert.
53                        guard.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            guard.0.set_len(write_i);
63
64            // `vec` is in a sane state again. Prevent the LeakGuard from leaking the data.
65            mem::forget(guard);
66        }
67    }
68}
69
70/// A vec-like type must implement these operations to support `flat_map_in_place`.
71///
72/// # Safety
73///
74/// The memory safety of the unsafe block in `flat_map_in_place` relies on impls of this trait
75/// implementing all the operations correctly.
76pub unsafe trait FlatMapInPlaceVec {
77    type Elem;
78
79    fn len(&self) -> usize;
80    unsafe fn set_len(&mut self, len: usize);
81    fn as_ptr(&self) -> *const Self::Elem;
82    fn as_mut_ptr(&mut self) -> *mut Self::Elem;
83    fn insert(&mut self, idx: usize, elem: Self::Elem);
84}
85
86unsafe impl<T> FlatMapInPlaceVec for Vec<T> {
87    type Elem = T;
88
89    fn len(&self) -> usize {
90        self.len()
91    }
92
93    unsafe fn set_len(&mut self, len: usize) {
94        unsafe {
95            self.set_len(len);
96        }
97    }
98
99    fn as_ptr(&self) -> *const Self::Elem {
100        self.as_ptr()
101    }
102
103    fn as_mut_ptr(&mut self) -> *mut Self::Elem {
104        self.as_mut_ptr()
105    }
106
107    fn insert(&mut self, idx: usize, elem: Self::Elem) {
108        self.insert(idx, elem);
109    }
110}
111
112unsafe impl<T> FlatMapInPlaceVec for ThinVec<T> {
113    type Elem = T;
114
115    fn len(&self) -> usize {
116        self.len()
117    }
118
119    unsafe fn set_len(&mut self, len: usize) {
120        unsafe {
121            self.set_len(len);
122        }
123    }
124
125    fn as_ptr(&self) -> *const Self::Elem {
126        self.as_slice().as_ptr()
127    }
128
129    fn as_mut_ptr(&mut self) -> *mut Self::Elem {
130        self.as_mut_slice().as_mut_ptr()
131    }
132
133    fn insert(&mut self, idx: usize, elem: Self::Elem) {
134        self.insert(idx, elem);
135    }
136}
137
138unsafe impl<T, const N: usize> FlatMapInPlaceVec for SmallVec<[T; N]> {
139    type Elem = T;
140
141    fn len(&self) -> usize {
142        self.len()
143    }
144
145    unsafe fn set_len(&mut self, len: usize) {
146        unsafe {
147            self.set_len(len);
148        }
149    }
150
151    fn as_ptr(&self) -> *const Self::Elem {
152        self.as_ptr()
153    }
154
155    fn as_mut_ptr(&mut self) -> *mut Self::Elem {
156        self.as_mut_ptr()
157    }
158
159    fn insert(&mut self, idx: usize, elem: Self::Elem) {
160        self.insert(idx, elem);
161    }
162}
rustc_data_structures/flat_map_in_place.rs

rustc_data_structures/
flat_map_in_place.rs
