From 8bafea57d8eff2f68c2bed6e74cfb5d9d5d82fcb Mon Sep 17 00:00:00 2001
From: Levi Morrison <levi.morrison@datadoghq.com>
Date: Mon, 23 Dec 2024 17:11:16 -0700
Subject: [PATCH] perf: inline and tweak ThinStr* internals

---
 thin-str/src/lib.rs         |  4 +---
 thin-str/src/thin_str.rs    |  8 --------
 thin-str/src/thin_string.rs | 32 +++++++++++++++++++++++---------
 3 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/thin-str/src/lib.rs b/thin-str/src/lib.rs
index 34bcfd0c66..8185f2eac2 100644
--- a/thin-str/src/lib.rs
+++ b/thin-str/src/lib.rs
@@ -122,10 +122,8 @@ impl<const N: usize> ConstStorage<N> {
         // ConstStorage only holds valid strs, and the lifetime is valid.
         unsafe { Storage::from_header(header_ptr) }
     }
-}
 
-impl<const N: usize> AsRef<str> for ConstStorage<N> {
-    fn as_ref(&self) -> &str {
+    pub const fn as_ref(&self) -> &str {
         // SAFETY: ConstStorage always hold valid strings.
         unsafe { core::str::from_utf8_unchecked(&self.data) }
     }
diff --git a/thin-str/src/thin_str.rs b/thin-str/src/thin_str.rs
index 2cd62f4d99..a288f97e2b 100644
--- a/thin-str/src/thin_str.rs
+++ b/thin-str/src/thin_str.rs
@@ -85,14 +85,6 @@ impl<'a> From<&ThinStr<'a>> for &'a Storage {
     }
 }
 
-impl<'a> From<&'a Storage> for ThinStr<'a> {
-    fn from(storage: &'a Storage) -> ThinStr<'a> {
-        let obj = ptr::NonNull::from(storage).cast::<ThinHeader>();
-        // SAFETY: pointer to Storage conforms to layout required for ThinStr.
-        unsafe { ThinStr::from_raw(obj) }
-    }
-}
-
 impl Deref for ThinStr<'_> {
     type Target = str;
 
diff --git a/thin-str/src/thin_string.rs b/thin-str/src/thin_string.rs
index adfd41ab84..d733ac62cb 100644
--- a/thin-str/src/thin_string.rs
+++ b/thin-str/src/thin_string.rs
@@ -42,6 +42,7 @@ unsafe impl<A: Allocator + Sync> Sync for ThinString<A> {}
 impl<A: Allocator> ThinString<A> {
     /// Creates the empty ThinString.
     /// This will not allocate.
+    #[inline]
     pub fn new_in(allocator: A) -> Self {
         Self::from_const_storage_in(&EMPTY, allocator)
     }
@@ -49,6 +50,7 @@ impl<A: Allocator> ThinString<A> {
     /// Create a ThinString from &static [ConstStorage]. Note that there's not
     /// _quite_ enough support to make this a const fn.
     /// This will not allocate.
+    #[inline]
     pub fn from_const_storage_in<const N: usize>(
         const_storage: &'static ConstStorage<N>,
         allocator: A,
@@ -64,6 +66,7 @@ impl<A: Allocator> ThinString<A> {
     /// This will allocate in the given allocator.
     /// # Panics
     /// This panics if allocation fails.
+    #[inline]
     pub fn from_str_in(string: &str, allocator: A) -> Self {
         Self::try_from_str_in(string, allocator).unwrap()
     }
@@ -73,7 +76,7 @@ impl<A: Allocator> ThinString<A> {
     pub fn try_from_str_in(string: &str, allocator: A) -> Result<Self, AllocError> {
         let header = Layout::new::<ThinHeader>();
         let data = Layout::for_value(string);
-        let Ok((layout, offset)) = header.extend(data) else {
+        let Ok((layout, _offset)) = header.extend(data) else {
             return Err(AllocError);
         };
         // Padding is important here, since the str could be an odd number of
@@ -81,8 +84,8 @@ impl<A: Allocator> ThinString<A> {
         let layout = layout.pad_to_align();
 
         // Sanity check some things to defend against refactoring.
-        debug_assert_eq!(offset, mem::size_of::<ThinHeader>());
-        debug_assert_eq!(offset, mem::size_of::<usize>());
+        debug_assert_eq!(_offset, mem::size_of::<ThinHeader>());
+        debug_assert_eq!(_offset, mem::size_of::<usize>());
 
         let obj = allocator.allocate(layout)?;
         let header = ThinHeader::from(string.len());
@@ -95,7 +98,7 @@ impl<A: Allocator> ThinString<A> {
 
         // SAFETY: the offset is at the correct place from the base pointer,
         // and again, the memory is valid for writes and Drop is irrelevant.
-        let data = unsafe { obj.cast::<u8>().as_ptr().add(offset) };
+        let data = unsafe { obj.cast::<u8>().as_ptr().add(mem::size_of::<usize>()) };
 
         // SAFETY: a new, non-zero sized allocated object must not overlap
         // with existing memory by definition. Both regions are at valid for
@@ -115,9 +118,22 @@ impl<A: Allocator> ThinString<A> {
     }
 
     /// Creates a ThinStr that borrows from the ThinString.
+    #[inline]
     pub fn as_thin_str(&self) -> ThinStr {
-        let storage: &Storage = self.into();
-        ThinStr::from(storage)
+        let obj = self.tagged_ptr.ptr();
+        // SAFETY: ThinString points to a valid header, no mutable references
+        // ever exist for ThinString, and the lifetime ensures the storage
+        // lives long enough.
+        unsafe { ThinStr::from_raw(obj) }
+    }
+
+    #[inline]
+    pub fn as_storage(&self) -> &Storage {
+        let obj = self.tagged_ptr.ptr();
+        // SAFETY: ThinString points to a valid header, no mutable references
+        // ever exist for ThinString, and the lifetime ensures the storage
+        // lives long enough.
+        unsafe { Storage::from_header(obj) }
     }
 }
 
@@ -159,9 +175,7 @@ impl<A: Allocator> Drop for ThinString<A> {
 
 impl<'a, A: Allocator> From<&'a ThinString<A>> for &'a Storage {
     fn from(thin_string: &'a ThinString<A>) -> Self {
-        let obj = thin_string.tagged_ptr.ptr();
-        // SAFETY: todo
-        unsafe { Storage::from_header(obj) }
+        thin_string.as_storage()
     }
 }
 impl<const N: usize> From<&'static ConstStorage<N>> for ThinString<Global> {