Seprate HeaderSliceWithLength into checked and unchecked variants

src/header.rs

@@ -245,15 +245,56 @@ impl<T> From<Vec<T>> for Arc<[T]> {
     }
 }
 
-pub(crate) type HeaderSliceWithLength<H, T> = HeaderSlice<HeaderWithLength<H>, T>;
+/// A type wrapping `HeaderSlice<HeaderWithLength<H>, T>` that is used internally in `ThinArc`.
+///
+/// # Safety
+///
+/// Safety-usable invariants:
+///
+/// - This is guaranteed to have the same representation as `HeaderSlice<HeaderWithLength<H>, T>` (i.e. `HeaderSliceWithLengthUnchecked`)
+/// - For `T` that is  `[U]` or `str`, the header length (`.length()` is checked to be the slice length)
+#[derive(Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
+#[repr(transparent)]
+pub struct HeaderSliceWithLengthProtected<H, T: ?Sized> {
+    // Invariant: if T is [U] or str, then the header's length field must be the slice length
+    // Currently no other DSTs are used with this type amd it has no invariants, but these may be added in the future
+    inner: HeaderSliceWithLengthUnchecked<H, T>,
+}
+
+pub(crate) type HeaderSliceWithLengthUnchecked<H, T> = HeaderSlice<HeaderWithLength<H>, T>;
+
+impl<H, T: ?Sized> HeaderSliceWithLengthProtected<H, T> {
+    pub fn header(&self) -> &H {
+        &self.inner.header.header
+    }
+    pub fn header_mut(&mut self) -> &mut H {
+        // Safety: only the length is unsafe to mutate
+        &mut self.inner.header.header
+    }
+    pub fn length(&self) -> usize {
+        self.inner.header.length
+    }
+
+    pub fn slice(&self) -> &T {
+        &self.inner.slice
+    }
+    pub fn slice_mut(&mut self) -> &mut T {
+        // Safety: only the length is unsafe to mutate
+        &mut self.inner.slice
+    }
+    pub(crate) fn inner(&self) -> &HeaderSliceWithLengthUnchecked<H, T> {
+        // This is safe in an immutable context
+        &self.inner
+    }
+}
 
-impl<H: PartialOrd, T: ?Sized + PartialOrd> PartialOrd for HeaderSliceWithLength<H, T> {
+impl<H: PartialOrd, T: ?Sized + PartialOrd> PartialOrd for HeaderSliceWithLengthUnchecked<H, T> {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
         (&self.header.header, &self.slice).partial_cmp(&(&other.header.header, &other.slice))
     }
 }
 
-impl<H: Ord, T: ?Sized + Ord> Ord for HeaderSliceWithLength<H, T> {
+impl<H: Ord, T: ?Sized + Ord> Ord for HeaderSliceWithLengthUnchecked<H, T> {
     fn cmp(&self, other: &Self) -> Ordering {
         (&self.header.header, &self.slice).cmp(&(&other.header.header, &other.slice))
     }

src/thin_arc.rs

@@ -4,11 +4,13 @@ use core::fmt;
 use core::hash::{Hash, Hasher};
 use core::iter::{ExactSizeIterator, Iterator};
 use core::marker::PhantomData;
+use core::mem;
 use core::mem::ManuallyDrop;
 use core::ops::Deref;
 use core::ptr;
 
-use super::{Arc, ArcInner, HeaderSliceWithLength, HeaderWithLength};
+use super::{Arc, ArcInner, HeaderSliceWithLengthProtected, HeaderWithLength};
+use crate::header::HeaderSliceWithLengthUnchecked;
 
 /// A "thin" `Arc` containing dynamically sized data
 ///
@@ -24,10 +26,10 @@ use super::{Arc, ArcInner, HeaderSliceWithLength, HeaderWithLength};
 /// Note that we use `[T; 0]` in order to have the right alignment for `T`.
 ///
 /// `ThinArc` solves this by storing the length in the allocation itself,
-/// via `HeaderSliceWithLength`.
+/// via `HeaderSliceWithLengthProtected`.
 #[repr(transparent)]
 pub struct ThinArc<H, T> {
-    ptr: ptr::NonNull<ArcInner<HeaderSliceWithLength<H, [T; 0]>>>,
+    ptr: ptr::NonNull<ArcInner<HeaderSliceWithLengthProtected<H, [T; 0]>>>,
     phantom: PhantomData<(H, T)>,
 }
 
@@ -39,12 +41,12 @@ unsafe impl<H: Sync + Send, T: Sync + Send> Sync for ThinArc<H, T> {}
 // See the comment around the analogous operation in from_header_and_iter.
 #[inline]
 fn thin_to_thick<H, T>(
-    thin: *mut ArcInner<HeaderSliceWithLength<H, [T; 0]>>,
-) -> *mut ArcInner<HeaderSliceWithLength<H, [T]>> {
-    let len = unsafe { (*thin).data.header.length };
+    thin: *mut ArcInner<HeaderSliceWithLengthProtected<H, [T; 0]>>,
+) -> *mut ArcInner<HeaderSliceWithLengthProtected<H, [T]>> {
+    let len = unsafe { (*thin).data.length() };
     let fake_slice = ptr::slice_from_raw_parts_mut(thin as *mut T, len);
 
-    fake_slice as *mut ArcInner<HeaderSliceWithLength<H, [T]>>
+    fake_slice as *mut ArcInner<HeaderSliceWithLengthProtected<H, [T]>>
 }
 
 impl<H, T> ThinArc<H, T> {
@@ -53,7 +55,27 @@ impl<H, T> ThinArc<H, T> {
     #[inline]
     pub fn with_arc<F, U>(&self, f: F) -> U
     where
-        F: FnOnce(&Arc<HeaderSliceWithLength<H, [T]>>) -> U,
+        F: FnOnce(&Arc<HeaderSliceWithLengthUnchecked<H, [T]>>) -> U,
+    {
+        // Synthesize transient Arc, which never touches the refcount of the ArcInner.
+        let transient = unsafe {
+            ManuallyDrop::new(Arc::from_protected(Arc {
+                p: ptr::NonNull::new_unchecked(thin_to_thick(self.ptr.as_ptr())),
+                phantom: PhantomData,
+            }))
+        };
+
+        // Expose the transient Arc to the callback, which may clone it if it wants
+        // and forward the result to the user
+        f(&transient)
+    }
+
+    /// Temporarily converts |self| into a bonafide Arc and exposes it to the
+    /// provided callback. The refcount is not modified.
+    #[inline]
+    fn with_protected_arc<F, U>(&self, f: F) -> U
+    where
+        F: FnOnce(&Arc<HeaderSliceWithLengthProtected<H, [T]>>) -> U,
     {
         // Synthesize transient Arc, which never touches the refcount of the ArcInner.
         let transient = unsafe {
@@ -73,19 +95,40 @@ impl<H, T> ThinArc<H, T> {
     #[inline]
     pub fn with_arc_mut<F, U>(&mut self, f: F) -> U
     where
-        F: FnOnce(&mut Arc<HeaderSliceWithLength<H, [T]>>) -> U,
+        F: FnOnce(&mut Arc<HeaderSliceWithLengthProtected<H, [T]>>) -> U,
     {
+        // It is possible for the user to replace the Arc entirely here. If so, we need to update the ThinArc as well
+        // whenever this method exits. We do this with a drop guard to handle the panicking case
+        struct DropGuard<'a, H, T> {
+            transient: ManuallyDrop<Arc<HeaderSliceWithLengthProtected<H, [T]>>>,
+            this: &'a mut ThinArc<H, T>,
+        }
+
+        impl<'a, H, T> Drop for DropGuard<'a, H, T> {
+            fn drop(&mut self) {
+                // Safety: We're still in the realm of Protected types so this cast is safe
+                self.this.ptr = self.transient.p.cast();
+            }
+        }
+
         // Synthesize transient Arc, which never touches the refcount of the ArcInner.
-        let mut transient = unsafe {
+        let transient = unsafe {
             ManuallyDrop::new(Arc {
                 p: ptr::NonNull::new_unchecked(thin_to_thick(self.ptr.as_ptr())),
                 phantom: PhantomData,
             })
         };
 
+        let mut guard = DropGuard {
+            transient,
+            this: self,
+        };
+
         // Expose the transient Arc to the callback, which may clone it if it wants
         // and forward the result to the user
-        f(&mut transient)
+        let ret = f(&mut guard.transient);
+
+        ret
     }
 
     /// Creates a `ThinArc` for a HeaderSlice using the given header struct and
@@ -171,21 +214,18 @@ impl<H, T> ThinArc<H, T> {
 }
 
 impl<H, T> Deref for ThinArc<H, T> {
-    type Target = HeaderSliceWithLength<H, [T]>;
+    type Target = HeaderSliceWithLengthUnchecked<H, [T]>;
 
     #[inline]
     fn deref(&self) -> &Self::Target {
-        unsafe { &(*thin_to_thick(self.ptr.as_ptr())).data }
+        unsafe { &(*thin_to_thick(self.ptr.as_ptr())).data.inner() }
     }
 }
 
 impl<H, T> Clone for ThinArc<H, T> {
     #[inline]
     fn clone(&self) -> Self {
-        ThinArc::with_arc(self, |a| {
-            // Safety: `a` isn't mutable thus the header length remains valid
-            unsafe { Arc::into_thin_unchecked(a.clone()) }
-        })
+        ThinArc::with_protected_arc(self, |a| Arc::protected_into_thin(a.clone()))
     }
 }
 
@@ -199,30 +239,19 @@ impl<H, T> Drop for ThinArc<H, T> {
     }
 }
 
-impl<H, T> Arc<HeaderSliceWithLength<H, [T]>> {
+impl<H, T> Arc<HeaderSliceWithLengthUnchecked<H, [T]>> {
     /// Converts an `Arc` into a `ThinArc`. This consumes the `Arc`, so the refcount
     /// is not modified.
     ///
     /// # Safety
     /// Assumes that the header length matches the slice length.
     #[inline]
     unsafe fn into_thin_unchecked(a: Self) -> ThinArc<H, T> {
-        let a = ManuallyDrop::new(a);
-        debug_assert_eq!(
-            a.header.length,
-            a.slice.len(),
-            "Length needs to be correct for ThinArc to work"
-        );
-        let fat_ptr: *mut ArcInner<HeaderSliceWithLength<H, [T]>> = a.ptr();
-        let thin_ptr = fat_ptr as *mut [usize] as *mut usize;
-        ThinArc {
-            ptr: unsafe {
-                ptr::NonNull::new_unchecked(
-                    thin_ptr as *mut ArcInner<HeaderSliceWithLength<H, [T; 0]>>,
-                )
-            },
-            phantom: PhantomData,
-        }
+        // Safety: invariant bubbled up
+        let this_protected: Arc<HeaderSliceWithLengthProtected<H, [T]>> =
+            unsafe { Arc::from_unprotected_unchecked(a) };
+
+        Arc::protected_into_thin(this_protected)
     }
 
     /// Converts an `Arc` into a `ThinArc`. This consumes the `Arc`, so the refcount
@@ -234,13 +263,55 @@ impl<H, T> Arc<HeaderSliceWithLength<H, [T]>> {
             a.slice.len(),
             "Length needs to be correct for ThinArc to work"
         );
+        // Safety: invariant checked in assertion above
         unsafe { Self::into_thin_unchecked(a) }
     }
 
     /// Converts a `ThinArc` into an `Arc`. This consumes the `ThinArc`, so the refcount
     /// is not modified.
     #[inline]
     pub fn from_thin(a: ThinArc<H, T>) -> Self {
+        Self::from_protected(Arc::<HeaderSliceWithLengthProtected<H, [T]>>::protected_from_thin(a))
+    }
+
+    /// Converts an `Arc` into a `ThinArc`. This consumes the `Arc`, so the refcount
+    /// is not modified.
+    #[inline]
+    fn from_protected(a: Arc<HeaderSliceWithLengthProtected<H, [T]>>) -> Self {
+        // Safety: HeaderSliceWithLengthProtected and HeaderSliceWithLengthUnchecked have the same layout
+        unsafe { mem::transmute(a) }
+    }
+}
+
+impl<H, T> Arc<HeaderSliceWithLengthProtected<H, [T]>> {
+    /// Converts an `Arc` into a `ThinArc`. This consumes the `Arc`, so the refcount
+    /// is not modified.
+    #[inline]
+    pub fn protected_into_thin(a: Self) -> ThinArc<H, T> {
+        debug_assert_eq!(
+            a.length(),
+            a.slice().len(),
+            "Length needs to be correct for ThinArc to work"
+        );
+
+        let a = ManuallyDrop::new(a);
+        let fat_ptr: *mut ArcInner<HeaderSliceWithLengthProtected<H, [T]>> = a.ptr();
+        let thin_ptr = fat_ptr as *mut [usize] as *mut usize;
+        ThinArc {
+            // Safety: The pointer comes from a valid Arc, and HeaderSliceWithLengthProtected has the correct length invariant
+            ptr: unsafe {
+                ptr::NonNull::new_unchecked(
+                    thin_ptr as *mut ArcInner<HeaderSliceWithLengthProtected<H, [T; 0]>>,
+                )
+            },
+            phantom: PhantomData,
+        }
+    }
+
+    /// Converts a `ThinArc` into an `Arc`. This consumes the `ThinArc`, so the refcount
+    /// is not modified.
+    #[inline]
+    pub fn protected_from_thin(a: ThinArc<H, T>) -> Self {
         let a = ManuallyDrop::new(a);
         let ptr = thin_to_thick(a.ptr.as_ptr());
         unsafe {
@@ -250,6 +321,17 @@ impl<H, T> Arc<HeaderSliceWithLength<H, [T]>> {
             }
         }
     }
+
+    /// Obtains a HeaderSliceWithLengthProtected from an unchecked HeaderSliceWithLengthUnchecked, wrapped in an Arc
+    ///
+    /// # Safety
+    /// Assumes that the header length matches the slice length.
+    #[inline]
+    unsafe fn from_unprotected_unchecked(a: Arc<HeaderSliceWithLengthUnchecked<H, [T]>>) -> Self {
+        // Safety: HeaderSliceWithLengthProtected and HeaderSliceWithLengthUnchecked have the same layout
+        // and the safety invariant on HeaderSliceWithLengthProtected.inner is bubbled up
+        unsafe { mem::transmute(a) }
+    }
 }
 
 impl<H: PartialEq, T: PartialEq> PartialEq for ThinArc<H, T> {
@@ -472,7 +554,7 @@ mod tests {
     #[test]
     fn with_arc_mut() {
         let mut arc: ThinArc<u8, u16> = ThinArc::from_header_and_slice(1u8, &[1, 2, 3]);
-        arc.with_arc_mut(|arc| Arc::get_mut(arc).unwrap().slice.fill(2));
+        arc.with_arc_mut(|arc| Arc::get_mut(arc).unwrap().slice_mut().fill(2));
         arc.with_arc_mut(|arc| assert!(Arc::get_unique(arc).is_some()));
         arc.with_arc(|arc| assert!(Arc::is_unique(arc)));
         // Using clone to that the layout generated in new_uninit_slice is compatible

src/unique_arc.rs

@@ -299,7 +299,7 @@ impl<T: Serialize> Serialize for UniqueArc<T> {
 
 #[cfg(test)]
 mod tests {
-    use crate::{Arc, HeaderSliceWithLength, HeaderWithLength, UniqueArc};
+    use crate::{Arc, HeaderSliceWithLengthUnchecked, HeaderWithLength, UniqueArc};
     use core::{convert::TryFrom, mem::MaybeUninit};
 
     #[test]
@@ -332,7 +332,7 @@ mod tests {
 
     #[test]
     fn from_header_and_uninit_slice() {
-        let mut uarc: UniqueArc<HeaderSliceWithLength<u8, [MaybeUninit<u16>]>> =
+        let mut uarc: UniqueArc<HeaderSliceWithLengthUnchecked<u8, [MaybeUninit<u16>]>> =
             UniqueArc::from_header_and_uninit_slice(HeaderWithLength::new(1, 3), 3);
         uarc.slice.fill(MaybeUninit::new(2));
         let arc = unsafe { uarc.assume_init_slice_with_header() }.shareable();