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`.
71pub trait FlatMapInPlaceVec {
72    type Elem;
73
74    fn len(&self) -> usize;
75    unsafe fn set_len(&mut self, len: usize);
76    fn as_ptr(&self) -> *const Self::Elem;
77    fn as_mut_ptr(&mut self) -> *mut Self::Elem;
78    fn insert(&mut self, idx: usize, elem: Self::Elem);
79}
80
81impl<T> FlatMapInPlaceVec for Vec<T> {
82    type Elem = T;
83
84    fn len(&self) -> usize {
85        self.len()
86    }
87
88    unsafe fn set_len(&mut self, len: usize) {
89        unsafe {
90            self.set_len(len);
91        }
92    }
93
94    fn as_ptr(&self) -> *const Self::Elem {
95        self.as_ptr()
96    }
97
98    fn as_mut_ptr(&mut self) -> *mut Self::Elem {
99        self.as_mut_ptr()
100    }
101
102    fn insert(&mut self, idx: usize, elem: Self::Elem) {
103        self.insert(idx, elem);
104    }
105}
106
107impl<T> FlatMapInPlaceVec for ThinVec<T> {
108    type Elem = T;
109
110    fn len(&self) -> usize {
111        self.len()
112    }
113
114    unsafe fn set_len(&mut self, len: usize) {
115        unsafe {
116            self.set_len(len);
117        }
118    }
119
120    fn as_ptr(&self) -> *const Self::Elem {
121        self.as_slice().as_ptr()
122    }
123
124    fn as_mut_ptr(&mut self) -> *mut Self::Elem {
125        self.as_mut_slice().as_mut_ptr()
126    }
127
128    fn insert(&mut self, idx: usize, elem: Self::Elem) {
129        self.insert(idx, elem);
130    }
131}
132
133impl<T, const N: usize> FlatMapInPlaceVec for SmallVec<[T; N]> {
134    type Elem = T;
135
136    fn len(&self) -> usize {
137        self.len()
138    }
139
140    unsafe fn set_len(&mut self, len: usize) {
141        unsafe {
142            self.set_len(len);
143        }
144    }
145
146    fn as_ptr(&self) -> *const Self::Elem {
147        self.as_ptr()
148    }
149
150    fn as_mut_ptr(&mut self) -> *mut Self::Elem {
151        self.as_mut_ptr()
152    }
153
154    fn insert(&mut self, idx: usize, elem: Self::Elem) {
155        self.insert(idx, elem);
156    }
157}
rustc_data_structures/flat_map_in_place.rs

rustc_data_structures/
flat_map_in_place.rs
