triomphe/
unique_arc.rs

1use alloc::vec::Vec;
2use alloc::{alloc::Layout, boxed::Box};
3use core::convert::TryFrom;
4use core::iter::FromIterator;
5use core::marker::PhantomData;
6use core::mem::{ManuallyDrop, MaybeUninit};
7use core::ops::{Deref, DerefMut};
8use core::ptr::{self, addr_of_mut, NonNull};
9use core::sync::atomic::AtomicUsize;
10
11use crate::iterator_as_exact_size_iterator::IteratorAsExactSizeIterator;
12use crate::HeaderSlice;
13
14use super::{Arc, ArcInner};
15
16/// An `Arc` that is known to be uniquely owned
17///
18/// When `Arc`s are constructed, they are known to be
19/// uniquely owned. In such a case it is safe to mutate
20/// the contents of the `Arc`. Normally, one would just handle
21/// this by mutating the data on the stack before allocating the
22/// `Arc`, however it's possible the data is large or unsized
23/// and you need to heap-allocate it earlier in such a way
24/// that it can be freely converted into a regular `Arc` once you're
25/// done.
26///
27/// `UniqueArc` exists for this purpose, when constructed it performs
28/// the same allocations necessary for an `Arc`, however it allows mutable access.
29/// Once the mutation is finished, you can call `.shareable()` and get a regular `Arc`
30/// out of it.
31///
32/// ```rust
33/// # use triomphe::UniqueArc;
34/// let data = [1, 2, 3, 4, 5];
35/// let mut x = UniqueArc::new(data);
36/// x[4] = 7; // mutate!
37/// let y = x.shareable(); // y is an Arc<T>
38/// ```
39#[repr(transparent)]
40pub struct UniqueArc<T: ?Sized>(Arc<T>);
41
42// Uniquene ownership means that we can support weaker bounds than `T: Send + Sync`.
43// Thus, these impls can follow the precedent of std's `Box`, not `Arc`.
44unsafe impl<T: ?Sized + Send> Send for UniqueArc<T> {}
45unsafe impl<T: ?Sized + Sync> Sync for UniqueArc<T> {}
46
47impl<T> UniqueArc<T> {
48    #[inline]
49    /// Construct a new UniqueArc
50    pub fn new(data: T) -> Self {
51        UniqueArc(Arc::new(data))
52    }
53
54    /// Construct an uninitialized arc
55    #[inline]
56    pub fn new_uninit() -> UniqueArc<MaybeUninit<T>> {
57        unsafe {
58            let layout = Layout::new::<ArcInner<MaybeUninit<T>>>();
59            let ptr = alloc::alloc::alloc(layout);
60            let mut p = NonNull::new(ptr)
61                .unwrap_or_else(|| alloc::alloc::handle_alloc_error(layout))
62                .cast::<ArcInner<MaybeUninit<T>>>();
63            ptr::write(&mut p.as_mut().count, AtomicUsize::new(1));
64
65            UniqueArc(Arc {
66                p,
67                phantom: PhantomData,
68            })
69        }
70    }
71
72    /// Gets the inner value of the unique arc
73    pub fn into_inner(this: Self) -> T {
74        // Wrap the Arc in a `ManuallyDrop` so that its drop routine never runs
75        let this = ManuallyDrop::new(this.0);
76        debug_assert!(
77            this.is_unique(),
78            "attempted to call `.into_inner()` on a `UniqueArc` with a non-zero ref count",
79        );
80
81        // Safety: We have exclusive access to the inner data and the
82        //         arc will not perform its drop routine since we've
83        //         wrapped it in a `ManuallyDrop`
84        unsafe { Box::from_raw(this.ptr()).data }
85    }
86}
87
88impl<T: ?Sized> UniqueArc<T> {
89    /// Convert to a shareable `Arc<T>` once we're done mutating it
90    #[inline]
91    pub fn shareable(self) -> Arc<T> {
92        self.0
93    }
94
95    /// Creates a new [`UniqueArc`] from the given [`Arc`].
96    ///
97    /// An unchecked alternative to `Arc::try_unique()`
98    ///
99    /// # Safety
100    ///
101    /// The given `Arc` must have a reference count of exactly one
102    ///
103    pub(crate) unsafe fn from_arc(arc: Arc<T>) -> Self {
104        debug_assert_eq!(Arc::count(&arc), 1);
105        Self(arc)
106    }
107
108    /// Creates a new `&mut `[`UniqueArc`] from the given `&mut `[`Arc`].
109    ///
110    /// An unchecked alternative to `Arc::try_as_unique()`
111    ///
112    /// # Safety
113    ///
114    /// The given `Arc` must have a reference count of exactly one
115    pub(crate) unsafe fn from_arc_ref(arc: &mut Arc<T>) -> &mut Self {
116        debug_assert_eq!(Arc::count(arc), 1);
117
118        // Safety: caller guarantees that `arc` is unique,
119        //         `UniqueArc` is `repr(transparent)`
120        &mut *(arc as *mut Arc<T> as *mut UniqueArc<T>)
121    }
122}
123
124impl<T> UniqueArc<MaybeUninit<T>> {
125    /// Calls `MaybeUninit::write` on the contained value.
126    pub fn write(&mut self, val: T) -> &mut T {
127        unsafe {
128            // Casting *mut MaybeUninit<T> -> *mut T is always fine
129            let ptr = self.as_mut_ptr() as *mut T;
130
131            // Safety: We have exclusive access to the inner data
132            ptr.write(val);
133
134            // Safety: the pointer was just written to
135            &mut *ptr
136        }
137    }
138
139    /// Obtain a mutable pointer to the stored `MaybeUninit<T>`.
140    pub fn as_mut_ptr(&mut self) -> *mut MaybeUninit<T> {
141        unsafe { &mut (*self.0.ptr()).data }
142    }
143
144    /// Convert to an initialized Arc.
145    ///
146    /// # Safety
147    ///
148    /// This function is equivalent to `MaybeUninit::assume_init` and has the
149    /// same safety requirements. You are responsible for ensuring that the `T`
150    /// has actually been initialized before calling this method.
151    #[inline]
152    pub unsafe fn assume_init(this: Self) -> UniqueArc<T> {
153        UniqueArc(Arc {
154            p: ManuallyDrop::new(this).0.p.cast(),
155            phantom: PhantomData,
156        })
157    }
158}
159
160impl<T> UniqueArc<[MaybeUninit<T>]> {
161    /// Create an Arc contains an array `[MaybeUninit<T>]` of `len`.
162    pub fn new_uninit_slice(len: usize) -> Self {
163        // Safety (although no unsafe is required):
164        // - `ArcInner` is properly allocated and initialized.
165        //   - `()` and `[MaybeUninit<T>]` do not require special initialization
166        // - The `Arc` is just created and so -- unique.
167        let arc: Arc<HeaderSlice<(), [MaybeUninit<T>]>> =
168            UniqueArc::from_header_and_uninit_slice((), len).0;
169        let arc: Arc<[MaybeUninit<T>]> = arc.into();
170        UniqueArc(arc)
171    }
172
173    /// # Safety
174    ///
175    /// Must initialize all fields before calling this function.
176    #[inline]
177    pub unsafe fn assume_init_slice(Self(this): Self) -> UniqueArc<[T]> {
178        UniqueArc(this.assume_init())
179    }
180}
181
182impl<H, T> UniqueArc<HeaderSlice<H, [MaybeUninit<T>]>> {
183    /// Creates an Arc for a HeaderSlice using the given header struct and allocated space
184    /// for an unitialized slice of length `len`.
185    #[inline]
186    pub fn from_header_and_uninit_slice(header: H, len: usize) -> Self {
187        let inner = Arc::<HeaderSlice<H, [MaybeUninit<T>]>>::allocate_for_header_and_slice(len);
188
189        unsafe {
190            // Safety: inner is a valid pointer, so this can't go out of bounds
191            let dst = addr_of_mut!((*inner.as_ptr()).data.header);
192
193            // Safety: `dst` is valid for writes (just allocated)
194            ptr::write(dst, header);
195        }
196
197        // Safety: ptr is valid & the inner structure is initialized.
198        // We wrote the header above and the slice can stay unitialized as it's [MaybeUninit<T>]
199        Self(Arc {
200            p: inner,
201            phantom: PhantomData,
202        })
203    }
204
205    /// # Safety
206    ///
207    /// Must initialize all fields before calling this function.
208    #[inline]
209    pub unsafe fn assume_init_slice_with_header(self) -> UniqueArc<HeaderSlice<H, [T]>> {
210        unsafe { core::mem::transmute(self) }
211    }
212}
213
214impl<T: ?Sized> TryFrom<Arc<T>> for UniqueArc<T> {
215    type Error = Arc<T>;
216
217    fn try_from(arc: Arc<T>) -> Result<Self, Self::Error> {
218        Arc::try_unique(arc)
219    }
220}
221
222impl<T: ?Sized> Deref for UniqueArc<T> {
223    type Target = T;
224
225    #[inline]
226    fn deref(&self) -> &T {
227        &self.0
228    }
229}
230
231impl<T: ?Sized> DerefMut for UniqueArc<T> {
232    #[inline]
233    fn deref_mut(&mut self) -> &mut T {
234        // We know this to be uniquely owned
235        unsafe { &mut (*self.0.ptr()).data }
236    }
237}
238
239impl<A> FromIterator<A> for UniqueArc<[A]> {
240    fn from_iter<T: IntoIterator<Item = A>>(iter: T) -> Self {
241        let iter = iter.into_iter();
242        let (lower, upper) = iter.size_hint();
243        let arc: Arc<[A]> = if Some(lower) == upper {
244            let iter = IteratorAsExactSizeIterator::new(iter);
245            Arc::from_header_and_iter((), iter).into()
246        } else {
247            let vec = iter.collect::<Vec<_>>();
248            Arc::from(vec)
249        };
250        // Safety: We just created an `Arc`, so it's unique.
251        unsafe { UniqueArc::from_arc(arc) }
252    }
253}
254
255// Safety:
256// This leverages the correctness of Arc's CoerciblePtr impl. Additionally, we must ensure that
257// this can not be used to violate the safety invariants of UniqueArc, which require that we can not
258// duplicate the Arc, such that replace_ptr returns a valid instance. This holds since it consumes
259// a unique owner of the contained ArcInner.
260#[cfg(feature = "unsize")]
261unsafe impl<T, U: ?Sized> unsize::CoerciblePtr<U> for UniqueArc<T> {
262    type Pointee = T;
263    type Output = UniqueArc<U>;
264
265    fn as_sized_ptr(&mut self) -> *mut T {
266        // Dispatch to the contained field.
267        unsize::CoerciblePtr::<U>::as_sized_ptr(&mut self.0)
268    }
269
270    unsafe fn replace_ptr(self, new: *mut U) -> UniqueArc<U> {
271        // Dispatch to the contained field, work around conflict of destructuring and Drop.
272        let inner = ManuallyDrop::new(self);
273        UniqueArc(ptr::read(&inner.0).replace_ptr(new))
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use crate::{Arc, HeaderSliceWithLength, HeaderWithLength, UniqueArc};
280    use core::{convert::TryFrom, mem::MaybeUninit};
281
282    #[test]
283    fn unique_into_inner() {
284        let unique = UniqueArc::new(10u64);
285        assert_eq!(UniqueArc::into_inner(unique), 10);
286    }
287
288    #[test]
289    fn try_from_arc() {
290        let x = Arc::new(10_000);
291        let y = x.clone();
292
293        assert!(UniqueArc::try_from(x).is_err());
294        assert_eq!(
295            UniqueArc::into_inner(UniqueArc::try_from(y).unwrap()),
296            10_000,
297        );
298    }
299
300    #[test]
301    #[allow(deprecated)]
302    fn maybeuninit_smoke() {
303        let mut arc: UniqueArc<MaybeUninit<_>> = UniqueArc::new_uninit();
304        arc.write(999);
305
306        let arc = unsafe { UniqueArc::assume_init(arc) };
307        assert_eq!(*arc, 999);
308    }
309
310    #[test]
311    fn from_header_and_uninit_slice() {
312        let mut uarc: UniqueArc<HeaderSliceWithLength<u8, [MaybeUninit<u16>]>> =
313            UniqueArc::from_header_and_uninit_slice(HeaderWithLength::new(1, 3), 3);
314        uarc.slice.fill(MaybeUninit::new(2));
315        let arc = unsafe { uarc.assume_init_slice_with_header() }.shareable();
316        assert!(arc.is_unique());
317        // Using clone to that the layout generated in new_uninit_slice is compatible
318        // with ArcInner.
319        let arcs = [
320            arc.clone(),
321            arc.clone(),
322            arc.clone(),
323            arc.clone(),
324            arc.clone(),
325        ];
326        // Similar for ThinArc
327        let thin = Arc::into_thin(arc.clone());
328        assert_eq!(7, Arc::count(&arc));
329        // If the layout is not compatible, then the data might be corrupted.
330        assert_eq!(arc.header.header, 1);
331        assert_eq!(&arc.slice, [2, 2, 2]);
332        assert_eq!(thin.header.header, 1);
333        assert_eq!(&thin.slice, [2, 2, 2]);
334
335        // Drop the arcs and check the count and the content to
336        // make sure it isn't corrupted.
337        drop(arcs);
338        drop(thin);
339        assert!(arc.is_unique());
340        assert_eq!(arc.header.header, 1);
341        assert_eq!(&arc.slice, [2, 2, 2]);
342    }
343}
triomphe/unique_arc.rs

triomphe/
unique_arc.rs
