1
- use std:: ptr;
1
+ use std:: { mem , ptr} ;
2
2
3
3
use smallvec:: { Array , SmallVec } ;
4
4
use thin_vec:: ThinVec ;
@@ -13,39 +13,44 @@ pub trait FlatMapInPlace<T>: Sized {
13
13
// The implementation of this method is syntactically identical for all the
14
14
// different vector types.
15
15
macro_rules! flat_map_in_place {
16
- ( ) => {
16
+ ( $vec : ident $ ( where T : $bound : path ) ? ) => {
17
17
fn flat_map_in_place<F , I >( & mut self , mut f: F )
18
18
where
19
19
F : FnMut ( T ) -> I ,
20
20
I : IntoIterator <Item = T >,
21
21
{
22
+ struct LeakGuard <' a, T $( : $bound) ?>( & ' a mut $vec<T >) ;
23
+
24
+ impl <' a, T $( : $bound) ?> Drop for LeakGuard <' a, T > {
25
+ fn drop( & mut self ) {
26
+ unsafe {
27
+ self . 0 . set_len( 0 ) ; // make sure we just leak elements in case of panic
28
+ }
29
+ }
30
+ }
31
+
32
+ let this = LeakGuard ( self ) ;
33
+
22
34
let mut read_i = 0 ;
23
35
let mut write_i = 0 ;
24
36
unsafe {
25
- let mut old_len = self . len( ) ;
26
- self . set_len( 0 ) ; // make sure we just leak elements in case of panic
27
-
28
- while read_i < old_len {
37
+ while read_i < this. 0 . len( ) {
29
38
// move the read_i'th item out of the vector and map it
30
39
// to an iterator
31
- let e = ptr:: read( self . as_ptr( ) . add( read_i) ) ;
40
+ let e = ptr:: read( this . 0 . as_ptr( ) . add( read_i) ) ;
32
41
let iter = f( e) . into_iter( ) ;
33
42
read_i += 1 ;
34
43
35
44
for e in iter {
36
45
if write_i < read_i {
37
- ptr:: write( self . as_mut_ptr( ) . add( write_i) , e) ;
46
+ ptr:: write( this . 0 . as_mut_ptr( ) . add( write_i) , e) ;
38
47
write_i += 1 ;
39
48
} else {
40
49
// If this is reached we ran out of space
41
50
// in the middle of the vector.
42
51
// However, the vector is in a valid state here,
43
52
// so we just do a somewhat inefficient insert.
44
- self . set_len( old_len) ;
45
- self . insert( write_i, e) ;
46
-
47
- old_len = self . len( ) ;
48
- self . set_len( 0 ) ;
53
+ this. 0 . insert( write_i, e) ;
49
54
50
55
read_i += 1 ;
51
56
write_i += 1 ;
@@ -54,20 +59,23 @@ macro_rules! flat_map_in_place {
54
59
}
55
60
56
61
// write_i tracks the number of actually written new items.
57
- self . set_len( write_i) ;
62
+ this. 0 . set_len( write_i) ;
63
+
64
+ // The ThinVec is in a sane state again. Prevent the LeakGuard from leaking the data.
65
+ mem:: forget( this) ;
58
66
}
59
67
}
60
68
} ;
61
69
}
62
70
63
71
impl < T > FlatMapInPlace < T > for Vec < T > {
64
- flat_map_in_place ! ( ) ;
72
+ flat_map_in_place ! ( Vec ) ;
65
73
}
66
74
67
75
impl < T , A : Array < Item = T > > FlatMapInPlace < T > for SmallVec < A > {
68
- flat_map_in_place ! ( ) ;
76
+ flat_map_in_place ! ( SmallVec where T : Array ) ;
69
77
}
70
78
71
79
impl < T > FlatMapInPlace < T > for ThinVec < T > {
72
- flat_map_in_place ! ( ) ;
80
+ flat_map_in_place ! ( ThinVec ) ;
73
81
}
0 commit comments