use helpers::*;
use std::{
any::Any,
cell::UnsafeCell,
marker::PhantomData,
ops::{Index, IndexMut},
};
pub struct CtxMap<S: Schema> {
schema: PhantomData<S>,
ptrs: Vec<Option<*const dyn Any>>,
values: UnsafeCell<Vec<Option<Box<dyn Any>>>>,
}
impl<S: Schema> CtxMap<S> {
pub fn new() -> Self {
Self {
schema: PhantomData,
values: UnsafeCell::new(Vec::new()),
ptrs: Vec::new(),
}
}
pub fn with<T: ?Sized + 'static, U>(
&mut self,
key: &'static Key<S, T>,
value: &T,
f: impl FnOnce(&mut CtxMapView<S>) -> U,
) -> U {
self.view().with(key, value, f)
}
pub fn with_mut<T: ?Sized + 'static, U, const MUT: bool>(
&mut self,
key: &'static Key<S, T, MUT>,
value: &mut T,
f: impl FnOnce(&mut CtxMapView<S>) -> U,
) -> U {
self.view().with_mut(key, value, f)
}
pub fn view(&mut self) -> CtxMapView<S> {
CtxMapView(self)
}
pub fn get<T: ?Sized, const MUT: bool>(&self, key: &'static Key<S, T, MUT>) -> Option<&T> {
let index = key.index;
unsafe {
if let Some(Some(p)) = self.ptrs.get(index) {
if let Some(p) = <dyn Any>::downcast_ref::<*const T>(&**p) {
Some(&**p)
} else if let Some(p) = <dyn Any>::downcast_ref::<*mut T>(&**p) {
Some(&**p)
} else {
unreachable!()
}
} else {
let data = key.data.as_ref()?.as_ref();
loop {
if let Some(Some(value)) = (*self.values.get()).get(index) {
let p: *const dyn Any = value.as_ref();
return Some(data.get(&*p));
}
self.init_value(index, data);
}
}
}
}
pub fn get_mut<T: ?Sized>(&mut self, key: &'static KeyMut<S, T>) -> Option<&mut T> {
let index = key.index;
unsafe {
if let Some(Some(p)) = self.ptrs.get(index) {
Some(&mut **<dyn Any>::downcast_ref::<*mut T>(&**p).unwrap())
} else {
let data = key.data.as_ref()?.as_ref();
loop {
if let Some(Some(value)) = (*self.values.get()).get_mut(index) {
let p: *mut dyn Any = value.as_mut();
return Some(data.get_mut(&mut *p));
}
self.init_value(index, data);
}
}
}
}
unsafe fn init_value<T: ?Sized>(&self, index: usize, data: &dyn KeyData<T>) {
let init = data.init();
let values = &mut *self.values.get();
if values.len() <= index {
values.resize_with(index + 1, || None);
}
values[index] = Some(init);
}
}
impl<S: Schema> Default for CtxMap<S> {
fn default() -> Self {
Self::new()
}
}
impl<S, T, const MUT: bool> Index<&'static Key<S, T, MUT>> for CtxMap<S>
where
S: Schema,
T: ?Sized + 'static,
{
type Output = T;
fn index(&self, index: &'static Key<S, T, MUT>) -> &Self::Output {
self.get(index).expect("no entry found for key")
}
}
impl<S, T> IndexMut<&'static KeyMut<S, T>> for CtxMap<S>
where
S: Schema,
T: ?Sized + 'static,
{
fn index_mut(&mut self, index: &'static KeyMut<S, T>) -> &mut Self::Output {
self.get_mut(index).expect("no entry found for key")
}
}
pub struct CtxMapView<'a, S: Schema>(&'a mut CtxMap<S>);
impl<'a, S: Schema> CtxMapView<'a, S> {
pub fn with<T: ?Sized + 'static, U>(
&mut self,
key: &'static Key<S, T>,
value: &T,
f: impl FnOnce(&mut CtxMapView<S>) -> U,
) -> U {
let ptr: *const T = value;
self.with_impl(key, ptr, f)
}
pub fn with_mut<T: ?Sized + 'static, U, const MUT: bool>(
&mut self,
key: &'static Key<S, T, MUT>,
value: &mut T,
f: impl FnOnce(&mut CtxMapView<S>) -> U,
) -> U {
let ptr: *mut T = value;
self.with_impl(key, ptr, f)
}
fn with_impl<T: ?Sized + 'static, U, P: 'static, const MUT: bool>(
&mut self,
key: &'static Key<S, T, MUT>,
ptr: P,
f: impl FnOnce(&mut CtxMapView<S>) -> U,
) -> U {
let index = key.index;
if self.0.ptrs.len() <= index {
self.0.ptrs.resize_with(index + 1, || None);
}
let old = self.0.ptrs[index];
self.0.ptrs[index] = Some(&ptr);
let retval = f(self);
self.0.ptrs[index] = old;
retval
}
pub fn view(&mut self) -> CtxMapView<S> {
CtxMapView(self.0)
}
pub fn get<T: ?Sized, const MUT: bool>(&self, key: &'static Key<S, T, MUT>) -> Option<&T> {
self.0.get(key)
}
pub fn get_mut<T: ?Sized>(&mut self, key: &'static KeyMut<S, T>) -> Option<&mut T> {
self.0.get_mut(key)
}
}
impl<'a, S, T, const MUT: bool> Index<&'static Key<S, T, MUT>> for CtxMapView<'a, S>
where
S: Schema,
T: ?Sized + 'static,
{
type Output = T;
fn index(&self, index: &'static Key<S, T, MUT>) -> &Self::Output {
&self.0[index]
}
}
impl<'a, S, T> IndexMut<&'static KeyMut<S, T>> for CtxMapView<'a, S>
where
S: Schema,
T: ?Sized + 'static,
{
fn index_mut(&mut self, index: &'static KeyMut<S, T>) -> &mut Self::Output {
&mut self.0[index]
}
}
pub struct Key<S: Schema, T: ?Sized + 'static, const MUT: bool = false> {
schema: PhantomData<S>,
index: usize,
data: Option<Box<dyn KeyData<T>>>,
}
pub type KeyMut<S, T> = Key<S, T, true>;
trait KeyData<T: ?Sized>: Send + Sync {
fn get<'a>(&self, value: &'a dyn Any) -> &'a T;
fn get_mut<'a>(&self, value: &'a mut dyn Any) -> &'a mut T;
fn init(&self) -> Box<dyn Any>;
}
pub trait Schema: 'static + Sized {
fn data() -> &'static SchemaData;
fn key<T: ?Sized>() -> Key<Self, T> {
Key {
schema: PhantomData,
index: Self::data().push_key(),
data: None,
}
}
fn key_mut<T: ?Sized>() -> KeyMut<Self, T> {
Key {
schema: PhantomData,
index: Self::data().push_key(),
data: None,
}
}
fn key_with_default<Init, ToRef, V, T>(init: Init, to_ref: ToRef) -> Key<Self, T>
where
Init: Send + Sync + Fn() -> V + 'static,
ToRef: Send + Sync + Fn(&V) -> &T + 'static,
V: 'static,
T: ?Sized,
{
fn to_mut_unreachable<V, T: ?Sized>(_: &mut V) -> &mut T {
unreachable!()
}
Key {
schema: PhantomData,
index: Self::data().push_key(),
data: Some(Box::new(KeyDataValue {
init,
to_ref,
to_mut: to_mut_unreachable,
})),
}
}
fn key_mut_with_default<Init, ToRef, ToMut, V, T>(
init: Init,
to_ref: ToRef,
to_mut: ToMut,
) -> KeyMut<Self, T>
where
Init: Send + Sync + Fn() -> V + 'static,
ToRef: Send + Sync + Fn(&V) -> &T + 'static,
ToMut: Send + Sync + Fn(&mut V) -> &mut T + 'static,
V: 'static,
T: ?Sized,
{
Key {
schema: PhantomData,
index: Self::data().push_key(),
data: Some(Box::new(KeyDataValue {
init,
to_ref,
to_mut,
})),
}
}
}
struct KeyDataValue<Init, ToRef, ToMut> {
init: Init,
to_ref: ToRef,
to_mut: ToMut,
}
impl<Init, ToRef, ToMut, V, T> KeyData<T> for KeyDataValue<Init, ToRef, ToMut>
where
Init: Send + Sync + Fn() -> V,
ToRef: Send + Sync + Fn(&V) -> &T,
ToMut: Send + Sync + Fn(&mut V) -> &mut T,
V: 'static,
T: ?Sized,
{
fn get<'a>(&self, value: &'a dyn Any) -> &'a T {
(self.to_ref)(<dyn Any>::downcast_ref::<V>(value).unwrap())
}
fn get_mut<'a>(&self, value: &'a mut dyn Any) -> &'a mut T {
(self.to_mut)(<dyn Any>::downcast_mut::<V>(value).unwrap())
}
fn init(&self) -> Box<dyn Any> {
Box::new((self.init)())
}
}
#[doc(hidden)]
pub mod helpers {
use crate::Schema;
pub use once_cell::sync::Lazy;
use std::{
ops::Deref,
sync::atomic::{AtomicUsize, Ordering},
};
pub struct SchemaData {
next: AtomicUsize,
}
impl SchemaData {
pub const fn new() -> Self {
SchemaData {
next: AtomicUsize::new(0),
}
}
pub(crate) fn push_key(&self) -> usize {
self.next.fetch_add(1, Ordering::SeqCst)
}
}
pub struct Key<S: Schema, T: ?Sized + 'static, const MUT: bool = false>(
Lazy<crate::Key<S, T, MUT>>,
);
pub struct KeyMut<S: Schema, T: ?Sized + 'static>(Lazy<crate::KeyMut<S, T>>);
impl<S: Schema, T: ?Sized + 'static> Key<S, T> {
pub const fn new(f: fn() -> crate::Key<S, T>) -> Self {
Self(Lazy::new(f))
}
}
impl<S: Schema, T: ?Sized + 'static> Deref for Key<S, T> {
type Target = crate::Key<S, T>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl<S: Schema, T: ?Sized + 'static> KeyMut<S, T> {
pub const fn new(f: fn() -> crate::KeyMut<S, T>) -> Self {
Self(Lazy::new(f))
}
}
impl<S: Schema, T: ?Sized + 'static> Deref for KeyMut<S, T> {
type Target = crate::KeyMut<S, T>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
}
#[macro_export]
macro_rules! schema {
($vis:vis $id:ident) => {
$vis struct $id;
impl $crate::Schema for $id {
fn data() -> &'static $crate::helpers::SchemaData {
static DATA: $crate::helpers::SchemaData = $crate::helpers::SchemaData::new();
&DATA
}
}
};
}
#[macro_export]
macro_rules! key {
($schema:ty { }) => { };
($schema:ty { $vis:vis $id:ident: $type:ty }) => {
$vis static $id: $crate::helpers::Key<$schema, $type> =
$crate::helpers::Key::new(|| <$schema as $crate::Schema>::key());
};
($schema:ty { $vis:vis mut $id:ident: $type:ty }) => {
$vis static $id: $crate::helpers::KeyMut<$schema, $type> =
$crate::helpers::KeyMut::new(|| <$schema as $crate::Schema>::key_mut());
};
($schema:ty { $vis:vis $id:ident: $type:ty = $init:expr }) => {
$vis static $id: $crate::helpers::Key<$schema, $type> =
$crate::helpers::Key::new(|| <$schema as $crate::Schema>::key_with_default::<_, _, _, $type>(
|| $init,
|x| x));
};
($schema:ty { $vis:vis mut $id:ident: $type:ty = $init:expr }) => {
$vis static $id: $crate::helpers::KeyMut<$schema, $type> =
$crate::helpers::KeyMut::new(|| <$schema as $crate::Schema>::key_mut_with_default::<_, _, _, _, $type>(
|| $init,
|x| x,
|x| x));
};
($schema:ty { $vis:vis $id:ident: $type:ty, $($tt:tt)* }) => {
$crate::key!($schema { $vis $id: $type });
$crate::key!($schema { $($tt)* });
};
($schema:ty { $vis:vis mut $id:ident: $type:ty, $($tt:tt)* }) => {
$crate::key!($schema { $vis mut $id: $type });
$crate::key!($schema { $($tt)* });
};
($schema:ty { $vis:vis $id:ident: $type:ty = $init:expr, $($tt:tt)* }) => {
$crate::key!($schema { $vis $id: $type = $init });
$crate::key!($schema { $($tt)* });
};
($schema:ty { $vis:vis mut $id:ident: $type:ty = $init:expr, $($tt:tt)* }) => {
$crate::key!($schema { $vis mut $id: $type = $init });
$crate::key!($schema { $($tt)* });
};
}