use generational_box::{AnyStorage, Owner, SyncStorage, UnsyncStorage};
use std::{
any::{Any, TypeId},
cell::RefCell,
};
pub fn with_owner<S: AnyStorage, F: FnOnce() -> R, R>(owner: Owner<S>, f: F) -> R {
let old_owner = set_owner(Some(owner));
let result = f();
set_owner(old_owner);
result
}
fn set_owner<S: AnyStorage>(owner: Option<Owner<S>>) -> Option<Owner<S>> {
let id = TypeId::of::<S>();
if id == TypeId::of::<SyncStorage>() {
SYNC_OWNER.with(|cell| {
std::mem::replace(
&mut *cell.borrow_mut(),
owner.map(|owner| {
*(Box::new(owner) as Box<dyn Any>)
.downcast::<Owner<SyncStorage>>()
.unwrap()
}),
)
.map(|owner| *(Box::new(owner) as Box<dyn Any>).downcast().unwrap())
})
} else {
UNSYNC_OWNER.with(|cell| {
std::mem::replace(
&mut *cell.borrow_mut(),
owner.map(|owner| {
*(Box::new(owner) as Box<dyn Any>)
.downcast::<Owner<UnsyncStorage>>()
.unwrap()
}),
)
.map(|owner| *(Box::new(owner) as Box<dyn Any>).downcast().unwrap())
})
}
}
thread_local! {
static SYNC_OWNER: RefCell<Option<Owner<SyncStorage>>> = const { RefCell::new(None) };
static UNSYNC_OWNER: RefCell<Option<Owner<UnsyncStorage>>> = const { RefCell::new(None) };
}
pub fn current_owner<S: AnyStorage>() -> Owner<S> {
let id = TypeId::of::<S>();
let override_owner = if id == TypeId::of::<SyncStorage>() {
SYNC_OWNER.with(|cell| {
let owner = cell.borrow();
owner.clone().map(|owner| {
*(Box::new(owner) as Box<dyn Any>)
.downcast::<Owner<S>>()
.unwrap()
})
})
} else {
UNSYNC_OWNER.with(|cell| {
cell.borrow().clone().map(|owner| {
*(Box::new(owner) as Box<dyn Any>)
.downcast::<Owner<S>>()
.unwrap()
})
})
};
if let Some(owner) = override_owner {
return owner;
}
crate::Runtime::current().current_owner()
}