use crate::{ucount, AtomicCounter};
use alloc::{alloc::dealloc, boxed::Box};
use branches::unlikely;
use core::{
alloc::Layout,
cell::UnsafeCell,
fmt,
hash::{Hash, Hasher},
marker::PhantomData,
mem::{self, forget, MaybeUninit},
ops::Deref,
pin::Pin,
ptr::NonNull,
sync::atomic::{fence, Ordering},
};
#[cfg(target_pointer_width = "64")]
const BARRIER: ucount = 512;
#[cfg(target_pointer_width = "32")]
const BARRIER: ucount = 64;
#[repr(C)]
struct ArcInner<T> {
data: UnsafeCell<T>,
counter: AtomicCounter,
}
pub struct Arc<T> {
ptr: NonNull<ArcInner<T>>,
phantom: PhantomData<Box<T>>,
}
unsafe impl<T: Sync + Send> Send for Arc<T> {}
unsafe impl<T: Sync + Send> Sync for Arc<T> {}
impl<T> Arc<T> {
#[inline]
pub fn new(data: T) -> Arc<T> {
let inner = Box::new(ArcInner {
data: UnsafeCell::new(data),
counter: AtomicCounter::new(1),
});
Arc {
ptr: unsafe { NonNull::new_unchecked(Box::leak(inner)) },
phantom: PhantomData,
}
}
#[inline]
#[must_use]
pub fn pin(data: T) -> Pin<Arc<T>> {
unsafe { Pin::new_unchecked(Arc::new(data)) }
}
#[inline]
#[must_use]
pub fn as_ptr(&self) -> *const T {
self.ptr.as_ptr() as *const T
}
#[inline]
pub fn into_raw(this: Self) -> *const T {
let ptr = Self::as_ptr(&this);
core::mem::forget(this);
ptr
}
#[inline]
pub unsafe fn from_raw(ptr: *const T) -> Self {
Arc {
ptr: NonNull::new_unchecked(ptr as *mut ArcInner<T>),
phantom: PhantomData,
}
}
#[inline]
#[must_use]
pub fn strong_count(&self) -> usize {
self.inner().counter.load(Ordering::Acquire) as usize
}
#[inline]
#[must_use]
pub fn ptr_eq(this: &Self, other: &Self) -> bool {
this.ptr.as_ptr() == other.ptr.as_ptr()
}
#[inline]
pub fn try_unwrap(this: Self) -> Result<T, Self> {
if this.is_unique() {
unsafe {
let inner = Box::from_raw(this.ptr.as_ptr());
core::mem::forget(this);
let val = core::ptr::read(inner.data.get());
core::mem::forget(inner.data);
Ok(val)
}
} else {
Err(this)
}
}
#[inline(always)]
fn inner(&self) -> &ArcInner<T> {
unsafe { self.ptr.as_ref() }
}
#[inline]
fn is_unique(&self) -> bool {
self.inner().counter.load(Ordering::Acquire) == 1
}
#[inline]
pub fn get_mut(this: &mut Self) -> Option<&mut T> {
if this.is_unique() {
unsafe { Some(Arc::get_mut_unchecked(this)) }
} else {
None
}
}
#[inline]
pub unsafe fn get_mut_unchecked(this: &mut Self) -> &mut T {
unsafe { &mut *(*this.ptr.as_ptr()).data.get() }
}
unsafe fn drop_slow(&mut self) {
let _ = Box::from_raw(self.ptr.as_ptr());
}
pub fn into_inner(this: Self) -> Option<T> {
let this = core::mem::ManuallyDrop::new(this);
let inner = this.inner();
if inner.counter.fetch_sub(1, Ordering::Release) != 1 {
return None;
}
inner.counter.load(Ordering::Acquire);
Some(unsafe {
let value = core::ptr::read(inner.data.get());
dealloc(this.ptr.as_ptr() as *mut u8, Layout::new::<ArcInner<T>>());
value
})
}
pub fn new_uninit() -> Arc<mem::MaybeUninit<T>> {
let inner = Box::new(ArcInner {
data: UnsafeCell::new(MaybeUninit::uninit()),
counter: AtomicCounter::new(1),
});
Arc {
ptr: unsafe { NonNull::new_unchecked(Box::leak(inner)) },
phantom: PhantomData,
}
}
}
impl<T> Arc<MaybeUninit<T>> {
pub unsafe fn assume_init(self) -> Arc<T> {
let ptr = self.ptr.as_ptr();
forget(self);
Arc {
ptr: NonNull::new_unchecked(ptr as *mut ArcInner<T>),
phantom: PhantomData,
}
}
}
impl<T: Clone> Arc<T> {
#[inline]
pub fn unwrap_or_clone(this: Self) -> T {
Arc::try_unwrap(this).unwrap_or_else(|rc| (*rc).clone())
}
#[inline]
fn optimized_clone(&self) -> Arc<T> {
let mut buffer: Box<MaybeUninit<ArcInner<T>>> = Box::new(MaybeUninit::uninit());
let ptr = unsafe {
(*buffer.as_ptr()).data.get().write(T::clone(self));
(*buffer.as_mut_ptr()).counter = AtomicCounter::new(1);
NonNull::new_unchecked(Box::leak(buffer) as *mut _ as *mut ArcInner<T>)
};
Arc {
ptr,
phantom: PhantomData,
}
}
#[inline]
pub fn make_mut(this: &mut Arc<T>) -> &mut T {
if !this.is_unique() {
*this = this.optimized_clone();
}
unsafe { Self::get_mut_unchecked(this) }
}
}
impl<T> Deref for Arc<T> {
type Target = T;
#[inline(always)]
fn deref(&self) -> &Self::Target {
unsafe { &*(self.inner().data.get() as *const T) }
}
}
impl<T> From<T> for Arc<T> {
#[inline(always)]
fn from(value: T) -> Self {
Arc::new(value)
}
}
#[inline(never)]
fn drop_arc_and_panic_no_inline<T>(ptr: NonNull<ArcInner<T>>) {
drop(Arc {
ptr,
phantom: PhantomData,
});
panic!("reference counter overflow");
}
impl<T> Clone for Arc<T> {
#[inline]
fn clone(&self) -> Self {
let count = self.inner().counter.fetch_add(1, Ordering::Relaxed);
if unlikely(count >= ucount::MAX - BARRIER) {
drop_arc_and_panic_no_inline(self.ptr);
}
Self {
ptr: self.ptr,
phantom: PhantomData,
}
}
}
impl<T> Drop for Arc<T> {
#[inline]
fn drop(&mut self) {
if self.inner().counter.fetch_sub(1, Ordering::Release) != 1 {
return;
}
fence(Ordering::Acquire);
unsafe { self.drop_slow() };
}
}
impl<T: Hash> Hash for Arc<T> {
#[inline]
fn hash<H: Hasher>(&self, state: &mut H) {
(**self).hash(state);
}
}
impl<T: fmt::Display> fmt::Display for Arc<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Display::fmt(&**self, f)
}
}
impl<T: fmt::Debug> fmt::Debug for Arc<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&**self, f)
}
}
impl<T> fmt::Pointer for Arc<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Pointer::fmt(&(&**self as *const T), f)
}
}
impl<T: Default> Default for Arc<T> {
#[inline]
fn default() -> Arc<T> {
Arc::new(Default::default())
}
}
impl<T: PartialEq> PartialEq for Arc<T> {
#[inline]
fn eq(&self, other: &Arc<T>) -> bool {
self.deref().eq(other)
}
}
impl<T: Eq> Eq for Arc<T> {}
impl<T: PartialOrd> PartialOrd for Arc<T> {
#[inline]
fn partial_cmp(&self, other: &Arc<T>) -> Option<core::cmp::Ordering> {
(**self).partial_cmp(&**other)
}
#[inline]
fn lt(&self, other: &Arc<T>) -> bool {
**self < **other
}
#[inline]
fn le(&self, other: &Arc<T>) -> bool {
**self <= **other
}
#[inline]
fn gt(&self, other: &Arc<T>) -> bool {
**self > **other
}
#[inline]
fn ge(&self, other: &Arc<T>) -> bool {
**self >= **other
}
}
impl<T: Ord> Ord for Arc<T> {
#[inline]
fn cmp(&self, other: &Arc<T>) -> core::cmp::Ordering {
(**self).cmp(&**other)
}
}
impl<T> core::borrow::Borrow<T> for Arc<T> {
#[inline(always)]
fn borrow(&self) -> &T {
self
}
}
impl<T> AsRef<T> for Arc<T> {
#[inline(always)]
fn as_ref(&self) -> &T {
self
}
}
impl<T> Unpin for Arc<T> {}