sync: ArcIn unsized data support

score/sync/arc_in.rs

@@ -15,8 +15,9 @@ use core::alloc::Layout;
 use core::cmp::Ordering;
 use core::fmt as core_fmt;
 use core::hash::{Hash, Hasher};
+use core::mem::ManuallyDrop;
 use core::ops::Deref;
-use core::ptr::{drop_in_place, NonNull};
+use core::ptr::{drop_in_place, read, NonNull};
 use core::sync::atomic;
 use core::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
 use score_log::fmt as score_fmt;
@@ -32,12 +33,15 @@ use allocator::BasicAllocator;
 ///  - This is a simplified version and does not include weak references.
 ///  - This provides limited functionality compared to `std::sync::Arc` and shall be used only when custom allocator support is required.
 ///
-pub struct ArcIn<T, A: BasicAllocator> {
+pub struct ArcIn<T: ?Sized, A: BasicAllocator> {
     ptr: NonNull<ArcInner<T>>,
     alloc: A,
 }
 
-struct ArcInner<T> {
+/// Heap layout backing an [`ArcIn`].
+/// Not a part of the stable public API.
+#[doc(hidden)]
+pub struct ArcInner<T: ?Sized> {
     strong: AtomicUsize,
     data: T,
 }
@@ -62,15 +66,45 @@ impl<T, A: BasicAllocator + Clone> ArcIn<T, A> {
 
         ArcIn { ptr, alloc }
     }
+}
 
+impl<T: ?Sized, A: BasicAllocator> ArcIn<T, A> {
     /// Get strong reference count
     pub fn strong_count(this: &Self) -> usize {
         // SAFETY: `this.ptr` is guaranteed to be valid because we keep at least one strong reference by `this`
         unsafe { this.ptr.as_ref().strong.load(AtomicOrdering::SeqCst) }
     }
+
+    /// Decompose an `ArcIn` into its raw parts without touching the reference count.
+    ///
+    /// The returned pointer keeps ownership of the one strong reference held by `this`.
+    /// Caller must eventually pass the parts back to [`ArcIn::from_raw_parts`] to avoid leaking.
+    ///
+    /// The pointer may be unsize-coerced before being handed back.
+    /// E.g., `NonNull<ArcInner<T>>` to `NonNull<ArcInner<dyn Trait>>`.
+    pub fn into_raw_parts(this: Self) -> (NonNull<ArcInner<T>>, A) {
+        // Suppress `Drop` so the strong count is preserved and `alloc` can be moved out.
+        let this = ManuallyDrop::new(this);
+        let ptr = this.ptr;
+        // SAFETY:
+        // `this` is wrapped in `ManuallyDrop` and never dropped.
+        // Reading `alloc` out of it does not risk a double drop.
+        let alloc = unsafe { read(&this.alloc) };
+        (ptr, alloc)
+    }
+
+    /// Reassemble an `ArcIn` from parts previously obtained via [`ArcIn::into_raw_parts`].
+    ///
+    /// # Safety
+    ///
+    /// Parameters must originate from a single [`ArcIn::into_raw_parts`] call.
+    /// They must not be consumed by another call to [`ArcIn::from_raw_parts`].
+    pub unsafe fn from_raw_parts(ptr: NonNull<ArcInner<T>>, alloc: A) -> Self {
+        ArcIn { ptr, alloc }
+    }
 }
 
-impl<T, A: BasicAllocator + Clone> Clone for ArcIn<T, A> {
+impl<T: ?Sized, A: BasicAllocator + Clone> Clone for ArcIn<T, A> {
     fn clone(&self) -> Self {
         // SAFETY: `self.ptr` is guaranteed to be valid because we keep at least one strong reference by `self`
         unsafe {
@@ -84,7 +118,7 @@ impl<T, A: BasicAllocator + Clone> Clone for ArcIn<T, A> {
     }
 }
 
-impl<T, A: BasicAllocator> Deref for ArcIn<T, A> {
+impl<T: ?Sized, A: BasicAllocator> Deref for ArcIn<T, A> {
     type Target = T;
     fn deref(&self) -> &T {
         unsafe { &self.ptr.as_ref().data }
@@ -97,62 +131,62 @@ impl<T: Default, A: BasicAllocator + Clone + Default> Default for ArcIn<T, A> {
     }
 }
 
-impl<T: core_fmt::Debug, A: BasicAllocator> core_fmt::Debug for ArcIn<T, A> {
+impl<T: ?Sized + core_fmt::Debug, A: BasicAllocator> core_fmt::Debug for ArcIn<T, A> {
     fn fmt(&self, f: &mut core_fmt::Formatter<'_>) -> core_fmt::Result {
         self.deref().fmt(f)
     }
 }
 
-impl<T: score_fmt::ScoreDebug, A: BasicAllocator> score_fmt::ScoreDebug for ArcIn<T, A> {
+impl<T: ?Sized + score_fmt::ScoreDebug, A: BasicAllocator> score_fmt::ScoreDebug for ArcIn<T, A> {
     fn fmt(&self, f: score_fmt::Writer, spec: &score_fmt::FormatSpec) -> score_fmt::Result {
         self.deref().fmt(f, spec)
     }
 }
 
-impl<T, A: BasicAllocator> AsRef<T> for ArcIn<T, A> {
+impl<T: ?Sized, A: BasicAllocator> AsRef<T> for ArcIn<T, A> {
     fn as_ref(&self) -> &T {
         self.deref()
     }
 }
 
-impl<T: PartialEq, A: BasicAllocator> PartialEq for ArcIn<T, A> {
+impl<T: ?Sized + PartialEq, A: BasicAllocator> PartialEq for ArcIn<T, A> {
     fn eq(&self, other: &Self) -> bool {
         **self == **other
     }
 }
 
-impl<T: Eq, A: BasicAllocator> Eq for ArcIn<T, A> {}
+impl<T: ?Sized + Eq, A: BasicAllocator> Eq for ArcIn<T, A> {}
 
-impl<T: PartialOrd, A: BasicAllocator> PartialOrd for ArcIn<T, A> {
+impl<T: ?Sized + PartialOrd, A: BasicAllocator> PartialOrd for ArcIn<T, A> {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
         (**self).partial_cmp(&**other)
     }
 }
 
-impl<T: Ord, A: BasicAllocator> Ord for ArcIn<T, A> {
+impl<T: ?Sized + Ord, A: BasicAllocator> Ord for ArcIn<T, A> {
     fn cmp(&self, other: &Self) -> Ordering {
         (**self).cmp(&**other)
     }
 }
 
-impl<T: Hash, A: BasicAllocator> Hash for ArcIn<T, A> {
+impl<T: ?Sized + Hash, A: BasicAllocator> Hash for ArcIn<T, A> {
     fn hash<H: Hasher>(&self, state: &mut H) {
         (**self).hash(state);
     }
 }
 
-unsafe impl<T: Send + Sync, A: BasicAllocator + Send> Send for ArcIn<T, A> {}
-unsafe impl<T: Send + Sync, A: BasicAllocator + Sync> Sync for ArcIn<T, A> {}
+unsafe impl<T: ?Sized + Send + Sync, A: BasicAllocator + Send> Send for ArcIn<T, A> {}
+unsafe impl<T: ?Sized + Send + Sync, A: BasicAllocator + Sync> Sync for ArcIn<T, A> {}
 
-impl<T, A: BasicAllocator> Drop for ArcIn<T, A> {
+impl<T: ?Sized, A: BasicAllocator> Drop for ArcIn<T, A> {
     fn drop(&mut self) {
         if unsafe { self.ptr.as_ref().strong.fetch_sub(1, AtomicOrdering::Release) } == 1 {
             // SYNC: Ensure all previous writes are visible before we drop the data. This is enough because
             // we are the last strong reference.
             atomic::fence(AtomicOrdering::Acquire);
             unsafe {
+                let layout = Layout::for_value(self.ptr.as_ref());
                 drop_in_place(&mut self.ptr.as_mut().data);
-                let layout = Layout::new::<ArcInner<T>>();
                 self.alloc.deallocate(self.ptr.cast(), layout);
             }
         }
@@ -166,6 +200,12 @@ mod tests {
     use core::hash::Hash;
     use std::collections::hash_map::DefaultHasher;
 
+    #[test]
+    fn arc_in_is_send_sync_with_thread_safe_allocators() {
+        fn assert_send_sync<T: Send + Sync>() {}
+        assert_send_sync::<ArcIn<u32, HeapAllocator>>();
+    }
+
     #[test]
     fn new_and_deref() {
         let alloc = HeapAllocator;
@@ -265,4 +305,77 @@ mod tests {
         }
         assert!(dropped);
     }
+
+    #[test]
+    fn unsized_trait_object() {
+        trait Speak {
+            fn value(&self) -> i32;
+        }
+        struct Concrete(i32);
+        impl Speak for Concrete {
+            fn value(&self) -> i32 {
+                self.0
+            }
+        }
+
+        let alloc = HeapAllocator;
+        let arc: ArcIn<Concrete, HeapAllocator> = ArcIn::new_in(Concrete(7), alloc);
+
+        // Coerce `ArcIn<Concrete>` -> `ArcIn<dyn Speak>` via the raw-parts primitive.
+        // The unsizing happens on `NonNull`, which is a stable coercion.
+        let (ptr, alloc) = ArcIn::into_raw_parts(arc);
+        let ptr: NonNull<ArcInner<dyn Speak>> = ptr;
+        let arc: ArcIn<dyn Speak, HeapAllocator> = unsafe { ArcIn::from_raw_parts(ptr, alloc) };
+
+        assert_eq!(arc.value(), 7);
+        assert_eq!(ArcIn::strong_count(&arc), 1);
+
+        let arc2 = arc.clone();
+        assert_eq!(ArcIn::strong_count(&arc), 2);
+        assert_eq!(arc2.value(), 7);
+        drop(arc2);
+        assert_eq!(ArcIn::strong_count(&arc), 1);
+    }
+
+    #[test]
+    fn unsized_trait_object_drops_once() {
+        // Check deallocation happens with the correct layout when accessed as `dyn`.
+        static DROPS: AtomicUsize = AtomicUsize::new(0);
+
+        trait Marker {}
+        struct Tracked;
+        impl Marker for Tracked {}
+        impl Drop for Tracked {
+            fn drop(&mut self) {
+                DROPS.fetch_add(1, AtomicOrdering::SeqCst);
+            }
+        }
+
+        let alloc = HeapAllocator;
+        let arc = ArcIn::new_in(Tracked, alloc);
+        let (ptr, alloc) = ArcIn::into_raw_parts(arc);
+        let ptr: NonNull<ArcInner<dyn Marker>> = ptr;
+        let arc: ArcIn<dyn Marker, HeapAllocator> = unsafe { ArcIn::from_raw_parts(ptr, alloc) };
+
+        let arc2 = arc.clone();
+        assert_eq!(DROPS.load(AtomicOrdering::SeqCst), 0);
+        drop(arc);
+        assert_eq!(DROPS.load(AtomicOrdering::SeqCst), 0);
+        drop(arc2);
+        assert_eq!(DROPS.load(AtomicOrdering::SeqCst), 1);
+    }
+
+    #[test]
+    fn unsized_slice() {
+        let alloc = HeapAllocator;
+        let arc: ArcIn<[i32; 3], HeapAllocator> = ArcIn::new_in([10, 20, 30], alloc);
+
+        let (ptr, alloc) = ArcIn::into_raw_parts(arc);
+        let ptr: NonNull<ArcInner<[i32]>> = ptr;
+        let arc: ArcIn<[i32], HeapAllocator> = unsafe { ArcIn::from_raw_parts(ptr, alloc) };
+
+        assert_eq!(arc.len(), 3);
+        assert_eq!(arc[1], 20);
+        assert_eq!(&*arc, &[10, 20, 30]);
+    }
 }

score/sync/lib.rs

@@ -16,3 +16,5 @@
 mod arc_in;
 
 pub use arc_in::ArcIn;
+#[doc(hidden)]
+pub use arc_in::ArcInner;