Skip to content

Commit 85fdf86

Browse files
committed
Add RawUniqueRc methods for sized values
1 parent cd4f99f commit 85fdf86

File tree

2 files changed

+140
-0
lines changed

2 files changed

+140
-0
lines changed

library/alloc/src/raw_rc/raw_rc.rs

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ use crate::boxed::Box;
2929
use crate::raw_rc::MakeMutStrategy;
3030
use crate::raw_rc::RefCounter;
3131
#[cfg(not(no_global_oom_handling))]
32+
use crate::raw_rc::raw_unique_rc::RawUniqueRc;
33+
#[cfg(not(no_global_oom_handling))]
3234
use crate::raw_rc::raw_weak;
3335
use crate::raw_rc::raw_weak::RawWeak;
3436
#[cfg(not(no_global_oom_handling))]
@@ -412,6 +414,45 @@ impl<T, A> RawRc<T, A> {
412414
unsafe { Self::from_raw_parts(ptr.as_ptr().cast(), alloc) }
413415
}
414416

417+
#[cfg(not(no_global_oom_handling))]
418+
unsafe fn new_cyclic_impl<F, R>(mut weak: RawWeak<T, A>, data_fn: F) -> Self
419+
where
420+
A: Allocator,
421+
F: FnOnce(&RawWeak<T, A>) -> T,
422+
R: RefCounter,
423+
{
424+
let guard = unsafe { raw_weak::new_weak_guard::<T, A, R>(&mut weak) };
425+
let data = data_fn(&guard);
426+
427+
mem::forget(guard);
428+
429+
unsafe { RawUniqueRc::from_weak_with_value(weak, data).into_rc::<R>() }
430+
}
431+
432+
#[cfg(not(no_global_oom_handling))]
433+
pub(crate) unsafe fn new_cyclic<F, R>(data_fn: F) -> Self
434+
where
435+
A: Allocator + Default,
436+
F: FnOnce(&RawWeak<T, A>) -> T,
437+
R: RefCounter,
438+
{
439+
let weak = RawWeak::new_uninit::<0>();
440+
441+
unsafe { Self::new_cyclic_impl::<F, R>(weak, data_fn) }
442+
}
443+
444+
#[cfg(not(no_global_oom_handling))]
445+
pub(crate) unsafe fn new_cyclic_in<F, R>(data_fn: F, alloc: A) -> Self
446+
where
447+
A: Allocator,
448+
F: FnOnce(&RawWeak<T, A>) -> T,
449+
R: RefCounter,
450+
{
451+
let weak = RawWeak::new_uninit_in::<0>(alloc);
452+
453+
unsafe { Self::new_cyclic_impl::<F, R>(weak, data_fn) }
454+
}
455+
415456
/// Maps the value in an `RawRc`, reusing the allocation if possible.
416457
///
417458
/// # Safety

library/alloc/src/raw_rc/raw_unique_rc.rs

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,14 @@
11
use core::alloc::Allocator;
22
use core::marker::PhantomData;
3+
#[cfg(not(no_global_oom_handling))]
4+
use core::mem::{self, SizedTypeProperties};
5+
#[cfg(not(no_global_oom_handling))]
6+
use core::ops::{ControlFlow, Try};
37

48
use crate::raw_rc::RefCounter;
59
use crate::raw_rc::raw_rc::RawRc;
10+
#[cfg(not(no_global_oom_handling))]
11+
use crate::raw_rc::raw_weak;
612
use crate::raw_rc::raw_weak::RawWeak;
713
use crate::raw_rc::rc_value_pointer::RcValuePointer;
814

@@ -73,3 +79,96 @@ where
7379
}
7480
}
7581
}
82+
83+
impl<T, A> RawUniqueRc<T, A> {
84+
#[cfg(not(no_global_oom_handling))]
85+
pub(super) unsafe fn from_weak_with_value(weak: RawWeak<T, A>, value: T) -> Self {
86+
unsafe { weak.as_ptr().write(value) };
87+
88+
Self { weak, _marker: PhantomData, _marker2: PhantomData }
89+
}
90+
91+
#[cfg(not(no_global_oom_handling))]
92+
pub(crate) fn new(value: T) -> Self
93+
where
94+
A: Allocator + Default,
95+
{
96+
unsafe { Self::from_weak_with_value(RawWeak::new_uninit::<0>(), value) }
97+
}
98+
99+
#[cfg(not(no_global_oom_handling))]
100+
pub(crate) fn new_in(value: T, alloc: A) -> Self
101+
where
102+
A: Allocator,
103+
{
104+
unsafe { Self::from_weak_with_value(RawWeak::new_uninit_in::<0>(alloc), value) }
105+
}
106+
107+
/// Maps the value in a `RawUniqueRc`, reusing the allocation if possible.
108+
///
109+
/// # Safety
110+
///
111+
/// All accesses to `self` must use the same `RefCounter` implementation for `R`.
112+
#[cfg(not(no_global_oom_handling))]
113+
pub(crate) unsafe fn map<R, U>(self, f: impl FnOnce(T) -> U) -> RawUniqueRc<U, A>
114+
where
115+
A: Allocator,
116+
R: RefCounter,
117+
{
118+
fn wrap_fn<T, U>(f: impl FnOnce(T) -> U) -> impl FnOnce(T) -> ControlFlow<!, U> {
119+
|x| ControlFlow::Continue(f(x))
120+
}
121+
122+
let f = wrap_fn(f);
123+
124+
match unsafe { self.try_map::<R, _>(f) } {
125+
ControlFlow::Continue(output) => output,
126+
}
127+
}
128+
129+
/// Attempts to map the value in a `RawUniqueRc`, reusing the allocation if possible.
130+
///
131+
/// # Safety
132+
///
133+
/// All accesses to `self` must use the same `RefCounter` implementation for `R`.
134+
#[cfg(not(no_global_oom_handling))]
135+
pub(crate) unsafe fn try_map<R, U>(
136+
mut self,
137+
f: impl FnOnce(T) -> U,
138+
) -> ControlFlow<U::Residual, RawUniqueRc<U::Output, A>>
139+
where
140+
A: Allocator,
141+
R: RefCounter,
142+
U: Try,
143+
{
144+
// Destruct `self` as a `RawWeak<T, A>` if `f` panics or returns a failure value.
145+
let guard = unsafe { raw_weak::new_weak_guard::<T, A, R>(&mut self.weak) };
146+
147+
let (allocation, mapped_value) = if T::LAYOUT == U::Output::LAYOUT
148+
&& R::unique_rc_weak_count(unsafe { R::from_raw_counter(guard.weak_count_unchecked()) })
149+
== 1
150+
{
151+
let mapped_value = f(unsafe { guard.as_ptr().read() }).branch()?;
152+
153+
// Avoid deallocation on success, reuse the allocation.
154+
mem::forget(guard);
155+
156+
let allocation = unsafe { self.weak.cast() };
157+
158+
(allocation, mapped_value)
159+
} else {
160+
let value = unsafe { guard.as_ptr().read() };
161+
162+
drop(guard);
163+
164+
let mapped_value = f(value).branch()?;
165+
let allocation = RawWeak::new_uninit_in::<0>(self.weak.into_raw_parts().1);
166+
167+
(allocation, mapped_value)
168+
};
169+
170+
ControlFlow::Continue(unsafe {
171+
RawUniqueRc::from_weak_with_value(allocation, mapped_value)
172+
})
173+
}
174+
}

0 commit comments

Comments
 (0)