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

triomphe/
unique_arc.rs
