use once_cell::sync::OnceCell;
use std::any::Any;
use std::mem::ManuallyDrop;
use std::thread::{self, ThreadId};
use self::list::AppendList;
pub struct UserData {
inner: OnceCell<UserDataInner>,
}
enum UserDataInner {
ThreadSafe(Box<dyn Any + Send + Sync + 'static>),
NonThreadSafe(Box<ManuallyDrop<dyn Any + 'static>>, ThreadId),
}
unsafe impl Send for UserData {}
unsafe impl Sync for UserData {}
impl UserData {
pub const fn new() -> UserData {
UserData {
inner: OnceCell::new(),
}
}
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())
});
}
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())));
}
pub fn get<T: 'static>(&self) -> Option<&T> {
match self.inner.get() {
Some(&UserDataInner::ThreadSafe(ref val)) => Any::downcast_ref::<T>(&**val),
Some(&UserDataInner::NonThreadSafe(ref val, threadid)) => {
if threadid == thread::current().id() {
Any::downcast_ref::<T>(&***val)
} else {
None
}
}
None => None,
}
}
}
impl Drop for UserData {
fn drop(&mut self) {
if let Some(&mut UserDataInner::NonThreadSafe(ref mut val, threadid)) = self.inner.get_mut() {
if threadid == thread::current().id() {
unsafe {
ManuallyDrop::drop(&mut **val);
}
}
}
}
}
pub struct UserDataMap {
list: AppendList<UserData>,
}
impl UserDataMap {
pub fn new() -> UserDataMap {
UserDataMap {
list: AppendList::new(),
}
}
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
}
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
}
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 {
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));
}
}