use core::any::Any;
use core::borrow::Borrow;
use core::cell::Cell;
use core::convert::Infallible;
use core::convert::TryFrom;
use core::hash::{Hash, Hasher};
use core::marker::PhantomData;
use core::ops::Deref;
use core::ptr::NonNull;
use core::sync::atomic;
use core::sync::atomic::Ordering;
use core::{fmt, mem, ptr};
use std::error::Error;
use std::panic::{RefUnwindSafe, UnwindSafe};
mod thread_id;
use thread_id::{AtomicOptionThreadId, ThreadId};
#[inline]
const fn senitel<T>() -> NonNull<T> {
unsafe { NonNull::new_unchecked(usize::MAX as *mut T) }
}
#[inline]
fn is_senitel<T: ?Sized>(ptr: *const T) -> bool {
ptr.cast::<()>() == senitel().as_ptr()
}
mod state_trait {
use core::fmt::Debug;
pub trait RcState: Debug {
const SHARED: bool;
}
}
use state_trait::RcState;
pub mod state {
#[derive(Debug, Clone, Copy)]
pub enum Shared {}
impl super::RcState for Shared {
const SHARED: bool = true;
}
#[derive(Debug, Clone, Copy)]
pub enum Local {}
impl super::RcState for Local {
const SHARED: bool = false;
}
}
use state::{Local, Shared};
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum UpgradeError {
ValueDropped,
WrongThread,
}
impl fmt::Display for UpgradeError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match *self {
Self::ValueDropped => f.write_str("value was already dropped"),
Self::WrongThread => {
f.write_str("tried to get a local reference while another thread was the owner")
}
}
}
}
impl Error for UpgradeError {}
impl From<Infallible> for UpgradeError {
fn from(x: Infallible) -> UpgradeError {
match x {}
}
}
#[derive(Debug)]
struct RcMeta {
owner: AtomicOptionThreadId,
strong_local: Cell<usize>,
strong_shared: atomic::AtomicUsize,
weak: atomic::AtomicUsize,
}
#[repr(C)]
struct RcBox<T: ?Sized> {
meta: RcMeta,
data: T,
}
impl RcMeta {
#[inline(always)]
fn inc_strong_local(&self) {
let counter = self.strong_local.get();
if counter == usize::MAX {
panic!("reference counter overflow");
}
self.strong_local.set(counter + 1);
}
#[inline]
fn try_inc_strong_local(&self) -> Result<(), ()> {
let counter = self.strong_local.get();
if counter == usize::MAX {
panic!("reference counter overflow");
} else if counter == 0 {
self.try_inc_strong_shared()?;
}
self.strong_local.set(counter + 1);
Ok(())
}
#[inline(always)]
fn dec_strong_local(&self) -> bool {
let counter = self.strong_local.get();
self.strong_local.set(counter - 1);
if counter == 1 {
self.remove_last_local_reference()
} else {
false
}
}
fn remove_last_local_reference(&self) -> bool {
let old_shared = self.strong_shared.fetch_sub(1, Ordering::Release);
if old_shared == 0 {
panic!("reference counter underflow");
}
self.owner.store(None, Ordering::Release);
old_shared == 1
}
#[inline]
fn inc_strong_shared(&self) {
let old_counter = self.strong_shared.fetch_add(1, Ordering::Relaxed);
if old_counter == usize::MAX {
panic!("reference counter overflow");
}
}
#[inline]
fn try_inc_strong_shared(&self) -> Result<(), ()> {
self.strong_shared
.fetch_update(
Ordering::Relaxed,
Ordering::Relaxed,
|old_counter| match old_counter {
0 => None,
usize::MAX => panic!("reference counter overflow"),
_ => Some(old_counter + 1),
},
)
.map(|_| ())
.map_err(|_| ())
}
#[inline]
fn dec_strong_shared(&self) -> bool {
let old_counter = self.strong_shared.fetch_sub(1, Ordering::Release);
if old_counter == 0 {
panic!("reference counter underflow");
}
old_counter == 1 && self.owner.load(Ordering::Relaxed).is_none()
}
#[inline]
fn inc_weak(&self) {
const MAX_COUNT: usize = usize::MAX - 1;
let mut counter = self.weak.load(Ordering::Relaxed);
loop {
match counter {
usize::MAX => {
core::hint::spin_loop();
counter = self.weak.load(Ordering::Relaxed);
continue;
}
MAX_COUNT => panic!("weak counter overflow"),
0 => panic!("BUG: weak resurrection of dead counted reference"),
_ => {
let result = self.weak.compare_exchange_weak(
counter,
counter + 1,
Ordering::Acquire,
Ordering::Relaxed,
);
match result {
Ok(_) => break,
Err(old) => counter = old,
}
}
}
}
}
#[inline]
fn inc_weak_nolock(&self) {
const MAX_COUNT: usize = usize::MAX - 1;
match self.weak.fetch_add(1, Ordering::Relaxed) {
usize::MAX => panic!("BUG: weak counter locked"),
MAX_COUNT => panic!("weak counter overflow"),
0 => panic!("BUG: weak resurrection of dead counted reference"),
_ => (),
}
}
#[inline]
fn dec_weak(&self) -> bool {
let old_counter = self.weak.fetch_sub(1, Ordering::Release);
if old_counter == 0 {
panic!("weak counter underflow");
}
old_counter == 1
}
#[inline]
fn has_unique_ref(&self) -> bool {
let result =
self.weak
.compare_exchange(1, usize::MAX, Ordering::Acquire, Ordering::Relaxed);
if result.is_ok() {
let mut count = self.strong_shared.load(Ordering::Acquire);
if count == 1 {
let owner = self.owner.load(Ordering::Relaxed);
if owner == Some(ThreadId::current_thread()) {
count = self.strong_local.get();
}
}
self.weak.store(1, Ordering::Release);
count == 1
} else {
false
}
}
}
#[must_use]
pub struct HybridRc<T: ?Sized, State: RcState> {
ptr: NonNull<RcBox<T>>,
phantom: PhantomData<(State, RcBox<T>)>,
}
pub type Rc<T> = HybridRc<T, Local>;
pub type Arc<T> = HybridRc<T, Shared>;
impl<T: ?Sized, State: RcState> HybridRc<T, State> {
#[inline(always)]
fn from_inner(ptr: NonNull<RcBox<T>>) -> Self {
Self {
ptr,
phantom: PhantomData,
}
}
#[inline(always)]
fn data(&self) -> &T {
unsafe { &(*self.ptr.as_ptr()).data }
}
#[inline(always)]
fn meta(&self) -> &RcMeta {
unsafe { &(*self.ptr.as_ptr()).meta }
}
#[must_use]
#[inline]
pub unsafe fn get_mut_unchecked(this: &mut Self) -> &mut T {
unsafe { &mut (*this.ptr.as_ptr()).data }
}
#[must_use]
pub fn get_mut(this: &mut Self) -> Option<&mut T> {
if this.meta().has_unique_ref() {
unsafe { Some(Self::get_mut_unchecked(this)) }
} else {
None
}
}
#[must_use]
#[inline]
pub fn as_ptr(this: &Self) -> *const T {
this.data()
}
#[inline]
pub fn downgrade(this: &Self) -> Weak<T> {
this.meta().inc_weak();
Weak { ptr: this.ptr }
}
unsafe fn drop_contents_and_maybe_box(&mut self) {
unsafe {
ptr::drop_in_place(Self::get_mut_unchecked(self));
}
if self.meta().dec_weak() {
let ptr: *mut RcBox<mem::ManuallyDrop<T>> =
unsafe { mem::transmute(self.ptr.as_mut()) };
mem::drop(unsafe { Box::from_raw(ptr) });
}
}
}
impl<State: RcState> HybridRc<dyn Any, State> {
#[inline]
pub fn downcast<T: Any>(self) -> Result<HybridRc<T, State>, Self> {
if (*self).is::<T>() {
let ptr = self.ptr.cast::<RcBox<T>>();
std::mem::forget(self);
Ok(HybridRc::from_inner(ptr))
} else {
Err(self)
}
}
}
impl<T: ?Sized> Rc<T> {
#[inline]
pub fn to_shared(this: &Self) -> Arc<T> {
this.meta().inc_strong_shared();
Arc::from_inner(this.ptr)
}
}
impl<T: ?Sized> Arc<T> {
#[must_use]
#[inline]
pub fn to_local(this: &Self) -> Option<Rc<T>> {
let meta = this.meta();
let current_thread = ThreadId::current_thread();
let owner = match meta.owner.compare_exchange(
None,
Some(current_thread),
Ordering::Acquire,
Ordering::Relaxed,
) {
Ok(_) => None,
Err(owner) => owner,
};
match owner {
None => {
meta.try_inc_strong_local()
.expect("inconsistent reference count (shared == 0)");
Some(Rc::from_inner(this.ptr))
}
Some(v) if v == current_thread => {
meta.inc_strong_local();
Some(Rc::from_inner(this.ptr))
}
Some(_) => None,
}
}
}
impl<T> Rc<T> {
#[inline]
pub fn new(data: T) -> Self {
let inner = Box::new(RcBox::<T> {
meta: RcMeta {
owner: ThreadId::current_thread().into(),
strong_local: Cell::new(1),
strong_shared: 1.into(),
weak: 1.into(),
},
data,
});
Self::from_inner(Box::leak(inner).into())
}
}
impl<T> Arc<T> {
#[inline]
pub fn new(data: T) -> Self {
let inner = Box::new(RcBox::<T> {
meta: RcMeta {
owner: None.into(),
strong_local: Cell::new(0),
strong_shared: 1.into(),
weak: 1.into(),
},
data,
});
Self::from_inner(Box::leak(inner).into())
}
}
impl<T: ?Sized> Clone for HybridRc<T, Local> {
#[inline]
fn clone(&self) -> Self {
self.meta().inc_strong_local();
Self::from_inner(self.ptr)
}
}
impl<T: ?Sized> Clone for HybridRc<T, Shared> {
#[inline]
fn clone(&self) -> Self {
self.meta().inc_strong_shared();
Self::from_inner(self.ptr)
}
}
impl<T: ?Sized, State: RcState> Drop for HybridRc<T, State> {
#[inline]
fn drop(&mut self) {
let no_more_strong_refs = if State::SHARED {
self.meta().dec_strong_shared()
} else {
self.meta().dec_strong_local()
};
if no_more_strong_refs {
unsafe {
self.drop_contents_and_maybe_box();
}
}
}
}
impl<T: ?Sized, State: RcState> Deref for HybridRc<T, State> {
type Target = T;
#[inline]
fn deref(&self) -> &T {
self.data()
}
}
impl<T: ?Sized, State: RcState> Borrow<T> for HybridRc<T, State> {
#[inline]
fn borrow(&self) -> &T {
&**self
}
}
impl<T: ?Sized, State: RcState> AsRef<T> for HybridRc<T, State> {
#[inline]
fn as_ref(&self) -> &T {
&**self
}
}
unsafe impl<T: ?Sized + Sync + Send> Send for HybridRc<T, Shared> {}
unsafe impl<T: ?Sized + Sync + Send> Sync for HybridRc<T, Shared> {}
impl<T: RefUnwindSafe + ?Sized, State: RcState> UnwindSafe for HybridRc<T, State> {}
impl<T: RefUnwindSafe> RefUnwindSafe for HybridRc<T, Shared> {}
impl<T: Any + 'static, State: RcState> From<HybridRc<T, State>>
for HybridRc<dyn Any + 'static, State>
{
#[inline]
fn from(src: HybridRc<T, State>) -> Self {
let ptr = src.ptr.as_ptr() as *mut RcBox<dyn Any>;
std::mem::forget(src);
Self::from_inner(unsafe { NonNull::new_unchecked(ptr) })
}
}
impl<T: ?Sized> From<Rc<T>> for HybridRc<T, Shared> {
#[inline]
fn from(src: Rc<T>) -> Self {
HybridRc::to_shared(&src)
}
}
impl<T: ?Sized> TryFrom<Arc<T>> for HybridRc<T, Local> {
type Error = Arc<T>;
#[inline]
fn try_from(src: Arc<T>) -> Result<Self, Self::Error> {
match HybridRc::to_local(&src) {
Some(result) => Ok(result),
None => Err(src),
}
}
}
impl<T: Default> Default for HybridRc<T, Local> {
#[inline]
fn default() -> Self {
Self::new(Default::default())
}
}
impl<T: Default> Default for HybridRc<T, Shared> {
#[inline]
fn default() -> Self {
Self::new(Default::default())
}
}
impl<T: ?Sized + PartialEq, S1: RcState, S2: RcState> PartialEq<HybridRc<T, S2>>
for HybridRc<T, S1>
{
#[inline]
fn eq(&self, other: &HybridRc<T, S2>) -> bool {
**self == **other
}
}
impl<T: ?Sized + Eq, State: RcState> Eq for HybridRc<T, State> {}
impl<T: ?Sized + Hash, State: RcState> Hash for HybridRc<T, State> {
#[inline]
fn hash<H: Hasher>(&self, state: &mut H) {
Self::data(self).hash(state);
}
}
impl<T: ?Sized + fmt::Display, State: RcState> fmt::Display for HybridRc<T, State> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Display::fmt(&Self::data(self), f)
}
}
impl<T: ?Sized + fmt::Debug, State: RcState> fmt::Debug for HybridRc<T, State> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&Self::data(self), f)
}
}
impl<T: ?Sized, State: RcState> fmt::Pointer for HybridRc<T, State> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if f.alternate() {
fmt::Pointer::fmt(&Self::as_ptr(self), f)?;
f.write_str(if State::SHARED {
" [shared]"
} else {
" [local]"
})
} else {
fmt::Pointer::fmt(&Self::as_ptr(self), f)
}
}
}
#[must_use]
pub struct Weak<T: ?Sized> {
ptr: NonNull<RcBox<T>>,
}
impl<T: ?Sized> Weak<T> {
#[inline]
fn meta(&self) -> Option<&RcMeta> {
if is_senitel(self.ptr.as_ptr()) {
None
} else {
Some(unsafe { &(*self.ptr.as_ptr()).meta })
}
}
#[must_use]
#[inline]
pub fn as_ptr(&self) -> *const T {
let ptr: *mut RcBox<T> = self.ptr.as_ptr();
if is_senitel(ptr) {
ptr as *const T
} else {
unsafe { ptr::addr_of_mut!((*ptr).data) }
}
}
#[inline]
pub fn upgrade_local(&self) -> Result<Rc<T>, UpgradeError> {
let meta = self.meta().ok_or(UpgradeError::ValueDropped)?;
let current_thread = ThreadId::current_thread();
let owner = match meta.owner.compare_exchange(
None,
Some(current_thread),
Ordering::Acquire,
Ordering::Relaxed,
) {
Ok(_) => Some(current_thread),
Err(owner) => owner,
};
if owner == Some(current_thread) {
if meta.try_inc_strong_local().is_ok() {
Ok(HybridRc::<T, Local>::from_inner(self.ptr))
} else {
meta.owner.store(None, Ordering::Relaxed);
Err(UpgradeError::ValueDropped)
}
} else {
Err(UpgradeError::WrongThread)
}
}
#[inline]
pub fn upgrade(&self) -> Result<Arc<T>, UpgradeError> {
let meta = self.meta().ok_or(UpgradeError::ValueDropped)?;
meta.try_inc_strong_shared()
.map_err(|_| UpgradeError::ValueDropped)?;
Ok(HybridRc::<T, Shared>::from_inner(self.ptr))
}
}
impl<T> Weak<T> {
pub fn new() -> Weak<T> {
Weak { ptr: senitel() }
}
}
impl<T: ?Sized> fmt::Debug for Weak<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "(Weak)")
}
}
impl<T: ?Sized> fmt::Pointer for Weak<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if f.alternate() {
fmt::Pointer::fmt(&Self::as_ptr(self), f)?;
f.write_str(" [weak]")
} else {
fmt::Pointer::fmt(&Self::as_ptr(self), f)
}
}
}
impl<T> Default for Weak<T> {
#[inline]
fn default() -> Self {
Self::new()
}
}
impl<T: ?Sized> Clone for Weak<T> {
#[inline]
fn clone(&self) -> Self {
if let Some(meta) = self.meta() {
meta.inc_weak_nolock();
}
Self { ptr: self.ptr }
}
}
impl<T: ?Sized> Drop for Weak<T> {
#[inline]
fn drop(&mut self) {
if let Some(meta) = self.meta() {
let last_reference = meta.dec_weak();
if last_reference {
unsafe {
let ptr: *mut RcBox<mem::ManuallyDrop<T>> = mem::transmute(self.ptr.as_mut());
mem::drop(Box::from_raw(ptr));
}
}
}
}
}
unsafe impl<T: ?Sized + Sync + Send> Send for Weak<T> {}
unsafe impl<T: ?Sized + Sync + Send> Sync for Weak<T> {}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::hash_map;
use std::convert::TryInto;
use std::thread;
thread_local! {
static DROP_COUNTER: Cell<usize> = Cell::new(0);
}
#[derive(Debug, PartialEq, Eq, Hash)]
struct Test {}
impl Drop for Test {
fn drop(&mut self) {
DROP_COUNTER.with(|x| x.set(x.get() + 1));
}
}
impl Default for Test {
fn default() -> Self {
Test {}
}
}
#[test]
fn test_traits() {
let a = Rc::<Test>::default();
let b = Arc::<Test>::default();
assert_eq!(a, b);
let mut map = hash_map::HashMap::new();
map.insert(a.clone(), ());
assert!(map.get(&a).is_some());
assert_eq!(&*a, a.borrow());
assert_eq!(a.deref(), a.as_ref());
}
#[test]
fn test_local() {
DROP_COUNTER.with(|x| x.set(0));
let a = Rc::new(Test {});
let b = a.clone();
mem::drop(a);
assert_eq!(DROP_COUNTER.with(|x| x.get()), 0);
mem::drop(b);
assert_eq!(DROP_COUNTER.with(|x| x.get()), 1);
}
#[test]
fn test_shared() {
DROP_COUNTER.with(|x| x.set(0));
let a = Arc::new(Test {});
let b = a.clone();
mem::drop(a);
assert_eq!(DROP_COUNTER.with(|x| x.get()), 0);
mem::drop(b);
assert_eq!(DROP_COUNTER.with(|x| x.get()), 1);
}
#[test]
fn test_shared_to_local() {
DROP_COUNTER.with(|x| x.set(0));
let a = Arc::new(Test {});
let b = Arc::to_local(&a).unwrap();
let _: Rc<Test> = Arc::try_into(a.clone()).unwrap();
mem::drop(a);
assert_eq!(DROP_COUNTER.with(|x| x.get()), 0);
mem::drop(b);
assert_eq!(DROP_COUNTER.with(|x| x.get()), 1);
}
#[test]
fn test_shared_to_local_on_wrong_thread() {
DROP_COUNTER.with(|x| x.set(0));
let a = Arc::new(Test {});
let b = Arc::to_local(&a).unwrap();
assert!(thread::spawn(
move || Arc::to_local(&a).is_none() && TryInto::<Rc<Test>>::try_into(a).is_err()
)
.join()
.unwrap());
assert_eq!(DROP_COUNTER.with(|x| x.get()), 0);
mem::drop(b);
assert_eq!(DROP_COUNTER.with(|x| x.get()), 1);
}
#[test]
fn test_local_to_shared() {
DROP_COUNTER.with(|x| x.set(0));
let a = Rc::new(Test {});
let b = Rc::to_shared(&a);
mem::drop(a);
assert_eq!(DROP_COUNTER.with(|x| x.get()), 0);
mem::drop(b);
assert_eq!(DROP_COUNTER.with(|x| x.get()), 1);
}
#[test]
fn test_get_mut() {
let mut a = Rc::new(Test {});
assert_eq!(
unsafe { Rc::get_mut_unchecked(&mut a) } as *const Test,
Rc::as_ptr(&a)
);
assert!(Rc::get_mut(&mut a).is_some());
let mut b = Rc::clone(&a);
assert!(Rc::get_mut(&mut a).is_none());
mem::drop(a);
assert!(Rc::get_mut(&mut b).is_some());
}
#[test]
fn test_dangling_weak() {
DROP_COUNTER.with(|x| x.set(0));
let w = Weak::<Test>::new();
assert!(w.upgrade_local().is_err());
assert!(w.upgrade().is_err());
assert_eq!(DROP_COUNTER.with(|x| x.get()), 0);
}
#[test]
fn test_weak() {
DROP_COUNTER.with(|x| x.set(0));
let a = Rc::new(Test {});
let a_ptr = Rc::as_ptr(&a);
let w = Rc::downgrade(&a);
assert!(thread::spawn(move || w.upgrade_local().is_err())
.join()
.unwrap());
let w = Rc::downgrade(&a);
assert!(thread::spawn(move || w.upgrade().is_ok()).join().unwrap());
let b = Rc::to_shared(&a);
mem::drop(a);
let w = Arc::downgrade(&b);
let _ = w.clone();
assert_eq!(DROP_COUNTER.with(|x| x.get()), 0);
assert!(w.upgrade_local().is_ok());
mem::drop(b);
assert_eq!(DROP_COUNTER.with(|x| x.get()), 1);
assert!(w.upgrade_local().is_err());
assert_eq!(w.as_ptr(), a_ptr);
assert!(!is_senitel(w.as_ptr()));
}
#[test]
fn test_fmt() {
let a = Rc::new(String::from("abc"));
let _ = format!("{0:?} {0:p} {0}", a);
assert!(format!("{0:#p}", a).find("local").is_some());
let b = Rc::to_shared(&a);
assert!(format!("{0:#p}", b).find("shared").is_some());
let _ = format!("{0:?} {0}", UpgradeError::ValueDropped);
let _ = format!("{0:?} {0}", UpgradeError::WrongThread);
}
#[test]
fn test_senitel() {
let x = senitel::<()>();
assert!(is_senitel(x.as_ptr()));
let w = Weak::<()>::default();
assert!(is_senitel(w.as_ptr()));
}
}