12
12
13
13
use collections:: BTreeMap ;
14
14
use ptr;
15
+ use sync:: atomic:: { AtomicUsize , ATOMIC_USIZE_INIT , Ordering } ;
15
16
16
17
pub type Key = usize ;
17
18
18
19
type Dtor = unsafe extern fn ( * mut u8 ) ;
19
20
20
- #[ thread_local]
21
- static mut NEXT_KEY : Key = 0 ;
21
+ static NEXT_KEY : AtomicUsize = ATOMIC_USIZE_INIT ;
22
+
23
+ static mut KEYS : * mut BTreeMap < Key , Option < Dtor > > = ptr:: null_mut ( ) ;
22
24
23
25
#[ thread_local]
24
- static mut LOCALS : * mut BTreeMap < Key , ( * mut u8 , Option < Dtor > ) > = ptr:: null_mut ( ) ;
26
+ static mut LOCALS : * mut BTreeMap < Key , * mut u8 > = ptr:: null_mut ( ) ;
27
+
28
+ unsafe fn keys ( ) -> & ' static mut BTreeMap < Key , Option < Dtor > > {
29
+ if KEYS == ptr:: null_mut ( ) {
30
+ KEYS = Box :: into_raw ( Box :: new ( BTreeMap :: new ( ) ) ) ;
31
+ }
32
+ & mut * KEYS
33
+ }
25
34
26
- unsafe fn locals ( ) -> & ' static mut BTreeMap < Key , ( * mut u8 , Option < Dtor > ) > {
35
+ unsafe fn locals ( ) -> & ' static mut BTreeMap < Key , * mut u8 > {
27
36
if LOCALS == ptr:: null_mut ( ) {
28
37
LOCALS = Box :: into_raw ( Box :: new ( BTreeMap :: new ( ) ) ) ;
29
38
}
@@ -32,26 +41,26 @@ unsafe fn locals() -> &'static mut BTreeMap<Key, (*mut u8, Option<Dtor>)> {
32
41
33
42
#[ inline]
34
43
pub unsafe fn create ( dtor : Option < Dtor > ) -> Key {
35
- let key = NEXT_KEY ;
36
- NEXT_KEY += 1 ;
37
- locals ( ) . insert ( key, ( 0 as * mut u8 , dtor) ) ;
44
+ let key = NEXT_KEY . fetch_add ( 1 , Ordering :: SeqCst ) ;
45
+ keys ( ) . insert ( key, dtor) ;
38
46
key
39
47
}
40
48
41
49
#[ inline]
42
- pub unsafe fn set ( key : Key , value : * mut u8 ) {
43
- locals ( ) . get_mut ( & key) . unwrap ( ) . 0 = value;
50
+ pub unsafe fn get ( key : Key ) -> * mut u8 {
51
+ if let Some ( & entry) = locals ( ) . get ( & key) {
52
+ entry
53
+ } else {
54
+ ptr:: null_mut ( )
55
+ }
44
56
}
45
57
46
58
#[ inline]
47
- pub unsafe fn get ( key : Key ) -> * mut u8 {
48
- locals ( ) [ & key ] . 0
59
+ pub unsafe fn set ( key : Key , value : * mut u8 ) {
60
+ locals ( ) . insert ( key , value ) ;
49
61
}
50
62
51
63
#[ inline]
52
64
pub unsafe fn destroy ( key : Key ) {
53
- let ( value, dtor) = locals ( ) . remove ( & key) . unwrap ( ) ;
54
- if let Some ( dtor_fn) = dtor {
55
- dtor_fn ( value) ;
56
- }
65
+ keys ( ) . remove ( & key) ;
57
66
}
0 commit comments