Refactor

Author: EFanZh

library/alloc/src/raw_rc.rs

@@ -27,11 +27,8 @@ use crate::string::String;
 use crate::vec::Vec;
 
 pub unsafe trait RcOps {
-    unsafe fn inc_strong(strong_count: &UnsafeCell<usize>);
-    unsafe fn dec_strong(strong_count: &UnsafeCell<usize>) -> bool;
-
-    unsafe fn inc_weak(weak_count: &UnsafeCell<usize>);
-    unsafe fn dec_weak(weak_count: &UnsafeCell<usize>) -> bool;
+    unsafe fn increment_ref_count(count: &UnsafeCell<usize>);
+    unsafe fn decrement_ref_count(count: &UnsafeCell<usize>) -> bool;
 
     unsafe fn upgrade(strong_count: &UnsafeCell<usize>) -> bool;
     unsafe fn downgrade(weak_count: &UnsafeCell<usize>);
@@ -129,11 +126,11 @@ unsafe fn weak_count_ptr_from_value_ptr(value_ptr: NonNull<()>) -> NonNull<Unsaf
     unsafe { value_ptr.byte_sub(WEAK_COUNT_OFFSET_BYTES).cast() }
 }
 
-unsafe fn write_rc_allocation<const STRONG_COUNT: usize>(
-    ptr: NonNull<[u8]>,
+unsafe fn init_rc_allocation<const STRONG_COUNT: usize>(
+    allocation_ptr: NonNull<[u8]>,
     rc_layout: &RcLayout,
 ) -> NonNull<()> {
-    let allocation_ptr = ptr.cast::<()>();
+    let allocation_ptr = allocation_ptr.cast::<()>();
     let value_ptr = unsafe { allocation_ptr.byte_add(rc_layout.allocation_offset_bytes) };
     let ref_counts = const { RefCounts::new(STRONG_COUNT) };
 
@@ -142,100 +139,112 @@ unsafe fn write_rc_allocation<const STRONG_COUNT: usize>(
     value_ptr
 }
 
-fn try_allocate_for_rc<A, F, const STRONG_COUNT: usize>(
-    alloc: &A,
-    allocate_fn: F,
+unsafe fn try_handle_rc_allocation<const STRONG_COUNT: usize>(
+    allocation_result: Result<NonNull<[u8]>, AllocError>,
     rc_layout: &RcLayout,
-) -> Result<NonNull<()>, AllocError>
-where
-    F: FnOnce(&A, Layout) -> Result<NonNull<[u8]>, AllocError>,
-{
-    allocate_fn(alloc, rc_layout.allocation_layout)
-        .map(|ptr| unsafe { write_rc_allocation::<STRONG_COUNT>(ptr, rc_layout) })
+) -> Result<NonNull<()>, AllocError> {
+    allocation_result.map(|allocation_ptr| unsafe {
+        init_rc_allocation::<STRONG_COUNT>(allocation_ptr, rc_layout)
+    })
 }
 
-#[cfg(not(no_global_oom_handling))]
-#[track_caller]
-fn allocate_for_rc<A, F, const STRONG_COUNT: usize>(
+fn try_allocate_uninit_for_rc<A, const STRONG_COUNT: usize>(
     alloc: &A,
-    allocate_fn: F,
     rc_layout: &RcLayout,
-) -> NonNull<()>
+) -> Result<NonNull<()>, AllocError>
 where
-    F: FnOnce(&A, Layout) -> Result<NonNull<[u8]>, AllocError>,
+    A: Allocator,
 {
-    match allocate_fn(alloc, rc_layout.allocation_layout) {
-        Ok(ptr) => unsafe { write_rc_allocation::<STRONG_COUNT>(ptr, rc_layout) },
-        Err(AllocError) => alloc::handle_alloc_error(rc_layout.allocation_layout),
+    unsafe {
+        try_handle_rc_allocation::<STRONG_COUNT>(
+            alloc.allocate(rc_layout.allocation_layout),
+            rc_layout,
+        )
     }
 }
 
-#[cfg(not(no_global_oom_handling))]
-#[track_caller]
-unsafe fn allocate_for_rc_with_value<T, A, const STRONG_COUNT: usize>(
-    value: &T,
+fn try_allocate_zeroed_for_rc<A, const STRONG_COUNT: usize>(
     alloc: &A,
-) -> NonNull<T>
+    rc_layout: &RcLayout,
+) -> Result<NonNull<()>, AllocError>
 where
     A: Allocator,
-    T: ?Sized,
 {
     unsafe {
-        let rc_layout = RcLayout::from_value_ptr(NonNull::from(value));
-        let ptr = allocate_for_rc::<A, _, STRONG_COUNT>(alloc, A::allocate, &rc_layout);
-
-        ptr::copy_nonoverlapping::<u8>(
-            ptr::from_ref(value).cast(),
-            ptr.as_ptr().cast(),
-            mem::size_of_val(value),
-        );
-
-        NonNull::new_unchecked(ptr.as_ptr().with_metadata_of(value))
+        try_handle_rc_allocation::<STRONG_COUNT>(
+            alloc.allocate_zeroed(rc_layout.allocation_layout),
+            rc_layout,
+        )
     }
 }
 
-fn try_allocate_uninit_for_rc<A, const STRONG_COUNT: usize>(
-    alloc: &A,
+#[cfg(not(no_global_oom_handling))]
+#[track_caller]
+unsafe fn handle_rc_allocation<const STRONG_COUNT: usize>(
+    allocation_result: Result<NonNull<[u8]>, AllocError>,
     rc_layout: &RcLayout,
-) -> Result<NonNull<()>, AllocError>
-where
-    A: Allocator,
-{
-    try_allocate_for_rc::<A, _, STRONG_COUNT>(alloc, A::allocate, rc_layout)
+) -> NonNull<()> {
+    match allocation_result {
+        Ok(allocation_ptr) => unsafe {
+            init_rc_allocation::<STRONG_COUNT>(allocation_ptr, rc_layout)
+        },
+        Err(AllocError) => alloc::handle_alloc_error(rc_layout.allocation_layout),
+    }
 }
 
-fn try_allocate_zeroed_for_rc<A, const STRONG_COUNT: usize>(
+#[cfg(not(no_global_oom_handling))]
+#[track_caller]
+fn allocate_uninit_for_rc<A, const STRONG_COUNT: usize>(
     alloc: &A,
     rc_layout: &RcLayout,
-) -> Result<NonNull<()>, AllocError>
+) -> NonNull<()>
 where
     A: Allocator,
 {
-    try_allocate_for_rc::<_, _, STRONG_COUNT>(alloc, A::allocate_zeroed, rc_layout)
+    unsafe {
+        handle_rc_allocation::<STRONG_COUNT>(alloc.allocate(rc_layout.allocation_layout), rc_layout)
+    }
 }
 
 #[cfg(not(no_global_oom_handling))]
 #[track_caller]
-fn allocate_uninit_for_rc<A, const STRONG_COUNT: usize>(
+fn allocate_zeroed_for_rc<A, const STRONG_COUNT: usize>(
     alloc: &A,
     rc_layout: &RcLayout,
 ) -> NonNull<()>
 where
     A: Allocator,
 {
-    allocate_for_rc::<A, _, STRONG_COUNT>(alloc, A::allocate, rc_layout)
+    unsafe {
+        handle_rc_allocation::<STRONG_COUNT>(
+            alloc.allocate_zeroed(rc_layout.allocation_layout),
+            rc_layout,
+        )
+    }
 }
 
 #[cfg(not(no_global_oom_handling))]
 #[track_caller]
-fn allocate_zeroed_for_rc<A, const STRONG_COUNT: usize>(
+unsafe fn allocate_for_rc_with_value<T, A, const STRONG_COUNT: usize>(
+    value: &T,
     alloc: &A,
-    rc_layout: &RcLayout,
-) -> NonNull<()>
+) -> NonNull<T>
 where
     A: Allocator,
+    T: ?Sized,
 {
-    allocate_for_rc::<A, _, STRONG_COUNT>(alloc, A::allocate_zeroed, rc_layout)
+    unsafe {
+        let rc_layout = RcLayout::from_value_ptr(NonNull::from(value));
+        let ptr = allocate_uninit_for_rc::<A, STRONG_COUNT>(alloc, &rc_layout);
+
+        ptr::copy_nonoverlapping::<u8>(
+            ptr::from_ref(value).cast(),
+            ptr.as_ptr().cast(),
+            mem::size_of_val(value),
+        );
+
+        NonNull::new_unchecked(ptr.as_ptr().with_metadata_of(value))
+    }
 }
 
 struct GuardedWeak<'a, T, A, R>
@@ -370,7 +379,7 @@ where
     {
         unsafe {
             if !self.is_dangling() {
-                R::inc_weak(self.weak_count_unchecked());
+                R::increment_ref_count(self.weak_count_unchecked());
             }
 
             self.clone_without_inc_ref()
@@ -397,7 +406,7 @@ where
         R: RcOps,
     {
         unsafe {
-            if R::dec_weak(self.weak_count_unchecked()) {
+            if R::decrement_ref_count(self.weak_count_unchecked()) {
                 self.deallocate();
             }
         };
@@ -750,7 +759,7 @@ where
         R: RcOps,
     {
         unsafe {
-            R::inc_strong(self.strong_count());
+            R::increment_ref_count(self.strong_count());
 
             Self::from_weak(self.weak.clone_without_inc_ref())
         }
@@ -771,7 +780,7 @@ where
     }
 
     pub unsafe fn increment_strong_count<R: RcOps>(ptr: NonNull<T>) {
-        unsafe { R::inc_strong(strong_count_ptr_from_value_ptr(ptr.cast()).as_ref()) };
+        unsafe { R::increment_ref_count(strong_count_ptr_from_value_ptr(ptr.cast()).as_ref()) };
     }
 
     #[inline(never)]
@@ -789,7 +798,7 @@ where
         R: RcOps,
     {
         unsafe {
-            if R::dec_strong(self.strong_count()) {
+            if R::decrement_ref_count(self.strong_count()) {
                 self.drop_slow::<R>();
             }
         };
@@ -991,11 +1000,7 @@ impl<T, A> RawRc<T, A> {
 
         mem::forget(guard);
 
-        unsafe {
-            weak.as_ptr().write(data);
-
-            RawUniqueRc::from_weak(weak).into_rc::<R>()
-        }
+        unsafe { RawUniqueRc::from_weak_with_value(weak, data).into_rc::<R>() }
     }
 
     #[cfg(not(no_global_oom_handling))]
@@ -1015,7 +1020,8 @@ impl<T, A> RawRc<T, A> {
         R: RcOps,
     {
         unsafe {
-            R::dec_strong(&self.strong_count()).then(|| self.weak.assume_init_into_inner::<R>())
+            R::decrement_ref_count(&self.strong_count())
+                .then(|| self.weak.assume_init_into_inner::<R>())
         }
     }
 
@@ -1841,10 +1847,6 @@ impl<T, A> RawUniqueRc<T, A>
 where
     T: ?Sized,
 {
-    unsafe fn from_weak(weak: RawWeak<T, A>) -> Self {
-        Self { weak }
-    }
-
     pub unsafe fn downgrade<R>(&self) -> RawWeak<T, A>
     where
         A: Clone,
@@ -1876,11 +1878,9 @@ where
 impl<T, A> RawUniqueRc<T, A> {
     #[cfg(not(no_global_oom_handling))]
     unsafe fn from_weak_with_value(weak: RawWeak<T, A>, value: T) -> Self {
-        unsafe {
-            weak.as_ptr().write(value);
+        unsafe { weak.as_ptr().write(value) };
 
-            Self::from_weak(weak)
-        }
+        Self { weak }
     }
 
     #[cfg(not(no_global_oom_handling))]

library/alloc/src/rc.rs

@@ -279,53 +279,37 @@ where
     move |raw_weak: &RawWeak<T, A>| f(unsafe { mem::transmute(raw_weak) })
 }
 
-unsafe fn inc_ref(count: &UnsafeCell<usize>) {
-    let count = unsafe { &mut *count.get() };
-    let strong = *count;
-
-    // We insert an `assume` here to hint LLVM at an otherwise
-    // missed optimization.
-    // SAFETY: The reference count will never be zero when this is
-    // called.
-    unsafe { hint::assert_unchecked(strong != 0) };
-
-    let strong = count.wrapping_add(1);
-
-    *count = strong;
-
-    // We want to abort on overflow instead of dropping the value.
-    // Checking for overflow after the store instead of before
-    // allows for slightly better code generation.
-    if intrinsics::unlikely(strong == 0) {
-        intrinsics::abort();
-    }
-}
+enum RcOps {}
 
-unsafe fn dec_ref(count: &UnsafeCell<usize>) -> bool {
-    let count = unsafe { &mut *count.get() };
+unsafe impl raw_rc::RcOps for RcOps {
+    unsafe fn increment_ref_count(count: &UnsafeCell<usize>) {
+        let count = unsafe { &mut *count.get() };
+        let strong = *count;
 
-    *count -= 1;
+        // We insert an `assume` here to hint LLVM at an otherwise
+        // missed optimization.
+        // SAFETY: The reference count will never be zero when this is
+        // called.
+        unsafe { hint::assert_unchecked(strong != 0) };
 
-    *count == 0
-}
+        let strong = count.wrapping_add(1);
 
-enum RcOps {}
+        *count = strong;
 
-unsafe impl raw_rc::RcOps for RcOps {
-    unsafe fn inc_strong(strong_count: &UnsafeCell<usize>) {
-        unsafe { inc_ref(strong_count) };
+        // We want to abort on overflow instead of dropping the value.
+        // Checking for overflow after the store instead of before
+        // allows for slightly better code generation.
+        if intrinsics::unlikely(strong == 0) {
+            intrinsics::abort();
+        }
     }
 
-    unsafe fn dec_strong(strong_count: &UnsafeCell<usize>) -> bool {
-        unsafe { dec_ref(strong_count) }
-    }
+    unsafe fn decrement_ref_count(count: &UnsafeCell<usize>) -> bool {
+        let count = unsafe { &mut *count.get() };
 
-    unsafe fn inc_weak(weak_count: &UnsafeCell<usize>) {
-        unsafe { inc_ref(weak_count) };
-    }
+        *count -= 1;
 
-    unsafe fn dec_weak(weak_count: &UnsafeCell<usize>) -> bool {
-        unsafe { dec_ref(weak_count) }
+        *count == 0
     }
 
     unsafe fn upgrade(strong_count: &UnsafeCell<usize>) -> bool {
@@ -341,7 +325,7 @@ unsafe impl raw_rc::RcOps for RcOps {
     }
 
     unsafe fn downgrade(weak_count: &UnsafeCell<usize>) {
-        unsafe { inc_ref(weak_count) };
+        unsafe { Self::increment_ref_count(weak_count) };
     }
 
     unsafe fn lock_strong_count(strong_count: &UnsafeCell<usize>) -> bool {