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