Implement alloc::rc::{Rc,Weak,UniqueRc} with alloc::raw_rc types

Author: EFanZh

library/alloc/src/raw_rc/mod.rs

@@ -60,12 +60,14 @@
 //! - Converting `Option<Rc<T>>` to `Option<&T>` without checking for `None` values.
 //! - Converting `&[Rc<T>]` to `&[&T]` with zero cost.
 
-#![allow(dead_code)]
-
 use core::cell::UnsafeCell;
 use core::mem;
 use core::sync::atomic::Atomic;
 
+pub(crate) use crate::raw_rc::raw_rc::RawRc;
+pub(crate) use crate::raw_rc::raw_unique_rc::RawUniqueRc;
+pub(crate) use crate::raw_rc::raw_weak::RawWeak;
+
 mod raw_rc;
 mod raw_unique_rc;
 mod raw_weak;
@@ -169,4 +171,9 @@ pub(crate) unsafe trait RefCounter: Sized {
     /// - If strong count is greater than 1, returns `Some(MakeMutStrategy::Clone)`.
     #[cfg(not(no_global_oom_handling))]
     fn make_mut(strong_count: &Self, weak_count: &Self) -> Option<MakeMutStrategy>;
+
+    /// Returns the weak count of an `RawUniqueRc`, used to determine whether there are any weak
+    /// pointers to the same allocation.
+    #[cfg(not(no_global_oom_handling))]
+    fn unique_rc_weak_count(weak_count: &Self) -> usize;
 }

library/alloc/src/raw_rc/raw_unique_rc.rs

@@ -1,12 +1,20 @@
 use core::alloc::Allocator;
+#[cfg(not(no_global_oom_handling))]
+use core::convert;
 use core::fmt::{self, Debug, Display, Formatter, Pointer};
 use core::hash::{Hash, Hasher};
 use core::marker::{PhantomData, Unsize};
+#[cfg(not(no_global_oom_handling))]
+use core::mem::{self, SizedTypeProperties};
 use core::ops::{CoerceUnsized, DispatchFromDyn};
+#[cfg(not(no_global_oom_handling))]
+use core::ops::{Residual, Try};
 
 use crate::alloc::Global;
 use crate::raw_rc::RefCounter;
 use crate::raw_rc::raw_rc::RawRc;
+#[cfg(not(no_global_oom_handling))]
+use crate::raw_rc::raw_weak;
 use crate::raw_rc::raw_weak::RawWeak;
 use crate::raw_rc::rc_value_pointer::RcValuePointer;
 
@@ -101,6 +109,62 @@ impl<T, A> RawUniqueRc<T, A> {
     {
         unsafe { Self::from_weak_with_value(RawWeak::new_uninit_in::<0>(alloc), value) }
     }
+
+    /// Maps the value in a `RawUniqueRc`, reusing the allocation if possible.
+    #[cfg(not(no_global_oom_handling))]
+    pub(crate) fn map<R, U>(self, f: impl FnOnce(T) -> U) -> RawUniqueRc<U, A>
+    where
+        A: Allocator,
+        R: RefCounter,
+    {
+        fn wrap_fn<T, U>(f: impl FnOnce(T) -> U) -> impl FnOnce(T) -> Result<U, !> {
+            |x| Ok(f(x))
+        }
+
+        self.try_map::<R, _, _>(wrap_fn(f), convert::identity).into_ok()
+    }
+
+    /// Attempts to map the value in a `RawUniqueRc`, reusing the allocation if possible.
+    #[cfg(not(no_global_oom_handling))]
+    pub(crate) fn try_map<R, U, V>(
+        mut self,
+        f: impl FnOnce(T) -> U,
+        map_result: impl FnOnce(RawUniqueRc<U::Output, A>) -> V, // How to remove this argument?
+    ) -> <U::Residual as Residual<V>>::TryType
+    where
+        A: Allocator,
+        R: RefCounter,
+        U: Try,
+        U::Residual: Residual<V>,
+    {
+        // Destruct `self` as a `RawWeak<T, A>` if `f` panics or returns `Err`.
+        let guard = unsafe { raw_weak::new_weak_guard::<T, A, R>(&mut self.weak) };
+
+        let (allocation, mapped_value) = if T::LAYOUT == U::Output::LAYOUT
+            && R::unique_rc_weak_count(unsafe { R::from_raw_counter(guard.weak_count_unchecked()) })
+                == 1
+        {
+            let mapped_value = f(unsafe { guard.as_ptr().read() })?;
+
+            // Avoid deallocation on success, reuse the allocation.
+            mem::forget(guard);
+
+            let allocation = unsafe { self.weak.cast() };
+
+            (allocation, mapped_value)
+        } else {
+            let value = unsafe { guard.as_ptr().read() };
+
+            drop(guard);
+
+            let mapped_value = f(value)?;
+            let allocation = RawWeak::new_uninit_in::<0>(self.weak.into_raw_parts().1);
+
+            (allocation, mapped_value)
+        };
+
+        try { map_result(unsafe { RawUniqueRc::from_weak_with_value(allocation, mapped_value) }) }
+    }
 }
 
 impl<T, A> AsMut<T> for RawUniqueRc<T, A>