1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
//! Various utilities used for other implementations

use once_cell::sync::OnceCell;

use std::any::Any;
use std::mem::ManuallyDrop;
use std::thread::{self, ThreadId};

use self::list::AppendList;

/// A wrapper for user data, able to store any type, and correctly
/// handling access from a wrong thread
#[derive(Debug)]
pub struct UserData {
    inner: OnceCell<UserDataInner>,
}

#[derive(Debug)]
enum UserDataInner {
    ThreadSafe(Box<dyn Any + Send + Sync + 'static>),
    NonThreadSafe(Box<ManuallyDrop<dyn Any + 'static>>, ThreadId),
}

// UserData itself is always threadsafe, as it only gives access to its
// content if it is send+sync or we are on the right thread
unsafe impl Send for UserData {}
unsafe impl Sync for UserData {}

impl UserData {
    /// Create a new UserData instance
    pub const fn new() -> UserData {
        UserData { inner: OnceCell::new() }
    }

    /// Sets the UserData to a given value
    ///
    /// The provided closure is called to init the UserData,
    /// does nothing is the UserData had already been set.
    pub fn set<T: Any + 'static, F: FnOnce() -> T>(&self, f: F) {
        self.inner.get_or_init(|| {
            UserDataInner::NonThreadSafe(Box::new(ManuallyDrop::new(f())), thread::current().id())
        });
    }

    /// Sets the UserData to a given threadsafe value
    ///
    /// The provided closure is called to init the UserData,
    /// does nothing is the UserData had already been set.
    pub fn set_threadsafe<T: Any + Send + Sync + 'static, F: FnOnce() -> T>(&self, f: F) {
        self.inner.get_or_init(|| UserDataInner::ThreadSafe(Box::new(f())));
    }

    /// Attempt to access the wrapped user data
    ///
    /// Will return `None` if either:
    ///
    /// - The requested type `T` does not match the type used for construction
    /// - This `UserData` has been created using the non-threadsafe variant and access
    ///   is attempted from an other thread than the one it was created on
    pub fn get<T: 'static>(&self) -> Option<&T> {
        match self.inner.get() {
            Some(&UserDataInner::ThreadSafe(ref val)) => <dyn Any>::downcast_ref::<T>(&**val),
            Some(&UserDataInner::NonThreadSafe(ref val, threadid)) => {
                // only give access if we are on the right thread
                if threadid == thread::current().id() {
                    <dyn Any>::downcast_ref::<T>(&***val)
                } else {
                    None
                }
            }
            None => None,
        }
    }
}

impl Drop for UserData {
    fn drop(&mut self) {
        // only drop non-Send user data if we are on the right thread, leak it otherwise
        if let Some(&mut UserDataInner::NonThreadSafe(ref mut val, threadid)) = self.inner.get_mut()
        {
            if threadid == thread::current().id() {
                unsafe {
                    ManuallyDrop::drop(&mut **val);
                }
            }
        }
    }
}

/// A storage able to store several values of `UserData`
/// of different types. It behaves similarly to a `TypeMap`.
#[derive(Debug)]
pub struct UserDataMap {
    list: AppendList<UserData>,
}

impl UserDataMap {
    /// Create a new map
    pub fn new() -> UserDataMap {
        UserDataMap { list: AppendList::new() }
    }

    /// Attempt to access the wrapped user data of a given type
    ///
    /// Will return `None` if no value of type `T` is stored in this `UserDataMap`
    /// and accessible from this thread
    pub fn get<T: 'static>(&self) -> Option<&T> {
        for user_data in &self.list {
            if let Some(val) = user_data.get::<T>() {
                return Some(val);
            }
        }
        None
    }

    /// Insert a value in the map if it is not already there
    ///
    /// This is the non-threadsafe variant, the type you insert don't have to be
    /// threadsafe, but they will not be visible from other threads (even if they are
    /// actually threadsafe).
    ///
    /// If the value does not already exists, the closure is called to create it and
    /// this function returns `true`. If the value already exists, the closure is not
    /// called, and this function returns `false`.
    pub fn insert_if_missing<T: 'static, F: FnOnce() -> T>(&self, init: F) -> bool {
        if self.get::<T>().is_some() {
            return false;
        }
        let data = UserData::new();
        data.set(init);
        self.list.append(data);
        true
    }

    /// Insert a value in the map if it is not already there
    ///
    /// This is the threadsafe variant, the type you insert must be threadsafe and will
    /// be visible from all threads.
    ///
    /// If the value does not already exists, the closure is called to create it and
    /// this function returns `true`. If the value already exists, the closure is not
    /// called, and this function returns `false`.
    pub fn insert_if_missing_threadsafe<T: Send + Sync + 'static, F: FnOnce() -> T>(
        &self,
        init: F,
    ) -> bool {
        if self.get::<T>().is_some() {
            return false;
        }
        let data = UserData::new();
        data.set_threadsafe(init);
        self.list.append(data);
        true
    }
}

impl Default for UserDataMap {
    fn default() -> UserDataMap {
        UserDataMap::new()
    }
}

mod list {
    /*
     * This is a lock-free append-only list, it is used as an implementation
     * detail of the UserDataMap.
     *
     * It was extracted from https://github.com/Diggsey/lockless under MIT license
     * Copyright © Diggory Blake <diggsey@googlemail.com>
     */

    use std::sync::atomic::{AtomicPtr, Ordering};
    use std::{mem, ptr};

    type NodePtr<T> = Option<Box<Node<T>>>;

    #[derive(Debug)]
    struct Node<T> {
        value: T,
        next: AppendList<T>,
    }

    #[derive(Debug)]
    pub struct AppendList<T>(AtomicPtr<Node<T>>);

    impl<T> AppendList<T> {
        fn node_into_raw(ptr: NodePtr<T>) -> *mut Node<T> {
            match ptr {
                Some(b) => Box::into_raw(b),
                None => ptr::null_mut(),
            }
        }
        unsafe fn node_from_raw(ptr: *mut Node<T>) -> NodePtr<T> {
            if ptr.is_null() {
                None
            } else {
                Some(Box::from_raw(ptr))
            }
        }

        fn new_internal(ptr: NodePtr<T>) -> Self {
            AppendList(AtomicPtr::new(Self::node_into_raw(ptr)))
        }

        pub fn new() -> Self {
            Self::new_internal(None)
        }

        pub fn append(&self, value: T) {
            self.append_list(AppendList::new_internal(Some(Box::new(Node {
                value,
                next: AppendList::new(),
            }))));
        }

        unsafe fn append_ptr(&self, p: *mut Node<T>) {
            loop {
                match self.0.compare_exchange_weak(
                    ptr::null_mut(),
                    p,
                    Ordering::AcqRel,
                    Ordering::Acquire,
                ) {
                    Ok(_) => return,
                    Err(head) => {
                        if !head.is_null() {
                            return (*head).next.append_ptr(p);
                        }
                    }
                }
            }
        }

        pub fn append_list(&self, other: AppendList<T>) {
            let p = other.0.load(Ordering::Acquire);
            mem::forget(other);
            unsafe { self.append_ptr(p) };
        }

        pub fn iter(&self) -> AppendListIterator<T> {
            AppendListIterator(&self.0)
        }

        pub fn iter_mut(&mut self) -> AppendListMutIterator<T> {
            AppendListMutIterator(&mut self.0)
        }
    }

    impl<'a, T> IntoIterator for &'a AppendList<T> {
        type Item = &'a T;
        type IntoIter = AppendListIterator<'a, T>;

        fn into_iter(self) -> AppendListIterator<'a, T> {
            self.iter()
        }
    }

    impl<'a, T> IntoIterator for &'a mut AppendList<T> {
        type Item = &'a mut T;
        type IntoIter = AppendListMutIterator<'a, T>;

        fn into_iter(self) -> AppendListMutIterator<'a, T> {
            self.iter_mut()
        }
    }

    impl<T> Drop for AppendList<T> {
        fn drop(&mut self) {
            unsafe { Self::node_from_raw(mem::replace(self.0.get_mut(), ptr::null_mut())) };
        }
    }

    #[derive(Debug)]
    pub struct AppendListIterator<'a, T: 'a>(&'a AtomicPtr<Node<T>>);

    impl<'a, T: 'a> Iterator for AppendListIterator<'a, T> {
        type Item = &'a T;

        fn next(&mut self) -> Option<&'a T> {
            let p = self.0.load(Ordering::Acquire);
            if p.is_null() {
                None
            } else {
                unsafe {
                    self.0 = &(*p).next.0;
                    Some(&(*p).value)
                }
            }
        }
    }

    #[derive(Debug)]
    pub struct AppendListMutIterator<'a, T: 'a>(&'a mut AtomicPtr<Node<T>>);

    impl<'a, T: 'a> Iterator for AppendListMutIterator<'a, T> {
        type Item = &'a mut T;

        fn next(&mut self) -> Option<&'a mut T> {
            let p = self.0.load(Ordering::Acquire);
            if p.is_null() {
                None
            } else {
                unsafe {
                    self.0 = &mut (*p).next.0;
                    Some(&mut (*p).value)
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::UserDataMap;

    #[test]
    fn insert_twice() {
        let map = UserDataMap::new();

        assert_eq!(map.get::<usize>(), None);
        assert!(map.insert_if_missing(|| 42usize));
        assert!(!map.insert_if_missing(|| 43usize));
        assert_eq!(map.get::<usize>(), Some(&42));
    }
}