use std::{mem::ManuallyDrop, ops::DerefMut};
use derive_more::From;
use facet::{Def, PointerFlags, PtrConst, PtrMut, ReadLockResult, Shape, WriteLockResult};
use facet_reflect::{Peek, Poke};
#[derive(From)]
#[repr(C)]
pub enum MaybeMut<'mem, 'facet> {
Not(Peek<'mem, 'facet>),
Mut(Poke<'mem, 'facet>),
}
impl<'mem, 'facet> MaybeMut<'mem, 'facet> {
pub fn as_peek(&'mem self) -> Peek<'mem, 'facet> {
match self {
Self::Not(peek) => *peek,
Self::Mut(poke) => poke.as_peek(),
}
}
pub fn into_peek(self) -> Peek<'mem, 'facet> {
match self {
MaybeMut::Not(n) => n,
MaybeMut::Mut(m) => m.into_peek(),
}
}
pub fn shape(&self) -> &'static Shape {
self.as_peek().shape()
}
}
#[derive(Debug, thiserror::Error)]
#[error("{kind}")]
pub struct MakeLockError<'mem, 'facet> {
pub unchanged: Peek<'mem, 'facet>,
pub kind: MakeLockErrorKind,
}
#[derive(Debug, thiserror::Error)]
pub enum MakeLockErrorKind {
#[error("type cannot be locked")]
NotLockable,
#[error("locking of type failed")]
LockFailure,
#[error("could not upgrade weak pointer, no strong references exist")]
NotUpgradable,
}
#[derive(From)]
pub(crate) enum LockGuardType {
Write(WriteLockResult),
Read(ReadLockResult),
Upgrade {
strong_shape: &'static Shape,
allocation: PtrMut,
},
}
impl LockGuardType {
pub fn data_const(&self) -> PtrConst {
match self {
Self::Write(w) => w.data_const(),
Self::Read(r) => *r.data(),
Self::Upgrade {
strong_shape,
allocation,
} => {
let borrow_fn = strong_shape
.def
.into_pointer()
.expect("only pointer types get this lock type")
.vtable
.borrow_fn
.expect("all strong pointers have a borrow function");
unsafe { borrow_fn(allocation.as_const()) }
}
}
}
}
impl Drop for LockGuardType {
fn drop(&mut self) {
if let Self::Upgrade {
strong_shape,
allocation,
} = self
{
unsafe {
strong_shape.call_drop_in_place(*allocation);
}
unsafe {
strong_shape
.deallocate_mut(*allocation)
.expect("strong pointer is sized");
}
}
}
}
pub struct Guard<'lock_mem, 'facet> {
data: ManuallyDrop<MaybeMut<'lock_mem, 'facet>>,
guards: Vec<LockGuardType>,
}
impl Drop for Guard<'_, '_> {
fn drop(&mut self) {
unsafe {
ManuallyDrop::drop(&mut self.data);
}
while let Some(pop) = self.guards.pop() {
drop(pop);
}
}
}
impl<'lock, 'facet> Guard<'lock, 'facet> {
pub fn shape(&self) -> &'static Shape {
self.data.shape()
}
pub fn as_peek<'s>(&'s self) -> Peek<'s, 'facet> {
match &*self.data {
MaybeMut::Mut(m) => m.as_peek(),
MaybeMut::Not(n) => *n,
}
}
pub fn as_poke<'s>(&'s mut self) -> Option<Poke<'s, 'facet>> {
if let MaybeMut::Mut(m) = self.data.deref_mut() {
m.try_reborrow()
} else {
None
}
}
pub fn as_maybe<'s>(&'s mut self) -> MaybeMut<'s, 'facet> {
match self.data.deref_mut() {
MaybeMut::Mut(m) => {
let data = m.data();
let shape: &'static Shape = m.shape();
if let Some(poke) = m.try_reborrow() {
MaybeMut::Mut(poke)
} else {
let peek = unsafe { Peek::unchecked_new(data, shape) };
MaybeMut::Not(peek)
}
}
MaybeMut::Not(peek) => MaybeMut::Not(*peek),
}
}
unsafe fn take(mut self) -> (Vec<LockGuardType>, MaybeMut<'lock, 'facet>) {
let v = (core::mem::take(&mut self.guards), unsafe {
ManuallyDrop::take(&mut self.data)
});
core::mem::forget(self);
v
}
}
impl<'mem, 'facet> MaybeMut<'mem, 'facet> {
pub fn write<'lock>(self) -> Result<Guard<'lock, 'facet>, MakeLockError<'mem, 'facet>>
where
'mem: 'lock,
'facet: 'lock,
{
match self {
MaybeMut::Mut(v) => {
if let Def::Pointer(p) = v.as_peek().innermost_peek().shape().def
&& (p.flags.contains(PointerFlags::LOCK) || p.flags.contains(PointerFlags::WEAK))
{
Self::Not(v.into_peek()).write()
} else {
Ok(Guard {
guards: Vec::new(),
data: ManuallyDrop::new(v.into()),
})
}
}
MaybeMut::Not(v) => {
let v = v.innermost_peek();
let shape = v.shape();
let def = shape.def;
let Def::Pointer(pointer) = def else {
return Err(MakeLockError {
unchanged: v,
kind: MakeLockErrorKind::NotLockable,
});
};
let lock_fn =
pointer
.vtable
.write_fn
.or(pointer.vtable.lock_fn)
.ok_or(MakeLockError {
unchanged: v,
kind: MakeLockErrorKind::NotLockable,
});
let (mut guards, mut value): (Vec<LockGuardType>, MaybeMut<'lock, 'facet>) =
match lock_fn {
Ok(lock_fn) => {
let res = unsafe { lock_fn(v.data()) };
let Ok(lock) = res else {
return Err(MakeLockError {
unchanged: v,
kind: MakeLockErrorKind::LockFailure,
});
};
let poke: Poke<'lock, 'facet> = unsafe {
Poke::from_raw_parts(
*lock.data(),
shape
.inner
.expect("a smart pointer always has an inner shape"),
)
};
(vec![lock.into()], MaybeMut::Mut(poke))
}
Err(MakeLockError {
unchanged,
kind: MakeLockErrorKind::NotLockable,
}) if let Def::Pointer(pointer) = unchanged.shape().def
&& let Some(upgrade_fn) = pointer.vtable.upgrade_into_fn
&& let Some(strong_shape) =
def.into_pointer().ok().and_then(|x| x.strong()) =>
{
let strong = strong_shape
.allocate()
.expect("strong pointer is always sized");
let ptr = unsafe { v.data().into_mut() };
let guard = unsafe { upgrade_fn(ptr, strong) }
.map(|strong_instance| LockGuardType::Upgrade {
strong_shape,
allocation: strong_instance,
})
.ok_or(MakeLockError {
kind: MakeLockErrorKind::NotUpgradable,
unchanged: v,
})?;
let peek: Peek<'lock, 'facet> = unsafe {
Peek::unchecked_new(
guard.data_const(),
shape
.inner
.expect("a smart pointer always has an inner shape"),
)
};
(vec![guard], MaybeMut::Not(peek.innermost_peek()))
}
Err(e) => {
return Err(e);
}
};
while let Some(_inner) = value.as_peek().shape().inner
&& let Def::Pointer(def) = value.as_peek().shape().def
&& (def.flags.contains(PointerFlags::LOCK) ||
def.flags.contains(PointerFlags::WEAK) ||
def.flags.contains(PointerFlags::ATOMIC))
{
let shorter_peek: Peek<'mem, 'facet> =
unsafe { Peek::unchecked_new(value.as_peek().data(), value.shape()) };
let shorter_maybe: MaybeMut<'mem, 'facet> = MaybeMut::Not(shorter_peek);
let guard: Guard<'lock, 'facet> = match shorter_maybe.write() {
Ok(g) => g,
Err(e) => {
drop(guards);
return Err(MakeLockError {
unchanged: v,
kind: e.kind,
});
}
};
let (inner_guards, data) = unsafe { guard.take() };
guards.extend(inner_guards);
value = data;
}
Ok(Guard {
data: ManuallyDrop::new(value),
guards,
})
}
}
}
pub fn read<'lock>(self) -> Result<Guard<'lock, 'facet>, MakeLockError<'mem, 'facet>>
where
'mem: 'lock,
{
let peek = self.into_peek();
let v = peek.innermost_peek();
let shape = v.shape();
let def = shape.def;
let Def::Pointer(pointer) = def else {
return Ok(Guard {
guards: Vec::new(),
data: ManuallyDrop::new(MaybeMut::Not(v)),
});
};
let res: Result<LockGuardType, _> = if let Some(read_fn) = pointer.vtable.read_fn {
unsafe { read_fn(v.data()) }.map(Into::into)
} else if let Some(lock_fn) = pointer.vtable.lock_fn {
unsafe { lock_fn(v.data()) }.map(Into::into)
} else if let Some(upgrade_fn) = pointer.vtable.upgrade_into_fn
&& let Some(strong_shape) = def.into_pointer().ok().and_then(|x| x.strong())
{
let strong = strong_shape
.allocate()
.expect("strong pointer is always sized");
let ptr = unsafe { v.data().into_mut() };
Ok(unsafe { upgrade_fn(ptr, strong) }
.map(|strong_instance| LockGuardType::Upgrade {
strong_shape,
allocation: strong_instance,
})
.ok_or(MakeLockError {
kind: MakeLockErrorKind::NotUpgradable,
unchanged: v,
})?)
} else {
return Err(MakeLockError {
unchanged: v,
kind: MakeLockErrorKind::NotLockable,
});
};
let Ok(lock) = res else {
return Err(MakeLockError {
unchanged: v,
kind: MakeLockErrorKind::LockFailure,
});
};
let peek: Peek<'lock, 'facet> = unsafe {
Peek::unchecked_new(
lock.data_const(),
shape
.inner
.expect("a smart pointer always has an inner shape"),
)
};
let (mut guards, mut value): (Vec<LockGuardType>, MaybeMut<'lock, 'facet>) =
(vec![lock], MaybeMut::Not(peek.innermost_peek()));
while let Some(_inner) = value.as_peek().shape().inner
&& let Def::Pointer(def) = value.as_peek().shape().def
&& (def.flags.contains(PointerFlags::LOCK) ||
def.flags.contains(PointerFlags::WEAK) ||
def.flags.contains(PointerFlags::ATOMIC))
{
let shorter_peek: Peek<'mem, 'facet> =
unsafe { Peek::unchecked_new(value.as_peek().data(), value.shape()) };
let shorter_maybe: MaybeMut<'mem, 'facet> = MaybeMut::Not(shorter_peek);
let guard: Guard<'lock, 'facet> = match shorter_maybe.read() {
Ok(g) => g,
Err(e) => {
drop(guards);
return Err(MakeLockError {
unchanged: v,
kind: e.kind,
});
}
};
let (inner_guards, data) = unsafe { guard.take() };
guards.extend(inner_guards);
value = data;
}
Ok(Guard {
data: ManuallyDrop::new(value),
guards,
})
}
}
#[cfg(test)]
mod tests {
use facet::{Def, Facet, KnownPointer};
use facet_reflect::Peek;
#[derive(Debug, Facet)]
struct Foo {
value: String,
}
#[facet_testhelpers::test]
fn shared_reference() {
let a = Foo {
value: "aaaa".to_string(),
};
println!("{:#?}", <&Foo as Facet<'_>>::SHAPE.def);
assert!(
matches!(<&Foo as Facet<'_>>::SHAPE.def, Def::Pointer(p) if p.known == Some(KnownPointer::SharedReference))
);
let ref_a: &Foo = &a;
let ref_ref_a: &&Foo = &ref_a;
let peek = Peek::new(ref_ref_a);
println!("{:#?}", peek.shape().def); assert!(
matches!(peek.shape().def, Def::Pointer(p) if p.known == Some(KnownPointer::SharedReference))
);
}
}