use core::cell::Cell;
use core::marker::PhantomData;
use core::ptr::NonNull;
use core::sync::atomic::{AtomicU32, Ordering};
use bun_core::ThreadLock;
fn type_base_name(name: &'static str) -> &'static str {
let bytes = name.as_bytes();
let end = bun_core::strings::index_of_char_usize(bytes, b'<').unwrap_or(bytes.len());
match bun_core::strings::last_index_of(&bytes[..end], b"::") {
Some(i) => &name[i + 2..],
None => name,
}
}
pub trait RefCounted: Sized {
fn debug_name() -> &'static str {
type_base_name(core::any::type_name::<Self>())
}
unsafe fn get_ref_count(this: *mut Self) -> *mut RefCount<Self>;
unsafe fn destructor(this: *mut Self);
}
pub trait ThreadSafeRefCounted: Sized {
fn debug_name() -> &'static str {
type_base_name(core::any::type_name::<Self>())
}
unsafe fn get_ref_count(this: *mut Self) -> *mut ThreadSafeRefCount<Self>;
#[inline]
unsafe fn destructor(this: *mut Self) {
drop(unsafe { Box::from_raw(this) });
}
}
pub trait AnyRefCounted: Sized {
unsafe fn rc_ref(this: *mut Self);
unsafe fn rc_deref(this: *mut Self);
unsafe fn rc_has_one_ref(this: *const Self) -> bool;
#[inline]
unsafe fn rc_assert_valid(_this: *const Self) {}
}
pub struct RefCount<T: RefCounted> {
raw_count: Cell<u32>,
thread: ThreadLock,
#[cfg(debug_assertions)]
debug: DebugData,
_phantom: PhantomData<*const T>,
}
impl<T: RefCounted> RefCount<T> {
pub fn init() -> Self {
Self::init_exact_refs(1)
}
pub fn init_exact_refs(count: u32) -> Self {
debug_assert!(count > 0);
Self {
raw_count: Cell::new(count),
thread: ThreadLock::init_locked_if_non_comptime(),
#[cfg(debug_assertions)]
debug: DebugData::empty(),
_phantom: PhantomData,
}
}
pub unsafe fn ref_(self_: *mut T) {
let count = unsafe { &*T::get_ref_count(self_) };
count.assert_valid();
bun_core::scoped_log!(
ref_count,
"0x{:x} ref {} -> {}:",
self_ as usize,
count.raw_count.get(),
count.raw_count.get() + 1,
);
count.assert_single_threaded();
count.raw_count.set(count.raw_count.get() + 1);
}
pub unsafe fn deref(self_: *mut T) {
let count = unsafe { &*T::get_ref_count(self_) };
count.assert_valid(); bun_core::scoped_log!(
ref_count,
"0x{:x} deref {} -> {}:",
self_ as usize,
count.raw_count.get(),
count.raw_count.get() - 1,
);
count.assert_single_threaded();
count.raw_count.set(count.raw_count.get() - 1);
if count.raw_count.get() == 0 {
#[cfg(debug_assertions)]
{
unsafe { (*T::get_ref_count(self_)).debug.deinit() };
}
unsafe { T::destructor(self_) };
}
}
pub fn has_one_ref(&self) -> bool {
self.assert_single_threaded();
self.raw_count.get() == 1
}
pub fn get(&self) -> u32 {
self.raw_count.get()
}
pub fn assert_no_refs(&self) {
assert!(self.raw_count.get() == 0);
}
fn assert_single_threaded(&self) {
self.thread.lock_or_assert();
}
#[inline]
pub fn assert_valid(&self) {
#[cfg(debug_assertions)]
self.debug.assert_valid();
}
}
impl<T: RefCounted> AnyRefCounted for T {
unsafe fn rc_ref(this: *mut Self) {
unsafe { RefCount::<T>::ref_(this) }
}
unsafe fn rc_deref(this: *mut Self) {
unsafe { RefCount::<T>::deref(this) }
}
unsafe fn rc_has_one_ref(this: *const Self) -> bool {
unsafe { (*T::get_ref_count(this.cast_mut())).has_one_ref() }
}
unsafe fn rc_assert_valid(this: *const Self) {
unsafe { (*T::get_ref_count(this.cast_mut())).assert_valid() }
}
}
pub struct ThreadSafeRefCount<T: ThreadSafeRefCounted> {
raw_count: AtomicU32,
#[cfg(debug_assertions)]
debug: DebugData,
_phantom: PhantomData<*const T>,
}
impl<T: ThreadSafeRefCounted> ThreadSafeRefCount<T> {
pub fn init() -> Self {
Self::init_exact_refs(1)
}
pub fn init_exact_refs(count: u32) -> Self {
debug_assert!(count > 0);
Self {
raw_count: AtomicU32::new(count),
#[cfg(debug_assertions)]
debug: DebugData::empty(),
_phantom: PhantomData,
}
}
pub unsafe fn ref_(self_: *mut T) {
let count = unsafe { &*T::get_ref_count(self_) };
count.assert_valid();
let old_count = count.raw_count.fetch_add(1, Ordering::SeqCst);
bun_core::scoped_log!(
ref_count,
"0x{:x} ref {} -> {}",
self_ as usize,
old_count,
old_count + 1,
);
debug_assert!(old_count > 0);
}
pub unsafe fn deref(self_: *mut T) {
let count = unsafe { &*T::get_ref_count(self_) };
count.assert_valid();
let old_count = count.raw_count.fetch_sub(1, Ordering::SeqCst);
bun_core::scoped_log!(
ref_count,
"0x{:x} deref {} -> {}",
self_ as usize,
old_count,
old_count - 1,
);
debug_assert!(old_count > 0);
if old_count == 1 {
#[cfg(debug_assertions)]
{
unsafe { (*T::get_ref_count(self_)).debug.deinit() };
}
unsafe { T::destructor(self_) };
}
}
pub unsafe fn release(self_: *mut T) -> bool {
let count = unsafe { &*T::get_ref_count(self_) };
count.assert_valid();
let old_count = count.raw_count.fetch_sub(1, Ordering::SeqCst);
bun_core::scoped_log!(
ref_count,
"0x{:x} deref {} -> {}",
self_ as usize,
old_count,
old_count - 1,
);
debug_assert!(old_count > 0);
if old_count == 1 {
#[cfg(debug_assertions)]
{
unsafe { (*T::get_ref_count(self_)).debug.deinit() };
}
true
} else {
false
}
}
pub fn get(&self) -> u32 {
self.raw_count.load(Ordering::SeqCst)
}
pub fn has_one_ref(&self) -> bool {
self.assert_valid();
self.get() == 1
}
pub fn assert_no_refs(&self) {
assert!(self.raw_count.load(Ordering::SeqCst) == 0);
}
#[inline]
pub fn assert_valid(&self) {
#[cfg(debug_assertions)]
self.debug.assert_valid();
}
}
pub unsafe trait CellRefCounted: Sized {
fn ref_count(&self) -> &Cell<u32>;
unsafe fn ref_count_raw<'a>(this: *const Self) -> &'a Cell<u32>;
#[inline]
unsafe fn destroy(this: *mut Self) {
drop(unsafe { Box::from_raw(this) });
}
#[inline]
fn ref_(&self) {
let rc = self.ref_count();
rc.set(rc.get() + 1);
}
#[inline]
unsafe fn deref(this: *mut Self) {
let rc = unsafe { Self::ref_count_raw(this) };
let n = rc.get() - 1;
rc.set(n);
if n == 0 {
unsafe { Self::destroy(this) };
}
}
#[inline]
fn deref_nn(this: NonNull<Self>) {
unsafe { Self::deref(this.as_ptr()) };
}
}
#[must_use = "dropping a RefPtr releases its ref"]
#[repr(transparent)]
pub struct RefPtr<T: AnyRefCounted>(NonNull<T>);
impl<T: AnyRefCounted> Drop for RefPtr<T> {
#[inline]
fn drop(&mut self) {
unsafe { T::rc_deref(self.0.as_ptr()) };
}
}
unsafe impl<T: AnyRefCounted + Send + Sync> Send for RefPtr<T> {}
unsafe impl<T: AnyRefCounted + Send + Sync> Sync for RefPtr<T> {}
impl<T: AnyRefCounted> RefPtr<T> {
#[inline]
pub fn new(value: T) -> Self {
let ptr = bun_core::heap::into_raw_nn(Box::new(value));
debug_assert!(unsafe { T::rc_has_one_ref(ptr.as_ptr()) });
Self(ptr)
}
#[inline]
pub unsafe fn init_ref(raw_ptr: *mut T) -> Self {
unsafe {
T::rc_assert_valid(raw_ptr);
T::rc_ref(raw_ptr);
Self(NonNull::new_unchecked(raw_ptr))
}
}
#[inline]
pub fn from_this(this: crate::ThisPtr<T>) -> Self {
unsafe { Self::init_ref(this.as_ptr()) }
}
#[inline]
pub unsafe fn from_raw(raw_ptr: *mut T) -> Self {
unsafe {
T::rc_assert_valid(raw_ptr);
Self(NonNull::new_unchecked(raw_ptr))
}
}
#[inline]
pub fn into_raw(self) -> *mut T {
self.into_non_null().as_ptr()
}
#[inline]
pub fn into_non_null(self) -> NonNull<T> {
core::mem::ManuallyDrop::new(self).0
}
#[inline]
pub fn this_ptr(&self) -> crate::ThisPtr<T> {
unsafe { crate::ThisPtr::new(self.0.as_ptr()) }
}
#[inline]
pub fn into_this_ptr(self) -> crate::ThisPtr<T> {
unsafe { crate::ThisPtr::new(self.into_raw()) }
}
#[inline]
pub fn as_ptr(&self) -> *mut T {
self.0.as_ptr()
}
#[inline]
pub fn as_non_null(&self) -> NonNull<T> {
self.0
}
}
impl<T: AnyRefCounted> Clone for RefPtr<T> {
#[inline]
fn clone(&self) -> Self {
unsafe { Self::init_ref(self.0.as_ptr()) }
}
}
impl<T: AnyRefCounted> core::ops::Deref for RefPtr<T> {
type Target = T;
#[inline]
fn deref(&self) -> &T {
unsafe { self.0.as_ref() }
}
}
#[cfg(debug_assertions)]
const MAGIC_VALID: u32 = 0x2f84_e51d;
#[cfg(debug_assertions)]
struct DebugData {
magic: u32,
}
#[cfg(debug_assertions)]
impl DebugData {
const fn empty() -> Self {
Self { magic: MAGIC_VALID }
}
fn assert_valid(&self) {
debug_assert!(
self.magic == MAGIC_VALID,
"ref/deref on a destroyed refcount"
);
}
fn deinit(&mut self) {
self.assert_valid();
self.magic = 0;
}
}
bun_core::declare_scope!(ref_count, hidden);
#[cfg(test)]
mod tests {
use super::*;
use core::sync::atomic::AtomicUsize;
use std::sync::{Mutex, MutexGuard, PoisonError};
#[allow(unused_imports)]
use bun_sys as _bun_sys_link;
static DROPS: AtomicUsize = AtomicUsize::new(0);
static SERIAL: Mutex<()> = Mutex::new(());
fn serial() -> MutexGuard<'static, ()> {
SERIAL.lock().unwrap_or_else(PoisonError::into_inner)
}
fn drops() -> usize {
DROPS.load(Ordering::SeqCst)
}
struct Thing {
ref_count: RefCount<Thing>,
payload: Box<u32>,
}
impl Thing {
fn new(payload: u32) -> *mut Thing {
bun_core::heap::into_raw(Box::new(Thing {
ref_count: RefCount::init(),
payload: Box::new(payload),
}))
}
}
impl Drop for Thing {
fn drop(&mut self) {
DROPS.fetch_add(1, Ordering::SeqCst);
}
}
impl RefCounted for Thing {
unsafe fn get_ref_count(this: *mut Self) -> *mut RefCount<Self> {
unsafe { &raw mut (*this).ref_count }
}
unsafe fn destructor(this: *mut Self) {
drop(unsafe { bun_core::heap::take(this) });
}
}
#[test]
fn ref_count_ref_deref_destroys_at_zero() {
let _serial = serial();
let before = drops();
let t = Thing::new(7);
unsafe {
RefCount::<Thing>::ref_(t);
assert_eq!((*RefCounted::get_ref_count(t)).get(), 2);
RefCount::<Thing>::deref(t);
assert!((*RefCounted::get_ref_count(t)).has_one_ref());
assert_eq!(*(*t).payload, 7);
RefCount::<Thing>::deref(t);
}
assert_eq!(drops(), before + 1);
}
#[test]
fn ref_count_init_exact_refs() {
let _serial = serial();
let before = drops();
let t = bun_core::heap::into_raw(Box::new(Thing {
ref_count: RefCount::init_exact_refs(3),
payload: Box::new(1),
}));
unsafe {
RefCount::<Thing>::deref(t);
RefCount::<Thing>::deref(t);
assert_eq!(drops(), before);
RefCount::<Thing>::deref(t);
}
assert_eq!(drops(), before + 1);
}
#[test]
fn ref_ptr_round_trip() {
let _serial = serial();
let before = drops();
let p = RefPtr::new(Thing {
ref_count: RefCount::init(),
payload: Box::new(42),
});
assert_eq!(*p.payload, 42);
let q = p.clone();
let r = p.clone();
assert_eq!(p.as_ptr(), q.as_ptr());
assert_eq!(p.as_ptr(), r.as_ptr());
drop(r);
drop(q);
assert_eq!(drops(), before);
let raw = p.into_raw();
let p = unsafe { RefPtr::from_raw(raw) };
assert_eq!(*p.payload, 42);
drop(p);
assert_eq!(drops(), before + 1);
}
#[test]
fn ref_ptr_init_ref_releases_on_drop() {
let _serial = serial();
let before = drops();
let t = Thing::new(3);
{
let _guard = unsafe { RefPtr::init_ref(t) };
assert_eq!(unsafe { (*RefCounted::get_ref_count(t)).get() }, 2);
}
assert!(unsafe { (*RefCounted::get_ref_count(t)).has_one_ref() });
unsafe { RefCount::<Thing>::deref(t) };
assert_eq!(drops(), before + 1);
}
#[test]
fn ref_ptr_from_raw_consumes_caller_ref() {
let _serial = serial();
let before = drops();
let t = Thing::new(3);
drop(unsafe { RefPtr::from_raw(t) });
assert_eq!(drops(), before + 1);
}
struct Shared {
ref_count: ThreadSafeRefCount<Shared>,
payload: Box<u32>,
}
impl Drop for Shared {
fn drop(&mut self) {
DROPS.fetch_add(1, Ordering::SeqCst);
}
}
impl ThreadSafeRefCounted for Shared {
unsafe fn get_ref_count(this: *mut Self) -> *mut ThreadSafeRefCount<Self> {
unsafe { &raw mut (*this).ref_count }
}
}
#[derive(Clone, Copy)]
struct SendPtr(*mut Shared);
unsafe impl Send for SendPtr {}
#[test]
fn thread_safe_ref_count_cross_thread_destroy() {
let _serial = serial();
let before = drops();
let s = bun_core::heap::into_raw(Box::new(Shared {
ref_count: ThreadSafeRefCount::init(),
payload: Box::new(5),
}));
const N: usize = 4;
let mut handles = Vec::with_capacity(N);
for _ in 0..N {
unsafe { ThreadSafeRefCount::<Shared>::ref_(s) };
let p = SendPtr(s);
handles.push(std::thread::spawn(move || {
let p = p;
unsafe {
assert_eq!(*(*p.0).payload, 5);
ThreadSafeRefCount::<Shared>::deref(p.0);
}
}));
}
for h in handles {
h.join().unwrap();
}
assert_eq!(drops(), before);
unsafe { ThreadSafeRefCount::<Shared>::deref(s) };
assert_eq!(drops(), before + 1);
}
#[test]
fn thread_safe_release_defers_destruction() {
let _serial = serial();
let before = drops();
let s = bun_core::heap::into_raw(Box::new(Shared {
ref_count: ThreadSafeRefCount::init_exact_refs(2),
payload: Box::new(8),
}));
unsafe {
assert!(!ThreadSafeRefCount::<Shared>::release(s));
assert_eq!(drops(), before);
assert!(ThreadSafeRefCount::<Shared>::release(s));
assert_eq!(drops(), before);
drop(bun_core::heap::take(s));
}
assert_eq!(drops(), before + 1);
}
struct Light {
ref_count: Cell<u32>,
payload: Box<u32>,
}
impl Drop for Light {
fn drop(&mut self) {
DROPS.fetch_add(1, Ordering::SeqCst);
}
}
unsafe impl CellRefCounted for Light {
fn ref_count(&self) -> &Cell<u32> {
&self.ref_count
}
unsafe fn ref_count_raw<'a>(this: *const Self) -> &'a Cell<u32> {
unsafe { &*(&raw const (*this).ref_count) }
}
}
#[test]
fn cell_ref_counted_destroys_at_zero() {
let _serial = serial();
let before = drops();
let l = bun_core::heap::into_raw(Box::new(Light {
ref_count: Cell::new(1),
payload: Box::new(11),
}));
unsafe {
(*l).ref_();
assert_eq!((*l).ref_count.get(), 2);
CellRefCounted::deref(l);
assert_eq!(*(*l).payload, 11);
assert_eq!(drops(), before);
CellRefCounted::deref(l);
}
assert_eq!(drops(), before + 1);
}
#[test]
fn type_base_name_strips_module_path() {
assert_eq!(type_base_name("a::b::Foo"), "Foo");
assert_eq!(type_base_name("a::b::Foo<c::Bar>"), "Foo<c::Bar>");
assert_eq!(type_base_name("Foo"), "Foo");
}
}