swc_visit/util/move_map.rs
1use std::{iter, ptr};
2
3/// Modifiers vector in-place.
4pub trait MoveMap<T>: Sized {
5 /// Map in place.
6 fn move_map<F>(self, mut f: F) -> Self
7 where
8 F: FnMut(T) -> T,
9 {
10 self.move_flat_map(|e| iter::once(f(e)))
11 }
12
13 /// This will be very slow if you try to extend vector using this method.
14 ///
15 /// This method exists to drop useless nodes. You can return Option to do
16 /// such shortening.
17 fn move_flat_map<F, I>(self, f: F) -> Self
18 where
19 F: FnMut(T) -> I,
20 I: IntoIterator<Item = T>;
21}
22
23impl<T> MoveMap<T> for Vec<T> {
24 /// This reduces binary size.
25 fn move_map<F>(mut self, mut f: F) -> Self
26 where
27 F: FnMut(T) -> T,
28 {
29 unsafe {
30 let old_len = self.len();
31 self.set_len(0); // make sure we just leak elements in case of panic
32
33 for index in 0..old_len {
34 let item_ptr = self.as_mut_ptr().add(index);
35 // move the item out of the vector and map it
36 let item = ptr::read(item_ptr);
37 let item = f(item);
38
39 ptr::write(item_ptr, item);
40 }
41
42 // restore the original length
43 self.set_len(old_len);
44 }
45
46 self
47 }
48
49 fn move_flat_map<F, I>(mut self, mut f: F) -> Self
50 where
51 F: FnMut(T) -> I,
52 I: IntoIterator<Item = T>,
53 {
54 let mut read_i = 0;
55 let mut write_i = 0;
56 unsafe {
57 let mut old_len = self.len();
58 self.set_len(0); // make sure we just leak elements in case of panic
59
60 while read_i < old_len {
61 // move the read_i'th item out of the vector and map it
62 // to an iterator
63 let e = ptr::read(self.as_ptr().add(read_i));
64 let iter = f(e).into_iter();
65 read_i += 1;
66
67 for e in iter {
68 if write_i < read_i {
69 ptr::write(self.as_mut_ptr().add(write_i), e);
70 write_i += 1;
71 } else {
72 // If this is reached we ran out of space
73 // in the middle of the vector.
74 // However, the vector is in a valid state here,
75 // so we just do a somewhat inefficient insert.
76 self.set_len(old_len);
77 self.insert(write_i, e);
78
79 old_len = self.len();
80 self.set_len(0);
81
82 read_i += 1;
83 write_i += 1;
84 }
85 }
86 }
87
88 // write_i tracks the number of actually written new items.
89 self.set_len(write_i);
90 }
91
92 self
93 }
94}