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}