add Unique{Rc,Arc}::(try_)map

Author: Qelxiros

library/alloc/src/rc.rs

@@ -252,8 +252,6 @@ use core::intrinsics::abort;
 #[cfg(not(no_global_oom_handling))]
 use core::iter;
 use core::marker::{PhantomData, Unsize};
-#[cfg(not(no_global_oom_handling))]
-use core::mem::MaybeUninit;
 use core::mem::{self, ManuallyDrop, align_of_val_raw};
 use core::num::NonZeroUsize;
 use core::ops::{CoerceUnsized, Deref, DerefMut, DerefPure, DispatchFromDyn, LegacyReceiver};
@@ -651,7 +649,7 @@ impl<T> Rc<T> {
     /// an `Rc`.
     ///
     /// Note: this is an associated function, which means that you have
-    /// to call it as `Rc::map(&r, f)` instead of `r.map(f)`. This
+    /// to call it as `Rc::map(r, f)` instead of `r.map(f)`. This
     /// is so that there is no conflict with a method on the inner type.
     ///
     /// # Examples
@@ -675,7 +673,7 @@ impl<T> Rc<T> {
             unsafe {
                 let ptr = Rc::into_raw(this);
                 let value = ptr.read();
-                let mut allocation = Rc::from_raw(ptr.cast::<MaybeUninit<U>>());
+                let mut allocation = Rc::from_raw(ptr.cast::<mem::MaybeUninit<U>>());
 
                 Rc::get_mut_unchecked(&mut allocation).write(f(&value));
                 allocation.assume_init()
@@ -687,11 +685,11 @@ impl<T> Rc<T> {
 
     /// Attempts to map the value in an `Rc`, reusing the allocation if possible.
     ///
-    /// `f` is called on a reference to the value in the box, and if the operation succeeds, the
+    /// `f` is called on a reference to the value in the `Rc`, and if the operation succeeds, the
     /// result is returned, also in an `Rc`.
     ///
     /// Note: this is an associated function, which means that you have
-    /// to call it as `Rc::try_map(&r, f)` instead of `r.try_map(f)`. This
+    /// to call it as `Rc::try_map(r, f)` instead of `r.try_map(f)`. This
     /// is so that there is no conflict with a method on the inner type.
     ///
     /// # Examples
@@ -722,7 +720,7 @@ impl<T> Rc<T> {
             unsafe {
                 let ptr = Rc::into_raw(this);
                 let value = ptr.read();
-                let mut allocation = Rc::from_raw(ptr.cast::<MaybeUninit<R::Output>>());
+                let mut allocation = Rc::from_raw(ptr.cast::<mem::MaybeUninit<R::Output>>());
 
                 Rc::get_mut_unchecked(&mut allocation).write(f(&value)?);
                 try { allocation.assume_init() }
@@ -4093,6 +4091,128 @@ impl<T> UniqueRc<T> {
     pub fn new(value: T) -> Self {
         Self::new_in(value, Global)
     }
+
+    /// Maps the value in a `UniqueRc`, reusing the allocation if possible.
+    ///
+    /// `f` is called on a reference to the value in the `UniqueRc`, and the result is returned,
+    /// also in a `UniqueRc`.
+    ///
+    /// Note: this is an associated function, which means that you have
+    /// to call it as `UniqueRc::map(u, f)` instead of `u.map(f)`. This
+    /// is so that there is no conflict with a method on the inner type.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(smart_pointer_try_map)]
+    /// #![feature(unique_rc_arc)]
+    ///
+    /// use std::rc::UniqueRc;
+    ///
+    /// let r = UniqueRc::new(7);
+    /// let new = UniqueRc::map(r, |i| i + 7);
+    /// assert_eq!(*new, 14);
+    /// ```
+    #[cfg(not(no_global_oom_handling))]
+    #[unstable(feature = "smart_pointer_try_map", issue = "144419")]
+    pub fn map<U>(this: Self, f: impl FnOnce(T) -> U) -> UniqueRc<U> {
+        if size_of::<T>() == size_of::<U>()
+            && align_of::<T>() == align_of::<U>()
+            && UniqueRc::weak_count(&this) == 0
+        {
+            unsafe {
+                let ptr = UniqueRc::into_raw(this);
+                let value = ptr.read();
+                let mut allocation = UniqueRc::from_raw(ptr.cast::<mem::MaybeUninit<U>>());
+
+                allocation.write(f(value));
+                allocation.assume_init()
+            }
+        } else {
+            UniqueRc::new(f(UniqueRc::unwrap(this)))
+        }
+    }
+
+    /// Attempts to map the value in a `UniqueRc`, reusing the allocation if possible.
+    ///
+    /// `f` is called on a reference to the value in the `UniqueRc`, and if the operation succeeds,
+    /// the result is returned, also in a `UniqueRc`.
+    ///
+    /// Note: this is an associated function, which means that you have
+    /// to call it as `UniqueRc::try_map(u, f)` instead of `u.try_map(f)`. This
+    /// is so that there is no conflict with a method on the inner type.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(smart_pointer_try_map)]
+    /// #![feature(unique_rc_arc)]
+    ///
+    /// use std::rc::UniqueRc;
+    ///
+    /// let b = UniqueRc::new(7);
+    /// let new = UniqueRc::try_map(b, u32::try_from).unwrap();
+    /// assert_eq!(*new, 7);
+    /// ```
+    #[cfg(not(no_global_oom_handling))]
+    #[unstable(feature = "smart_pointer_try_map", issue = "144419")]
+    pub fn try_map<R>(
+        this: Self,
+        f: impl FnOnce(T) -> R,
+    ) -> <R::Residual as Residual<UniqueRc<R::Output>>>::TryType
+    where
+        R: Try,
+        R::Residual: Residual<UniqueRc<R::Output>>,
+    {
+        if size_of::<T>() == size_of::<R::Output>()
+            && align_of::<T>() == align_of::<R::Output>()
+            && UniqueRc::weak_count(&this) == 0
+        {
+            unsafe {
+                let ptr = UniqueRc::into_raw(this);
+                let value = ptr.read();
+                let mut allocation = UniqueRc::from_raw(ptr.cast::<mem::MaybeUninit<R::Output>>());
+
+                allocation.write(f(value)?);
+                try { allocation.assume_init() }
+            }
+        } else {
+            try { UniqueRc::new(f(UniqueRc::unwrap(this))?) }
+        }
+    }
+
+    #[cfg(not(no_global_oom_handling))]
+    fn unwrap(this: Self) -> T {
+        let this = ManuallyDrop::new(this);
+        let val: T = unsafe { ptr::read(&**this) };
+
+        let _weak = Weak { ptr: this.ptr, alloc: Global };
+
+        val
+    }
+}
+
+impl<T: ?Sized> UniqueRc<T> {
+    #[cfg(not(no_global_oom_handling))]
+    unsafe fn from_raw(ptr: *const T) -> Self {
+        let offset = unsafe { data_offset(ptr) };
+
+        // Reverse the offset to find the original RcInner.
+        let rc_ptr = unsafe { ptr.byte_sub(offset) as *mut RcInner<T> };
+
+        Self {
+            ptr: unsafe { NonNull::new_unchecked(rc_ptr) },
+            _marker: PhantomData,
+            _marker2: PhantomData,
+            alloc: Global,
+        }
+    }
+
+    #[cfg(not(no_global_oom_handling))]
+    fn into_raw(this: Self) -> *const T {
+        let this = ManuallyDrop::new(this);
+        Self::as_ptr(&*this)
+    }
 }
 
 impl<T, A: Allocator> UniqueRc<T, A> {
@@ -4143,6 +4263,40 @@ impl<T: ?Sized, A: Allocator> UniqueRc<T, A> {
             Rc::from_inner_in(this.ptr, alloc)
         }
     }
+
+    #[cfg(not(no_global_oom_handling))]
+    fn weak_count(this: &Self) -> usize {
+        this.inner().weak() - 1
+    }
+
+    #[cfg(not(no_global_oom_handling))]
+    fn inner(&self) -> &RcInner<T> {
+        // SAFETY: while this UniqueRc is alive we're guaranteed that the inner pointer is valid.
+        unsafe { self.ptr.as_ref() }
+    }
+
+    #[cfg(not(no_global_oom_handling))]
+    fn as_ptr(this: &Self) -> *const T {
+        let ptr: *mut RcInner<T> = NonNull::as_ptr(this.ptr);
+
+        // SAFETY: This cannot go through Deref::deref or UniqueRc::inner because
+        // this is required to retain raw/mut provenance such that e.g. `get_mut` can
+        // write through the pointer after the Rc is recovered through `from_raw`.
+        unsafe { &raw mut (*ptr).value }
+    }
+
+    #[inline]
+    #[cfg(not(no_global_oom_handling))]
+    fn into_inner_with_allocator(this: Self) -> (NonNull<RcInner<T>>, A) {
+        let this = mem::ManuallyDrop::new(this);
+        (this.ptr, unsafe { ptr::read(&this.alloc) })
+    }
+
+    #[inline]
+    #[cfg(not(no_global_oom_handling))]
+    unsafe fn from_inner_in(ptr: NonNull<RcInner<T>>, alloc: A) -> Self {
+        Self { ptr, _marker: PhantomData, _marker2: PhantomData, alloc }
+    }
 }
 
 impl<T: ?Sized, A: Allocator + Clone> UniqueRc<T, A> {
@@ -4161,6 +4315,14 @@ impl<T: ?Sized, A: Allocator + Clone> UniqueRc<T, A> {
     }
 }
 
+#[cfg(not(no_global_oom_handling))]
+impl<T, A: Allocator> UniqueRc<mem::MaybeUninit<T>, A> {
+    unsafe fn assume_init(self) -> UniqueRc<T, A> {
+        let (ptr, alloc) = UniqueRc::into_inner_with_allocator(self);
+        unsafe { UniqueRc::from_inner_in(ptr.cast(), alloc) }
+    }
+}
+
 #[unstable(feature = "unique_rc_arc", issue = "112566")]
 impl<T: ?Sized, A: Allocator> Deref for UniqueRc<T, A> {
     type Target = T;