Be smarter and less UB about pointer alignment.

Author: bodil

src/boxed.rs

@@ -4,7 +4,7 @@
 
 use alloc::{alloc::Layout, string::String};
 use core::{
-    mem::forget,
+    mem::{align_of, forget},
     ops::{Deref, DerefMut},
     ptr::NonNull,
 };
@@ -22,7 +22,7 @@ pub(crate) struct BoxedString {
 #[cfg(endian = "big")]
 #[repr(C)]
 pub(crate) struct BoxedString {
-    length: usize,
+    len: usize,
     cap: usize,
     ptr: NunNull<u8>,
 }
@@ -44,8 +44,17 @@ impl GenericString for BoxedString {
 impl BoxedString {
     const MINIMAL_CAPACITY: usize = MAX_INLINE * 2;
 
+    pub(crate) fn check_alignment(this: &Self) -> usize {
+        let ptr: *const u8 = this.ptr.as_ptr();
+        ptr.align_offset(2)
+    }
+
     fn layout_for(cap: usize) -> Layout {
-        let layout = Layout::array::<u8>(cap).unwrap();
+        // Always request memory that is specifically aligned to at least 2, so
+        // the least significant bit is guaranteed to be 0.
+        let layout = Layout::array::<u8>(cap)
+            .and_then(|layout| layout.align_to(align_of::<u16>()))
+            .unwrap();
         assert!(
             layout.size() <= isize::MAX as usize,
             "allocation too large!"
@@ -56,11 +65,12 @@ impl BoxedString {
     fn alloc(cap: usize) -> NonNull<u8> {
         let layout = Self::layout_for(cap);
         #[allow(unsafe_code)]
-        let ptr = unsafe { alloc::alloc::alloc(layout) };
-        match NonNull::new(ptr) {
+        let ptr = match NonNull::new(unsafe { alloc::alloc::alloc(layout) }) {
             Some(ptr) => ptr,
             None => alloc::alloc::handle_alloc_error(layout),
-        }
+        };
+        debug_assert!(ptr.as_ptr().align_offset(2) == 0);
+        ptr
     }
 
     fn realloc(&mut self, cap: usize) {
@@ -74,6 +84,7 @@ impl BoxedString {
             None => alloc::alloc::handle_alloc_error(layout),
         };
         self.cap = cap;
+        debug_assert!(self.ptr.as_ptr().align_offset(2) == 0);
     }
 
     pub(crate) fn ensure_capacity(&mut self, target_cap: usize) {

src/config.rs

@@ -5,7 +5,7 @@
 use crate::{boxed::BoxedString, inline::InlineString, SmartString};
 use alloc::string::String;
 use core::mem::{align_of, size_of};
-use static_assertions::{assert_eq_size, const_assert, const_assert_eq};
+use static_assertions::{assert_eq_align, assert_eq_size, const_assert, const_assert_eq};
 
 /// A compact string representation equal to [`String`] in size with guaranteed inlining.
 ///
@@ -71,6 +71,11 @@ assert_eq_size!(BoxedString, SmartString<LazyCompact>);
 assert_eq_size!(InlineString, SmartString<Compact>);
 assert_eq_size!(InlineString, SmartString<LazyCompact>);
 
+assert_eq_align!(BoxedString, String);
+assert_eq_align!(InlineString, String);
+assert_eq_align!(SmartString<Compact>, String);
+assert_eq_align!(SmartString<LazyCompact>, String);
+
 assert_eq_size!(String, SmartString<Compact>);
 assert_eq_size!(String, SmartString<LazyCompact>);
 

src/inline.rs

@@ -10,13 +10,17 @@ use core::{
 
 #[cfg(not(endian = "big"))]
 #[repr(C)]
+#[cfg_attr(target_pointer_width = "64", repr(align(8)))]
+#[cfg_attr(target_pointer_width = "32", repr(align(4)))]
 pub(crate) struct InlineString {
     pub(crate) marker: Marker,
     pub(crate) data: [u8; MAX_INLINE],
 }
 
 #[cfg(endian = "big")]
 #[repr(C)]
+#[cfg_attr(target_pointer_width = "64", repr(align(8)))]
+#[cfg_attr(target_pointer_width = "32", repr(align(4)))]
 pub(crate) struct InlineString {
     pub(crate) data: [u8; MAX_INLINE],
     pub(crate) marker: Marker,

src/lib.rs

@@ -188,9 +188,6 @@ pub mod alias {
 /// state changes wouldn't carry over if the inline string is promoted to a boxed
 /// one - not without also storing that state in the inline representation, which
 /// would waste precious bytes for inline string data.
-#[repr(C)]
-#[cfg_attr(target_pointer_width = "64", repr(align(8)))]
-#[cfg_attr(target_pointer_width = "32", repr(align(4)))]
 pub struct SmartString<Mode: SmartStringMode> {
     data: MaybeUninit<InlineString>,
     mode: PhantomData<Mode>,
@@ -300,8 +297,11 @@ impl<Mode: SmartStringMode> SmartString<Mode> {
     }
 
     fn discriminant(&self) -> Discriminant {
+        // unsafe { self.data.assume_init() }.marker.discriminant()
+        let str_ptr: *const BoxedString =
+            self.data.as_ptr().cast() as *const _ as *const BoxedString;
         #[allow(unsafe_code)]
-        unsafe { self.data.assume_init() }.marker.discriminant()
+        Discriminant::from_bit(BoxedString::check_alignment(unsafe { &*str_ptr }) == 1)
     }
 
     fn cast(&self) -> StringCast<'_> {

src/marker_byte.rs

@@ -10,7 +10,7 @@ pub(crate) enum Discriminant {
 
 impl Discriminant {
     #[inline(always)]
-    const fn from_bit(bit: bool) -> Self {
+    pub(crate) const fn from_bit(bit: bool) -> Self {
         if bit {
             Self::Inline
         } else {

src/test.rs

@@ -523,4 +523,32 @@ mod tests {
             SmartString::<Compact>::from("\u{323}\u{323}\u{323}ω\u{323}\u{323}\u{323}㌣\u{e323}㤘");
         s.remove(20);
     }
+
+    #[test]
+    fn check_alignment() {
+        use crate::boxed::BoxedString;
+        use crate::inline::InlineString;
+        use crate::marker_byte::Discriminant;
+
+        let inline = InlineString::new();
+        let inline_ptr: *const InlineString = &inline;
+        let boxed_ptr: *const BoxedString = inline_ptr.cast();
+        #[allow(unsafe_code)]
+        let discriminant =
+            Discriminant::from_bit(BoxedString::check_alignment(unsafe { &*boxed_ptr }) == 1);
+        assert_eq!(Discriminant::Inline, discriminant);
+
+        let boxed = BoxedString::from_str(32, "welp");
+        let discriminant = Discriminant::from_bit(BoxedString::check_alignment(&boxed) == 1);
+        assert_eq!(Discriminant::Boxed, discriminant);
+
+        let mut s = SmartString::<Compact>::new();
+        assert_eq!(Discriminant::Inline, s.discriminant());
+        let big_str = "1234567890123456789012345678901234567890";
+        assert!(big_str.len() > MAX_INLINE);
+        s.push_str(big_str);
+        assert_eq!(Discriminant::Boxed, s.discriminant());
+        s.clear();
+        assert_eq!(Discriminant::Inline, s.discriminant());
+    }
 }