]>
Commit | Line | Data |
---|---|---|
0bf4aa26 | 1 | use smallvec::{Array, SmallVec}; |
dfeec247 | 2 | use std::ptr; |
476ff2be | 3 | |
9fa01778 | 4 | pub trait MapInPlace<T>: Sized { |
dfeec247 XL |
5 | fn map_in_place<F>(&mut self, mut f: F) |
6 | where | |
7 | F: FnMut(T) -> T, | |
8 | { | |
9fa01778 | 9 | self.flat_map_in_place(|e| Some(f(e))) |
92a42be0 SL |
10 | } |
11 | ||
9fa01778 | 12 | fn flat_map_in_place<F, I>(&mut self, f: F) |
dfeec247 XL |
13 | where |
14 | F: FnMut(T) -> I, | |
15 | I: IntoIterator<Item = T>; | |
92a42be0 SL |
16 | } |
17 | ||
9fa01778 XL |
18 | impl<T> MapInPlace<T> for Vec<T> { |
19 | fn flat_map_in_place<F, I>(&mut self, mut f: F) | |
dfeec247 XL |
20 | where |
21 | F: FnMut(T) -> I, | |
22 | I: IntoIterator<Item = T>, | |
92a42be0 SL |
23 | { |
24 | let mut read_i = 0; | |
25 | let mut write_i = 0; | |
26 | unsafe { | |
27 | let mut old_len = self.len(); | |
28 | self.set_len(0); // make sure we just leak elements in case of panic | |
29 | ||
30 | while read_i < old_len { | |
31 | // move the read_i'th item out of the vector and map it | |
32 | // to an iterator | |
5e7ed085 | 33 | let e = ptr::read(self.as_ptr().add(read_i)); |
7cac9316 | 34 | let iter = f(e).into_iter(); |
92a42be0 SL |
35 | read_i += 1; |
36 | ||
7cac9316 | 37 | for e in iter { |
92a42be0 | 38 | if write_i < read_i { |
5e7ed085 | 39 | ptr::write(self.as_mut_ptr().add(write_i), e); |
92a42be0 SL |
40 | write_i += 1; |
41 | } else { | |
42 | // If this is reached we ran out of space | |
43 | // in the middle of the vector. | |
44 | // However, the vector is in a valid state here, | |
45 | // so we just do a somewhat inefficient insert. | |
46 | self.set_len(old_len); | |
47 | self.insert(write_i, e); | |
48 | ||
49 | old_len = self.len(); | |
50 | self.set_len(0); | |
51 | ||
52 | read_i += 1; | |
53 | write_i += 1; | |
54 | } | |
55 | } | |
56 | } | |
57 | ||
58 | // write_i tracks the number of actually written new items. | |
59 | self.set_len(write_i); | |
60 | } | |
92a42be0 SL |
61 | } |
62 | } | |
476ff2be | 63 | |
9fa01778 XL |
64 | impl<T, A: Array<Item = T>> MapInPlace<T> for SmallVec<A> { |
65 | fn flat_map_in_place<F, I>(&mut self, mut f: F) | |
dfeec247 XL |
66 | where |
67 | F: FnMut(T) -> I, | |
68 | I: IntoIterator<Item = T>, | |
476ff2be SL |
69 | { |
70 | let mut read_i = 0; | |
71 | let mut write_i = 0; | |
72 | unsafe { | |
73 | let mut old_len = self.len(); | |
74 | self.set_len(0); // make sure we just leak elements in case of panic | |
75 | ||
76 | while read_i < old_len { | |
77 | // move the read_i'th item out of the vector and map it | |
78 | // to an iterator | |
5e7ed085 | 79 | let e = ptr::read(self.as_ptr().add(read_i)); |
7cac9316 | 80 | let iter = f(e).into_iter(); |
476ff2be SL |
81 | read_i += 1; |
82 | ||
7cac9316 | 83 | for e in iter { |
476ff2be | 84 | if write_i < read_i { |
5e7ed085 | 85 | ptr::write(self.as_mut_ptr().add(write_i), e); |
476ff2be SL |
86 | write_i += 1; |
87 | } else { | |
88 | // If this is reached we ran out of space | |
89 | // in the middle of the vector. | |
90 | // However, the vector is in a valid state here, | |
91 | // so we just do a somewhat inefficient insert. | |
92 | self.set_len(old_len); | |
93 | self.insert(write_i, e); | |
94 | ||
95 | old_len = self.len(); | |
96 | self.set_len(0); | |
97 | ||
98 | read_i += 1; | |
99 | write_i += 1; | |
100 | } | |
101 | } | |
102 | } | |
103 | ||
104 | // write_i tracks the number of actually written new items. | |
105 | self.set_len(write_i); | |
106 | } | |
476ff2be SL |
107 | } |
108 | } |