Skip to content

Commit 91cc3b0

Browse files
QelxirosQelxiros
committed
add Unique{Rc,Arc}::(try_)map
1 parent cdb8d49 commit 91cc3b0

File tree

2 files changed

+315
-2
lines changed

2 files changed

+315
-2
lines changed

‎library/alloc/src/rc.rs‎

Lines changed: 157 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -687,7 +687,7 @@ impl<T> Rc<T> {
687687

688688
/// Attempts to map the value in an `Rc`, reusing the allocation if possible.
689689
///
690-
/// `f` is called on a reference to the value in the box, and if the operation succeeds, the
690+
/// `f` is called on a reference to the value in the `Rc`, and if the operation succeeds, the
691691
/// result is returned, also in an `Rc`.
692692
///
693693
/// Note: this is an associated function, which means that you have
@@ -4093,6 +4093,126 @@ impl<T> UniqueRc<T> {
40934093
pub fn new(value: T) -> Self {
40944094
Self::new_in(value, Global)
40954095
}
4096+
4097+
/// Maps the value in a `UniqueRc`, reusing the allocation if possible.
4098+
///
4099+
/// `f` is called on a reference to the value in the `UniqueRc`, and the result is returned,
4100+
/// also in a `UniqueRc`.
4101+
///
4102+
/// Note: this is an associated function, which means that you have
4103+
/// to call it as `UniqueRc::map(&u, f)` instead of `u.map(f)`. This
4104+
/// is so that there is no conflict with a method on the inner type.
4105+
///
4106+
/// # Examples
4107+
///
4108+
/// ```
4109+
/// #![feature(smart_pointer_try_map)]
4110+
/// #![feature(unique_rc_arc)]
4111+
///
4112+
/// use std::rc::UniqueRc;
4113+
///
4114+
/// let r = UniqueRc::new(7);
4115+
/// let new = UniqueRc::map(r, |i| i + 7);
4116+
/// assert_eq!(*new, 14);
4117+
/// ```
4118+
#[cfg(not(no_global_oom_handling))]
4119+
#[unstable(feature = "smart_pointer_try_map", issue = "144419")]
4120+
pub fn map<U>(this: Self, f: impl FnOnce(T) -> U) -> UniqueRc<U> {
4121+
if size_of::<T>() == size_of::<U>()
4122+
&& align_of::<T>() == align_of::<U>()
4123+
&& UniqueRc::weak_count(&this) == 0
4124+
{
4125+
unsafe {
4126+
let ptr = UniqueRc::into_raw(this);
4127+
let value = ptr.read();
4128+
let mut allocation = UniqueRc::from_raw(ptr.cast::<MaybeUninit<U>>());
4129+
4130+
allocation.write(f(value));
4131+
allocation.assume_init()
4132+
}
4133+
} else {
4134+
UniqueRc::new(f(UniqueRc::unwrap(this)))
4135+
}
4136+
}
4137+
4138+
/// Attempts to map the value in a `UniqueRc`, reusing the allocation if possible.
4139+
///
4140+
/// `f` is called on a reference to the value in the `UniqueRc`, and if the operation succeeds,
4141+
/// the result is returned, also in a `UniqueRc`.
4142+
///
4143+
/// Note: this is an associated function, which means that you have
4144+
/// to call it as `UniqueRc::try_map(&u, f)` instead of `u.try_map(f)`. This
4145+
/// is so that there is no conflict with a method on the inner type.
4146+
///
4147+
/// # Examples
4148+
///
4149+
/// ```
4150+
/// #![feature(smart_pointer_try_map)]
4151+
/// #![feature(unique_rc_arc)]
4152+
///
4153+
/// use std::rc::UniqueRc;
4154+
///
4155+
/// let b = UniqueRc::new(7);
4156+
/// let new = UniqueRc::try_map(b, u32::try_from).unwrap();
4157+
/// assert_eq!(*new, 7);
4158+
/// ```
4159+
#[cfg(not(no_global_oom_handling))]
4160+
#[unstable(feature = "smart_pointer_try_map", issue = "144419")]
4161+
pub fn try_map<R>(
4162+
this: Self,
4163+
f: impl FnOnce(T) -> R,
4164+
) -> <R::Residual as Residual<UniqueRc<R::Output>>>::TryType
4165+
where
4166+
R: Try,
4167+
R::Residual: Residual<UniqueRc<R::Output>>,
4168+
{
4169+
if size_of::<T>() == size_of::<R::Output>()
4170+
&& align_of::<T>() == align_of::<R::Output>()
4171+
&& UniqueRc::weak_count(&this) == 0
4172+
{
4173+
unsafe {
4174+
let ptr = UniqueRc::into_raw(this);
4175+
let value = ptr.read();
4176+
let mut allocation = UniqueRc::from_raw(ptr.cast::<MaybeUninit<R::Output>>());
4177+
4178+
allocation.write(f(value)?);
4179+
try { allocation.assume_init() }
4180+
}
4181+
} else {
4182+
try { UniqueRc::new(f(UniqueRc::unwrap(this))?) }
4183+
}
4184+
}
4185+
4186+
fn unwrap(this: Self) -> T {
4187+
let this = ManuallyDrop::new(this);
4188+
let val: T = unsafe { ptr::read(&**this) };
4189+
4190+
this.inner().dec_strong();
4191+
let _weak = Weak { ptr: this.ptr, alloc: Global };
4192+
4193+
val
4194+
}
4195+
}
4196+
4197+
impl<T: ?Sized> UniqueRc<T> {
4198+
unsafe fn from_raw(ptr: *const T) -> Self {
4199+
let offset = unsafe { data_offset(ptr) };
4200+
4201+
// Reverse the offset to find the original RcInner.
4202+
let rc_ptr = unsafe { ptr.byte_sub(offset) as *mut RcInner<T> };
4203+
4204+
Self {
4205+
ptr: unsafe { NonNull::new_unchecked(rc_ptr) },
4206+
_marker: PhantomData,
4207+
_marker2: PhantomData,
4208+
alloc: Global,
4209+
}
4210+
}
4211+
4212+
fn into_raw(this: Self) -> *const T {
4213+
let this = ManuallyDrop::new(this);
4214+
Self::as_ptr(&*this)
4215+
}
40964216
}
40974217

40984218
impl<T, A: Allocator> UniqueRc<T, A> {
@@ -4143,6 +4263,35 @@ impl<T: ?Sized, A: Allocator> UniqueRc<T, A> {
41434263
Rc::from_inner_in(this.ptr, alloc)
41444264
}
41454265
}
4266+
4267+
fn weak_count(this: &Self) -> usize {
4268+
this.inner().weak() - 1
4269+
}
4270+
4271+
fn inner(&self) -> &RcInner<T> {
4272+
// SAFETY: while this UniqueRc is alive we're guaranteed that the inner pointer is valid.
4273+
unsafe { self.ptr.as_ref() }
4274+
}
4275+
4276+
fn as_ptr(this: &Self) -> *const T {
4277+
let ptr: *mut RcInner<T> = NonNull::as_ptr(this.ptr);
4278+
4279+
// SAFETY: This cannot go through Deref::deref or UniqueRc::inner because
4280+
// this is required to retain raw/mut provenance such that e.g. `get_mut` can
4281+
// write through the pointer after the Rc is recovered through `from_raw`.
4282+
unsafe { &raw mut (*ptr).value }
4283+
}
4284+
4285+
#[inline]
4286+
fn into_inner_with_allocator(this: Self) -> (NonNull<RcInner<T>>, A) {
4287+
let this = mem::ManuallyDrop::new(this);
4288+
(this.ptr, unsafe { ptr::read(&this.alloc) })
4289+
}
4290+
4291+
#[inline]
4292+
unsafe fn from_inner_in(ptr: NonNull<RcInner<T>>, alloc: A) -> Self {
4293+
Self { ptr, _marker: PhantomData, _marker2: PhantomData, alloc }
4294+
}
41464295
}
41474296

41484297
impl<T: ?Sized, A: Allocator + Clone> UniqueRc<T, A> {
@@ -4161,6 +4310,13 @@ impl<T: ?Sized, A: Allocator + Clone> UniqueRc<T, A> {
41614310
}
41624311
}
41634312

4313+
impl<T, A: Allocator> UniqueRc<MaybeUninit<T>, A> {
4314+
unsafe fn assume_init(self) -> UniqueRc<T, A> {
4315+
let (ptr, alloc) = UniqueRc::into_inner_with_allocator(self);
4316+
unsafe { UniqueRc::from_inner_in(ptr.cast(), alloc) }
4317+
}
4318+
}
4319+
41644320
#[unstable(feature = "unique_rc_arc", issue = "112566")]
41654321
impl<T: ?Sized, A: Allocator> Deref for UniqueRc<T, A> {
41664322
type Target = T;

‎library/alloc/src/sync.rs‎

Lines changed: 158 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -693,7 +693,7 @@ impl<T> Arc<T> {
693693

694694
/// Attempts to map the value in an `Arc`, reusing the allocation if possible.
695695
///
696-
/// `f` is called on a reference to the value in the box, and if the operation succeeds, the
696+
/// `f` is called on a reference to the value in the `Arc`, and if the operation succeeds, the
697697
/// result is returned, also in an `Arc`.
698698
///
699699
/// Note: this is an associated function, which means that you have
@@ -4505,6 +4505,127 @@ impl<T> UniqueArc<T, Global> {
45054505
pub fn new(value: T) -> Self {
45064506
Self::new_in(value, Global)
45074507
}
4508+
4509+
/// Maps the value in a `UniqueArc`, reusing the allocation if possible.
4510+
///
4511+
/// `f` is called on a reference to the value in the `UniqueArc`, and the result is returned,
4512+
/// also in a `UniqueArc`.
4513+
///
4514+
/// Note: this is an associated function, which means that you have
4515+
/// to call it as `UniqueArc::map(&u, f)` instead of `u.map(f)`. This
4516+
/// is so that there is no conflict with a method on the inner type.
4517+
///
4518+
/// # Examples
4519+
///
4520+
/// ```
4521+
/// #![feature(smart_pointer_try_map)]
4522+
/// #![feature(unique_rc_arc)]
4523+
///
4524+
/// use std::sync::UniqueArc;
4525+
///
4526+
/// let r = UniqueArc::new(7);
4527+
/// let new = UniqueArc::map(r, |i| i + 7);
4528+
/// assert_eq!(*new, 14);
4529+
/// ```
4530+
#[cfg(not(no_global_oom_handling))]
4531+
#[unstable(feature = "smart_pointer_try_map", issue = "144419")]
4532+
pub fn map<U>(this: Self, f: impl FnOnce(T) -> U) -> UniqueArc<U> {
4533+
if size_of::<T>() == size_of::<U>()
4534+
&& align_of::<T>() == align_of::<U>()
4535+
&& UniqueArc::weak_count(&this) == 0
4536+
{
4537+
unsafe {
4538+
let ptr = UniqueArc::into_raw(this);
4539+
let value = ptr.read();
4540+
let mut allocation = UniqueArc::from_raw(ptr.cast::<mem::MaybeUninit<U>>());
4541+
4542+
allocation.write(f(value));
4543+
allocation.assume_init()
4544+
}
4545+
} else {
4546+
UniqueArc::new(f(UniqueArc::unwrap(this)))
4547+
}
4548+
}
4549+
4550+
/// Attempts to map the value in a `UniqueArc`, reusing the allocation if possible.
4551+
///
4552+
/// `f` is called on a reference to the value in the `UniqueArc`, and if the operation succeeds,
4553+
/// the result is returned, also in a `UniqueArc`.
4554+
///
4555+
/// Note: this is an associated function, which means that you have
4556+
/// to call it as `UniqueArc::try_map(&u, f)` instead of `u.try_map(f)`. This
4557+
/// is so that there is no conflict with a method on the inner type.
4558+
///
4559+
/// # Examples
4560+
///
4561+
/// ```
4562+
/// #![feature(smart_pointer_try_map)]
4563+
/// #![feature(unique_rc_arc)]
4564+
///
4565+
/// use std::sync::UniqueArc;
4566+
///
4567+
/// let b = UniqueArc::new(7);
4568+
/// let new = UniqueArc::try_map(b, u32::try_from).unwrap();
4569+
/// assert_eq!(*new, 7);
4570+
/// ```
4571+
#[cfg(not(no_global_oom_handling))]
4572+
#[unstable(feature = "smart_pointer_try_map", issue = "144419")]
4573+
pub fn try_map<R>(
4574+
this: Self,
4575+
f: impl FnOnce(T) -> R,
4576+
) -> <R::Residual as Residual<UniqueArc<R::Output>>>::TryType
4577+
where
4578+
R: Try,
4579+
R::Residual: Residual<UniqueArc<R::Output>>,
4580+
{
4581+
if size_of::<T>() == size_of::<R::Output>()
4582+
&& align_of::<T>() == align_of::<R::Output>()
4583+
&& UniqueArc::weak_count(&this) == 0
4584+
{
4585+
unsafe {
4586+
let ptr = UniqueArc::into_raw(this);
4587+
let value = ptr.read();
4588+
let mut allocation = UniqueArc::from_raw(ptr.cast::<mem::MaybeUninit<R::Output>>());
4589+
4590+
allocation.write(f(value)?);
4591+
try { allocation.assume_init() }
4592+
}
4593+
} else {
4594+
try { UniqueArc::new(f(UniqueArc::unwrap(this))?) }
4595+
}
4596+
}
4597+
4598+
fn unwrap(this: Self) -> T {
4599+
let this = ManuallyDrop::new(this);
4600+
let val: T = unsafe { ptr::read(&**this) };
4601+
4602+
// We know there's only one strong reference, so the ordering doesn't matter.
4603+
this.inner().strong.fetch_sub(1, Relaxed);
4604+
let _weak = Weak { ptr: this.ptr, alloc: Global };
4605+
4606+
val
4607+
}
4608+
}
4609+
4610+
impl<T: ?Sized> UniqueArc<T> {
4611+
unsafe fn from_raw(ptr: *const T) -> Self {
4612+
let offset = unsafe { data_offset(ptr) };
4613+
4614+
// Reverse the offset to find the original ArcInner.
4615+
let rc_ptr = unsafe { ptr.byte_sub(offset) as *mut ArcInner<T> };
4616+
4617+
Self {
4618+
ptr: unsafe { NonNull::new_unchecked(rc_ptr) },
4619+
_marker: PhantomData,
4620+
_marker2: PhantomData,
4621+
alloc: Global,
4622+
}
4623+
}
4624+
4625+
fn into_raw(this: Self) -> *const T {
4626+
let this = ManuallyDrop::new(this);
4627+
Self::as_ptr(&*this)
4628+
}
45084629
}
45094630

45104631
impl<T, A: Allocator> UniqueArc<T, A> {
@@ -4558,6 +4679,35 @@ impl<T: ?Sized, A: Allocator> UniqueArc<T, A> {
45584679
Arc::from_inner_in(this.ptr, alloc)
45594680
}
45604681
}
4682+
4683+
fn weak_count(this: &Self) -> usize {
4684+
this.inner().weak.load(Acquire) - 1
4685+
}
4686+
4687+
fn inner(&self) -> &ArcInner<T> {
4688+
// SAFETY: while this UniqueArc is alive we're guaranteed that the inner pointer is valid.
4689+
unsafe { self.ptr.as_ref() }
4690+
}
4691+
4692+
fn as_ptr(this: &Self) -> *const T {
4693+
let ptr: *mut ArcInner<T> = NonNull::as_ptr(this.ptr);
4694+
4695+
// SAFETY: This cannot go through Deref::deref or UniqueArc::inner because
4696+
// this is required to retain raw/mut provenance such that e.g. `get_mut` can
4697+
// write through the pointer after the Rc is recovered through `from_raw`.
4698+
unsafe { &raw mut (*ptr).data }
4699+
}
4700+
4701+
#[inline]
4702+
fn into_inner_with_allocator(this: Self) -> (NonNull<ArcInner<T>>, A) {
4703+
let this = mem::ManuallyDrop::new(this);
4704+
(this.ptr, unsafe { ptr::read(&this.alloc) })
4705+
}
4706+
4707+
#[inline]
4708+
unsafe fn from_inner_in(ptr: NonNull<ArcInner<T>>, alloc: A) -> Self {
4709+
Self { ptr, _marker: PhantomData, _marker2: PhantomData, alloc }
4710+
}
45614711
}
45624712

45634713
impl<T: ?Sized, A: Allocator + Clone> UniqueArc<T, A> {
@@ -4588,6 +4738,13 @@ impl<T: ?Sized, A: Allocator + Clone> UniqueArc<T, A> {
45884738
}
45894739
}
45904740

4741+
impl<T, A: Allocator> UniqueArc<mem::MaybeUninit<T>, A> {
4742+
unsafe fn assume_init(self) -> UniqueArc<T, A> {
4743+
let (ptr, alloc) = UniqueArc::into_inner_with_allocator(self);
4744+
unsafe { UniqueArc::from_inner_in(ptr.cast(), alloc) }
4745+
}
4746+
}
4747+
45914748
#[unstable(feature = "unique_rc_arc", issue = "112566")]
45924749
impl<T: ?Sized, A: Allocator> Deref for UniqueArc<T, A> {
45934750
type Target = T;

0 commit comments

Comments
 (0)