Change interned strings to be thin pointers

This moves the length of the associated string into the arena allocation.

Author: Mark-Simulacrum

collector/src/intern.rs

@@ -1,28 +1,43 @@
 use bumpalo::Bump;
 use hashbrown::HashSet;
+use std::alloc::Layout;
+use std::fmt;
+use std::ptr;
 use std::sync::Mutex;
 
 pub trait InternString {
-    fn to_interned(s: &'static str) -> Self;
+    unsafe fn to_interned(s: ArenaStr) -> Self;
 }
 
 #[macro_export]
 macro_rules! intern {
     (pub struct $for_ty:ident) => {
-        #[derive(Serialize, Debug, PartialOrd, Ord, Copy, Clone)]
-        pub struct $for_ty(&'static str);
+        #[derive(Serialize, Debug, Copy, Clone)]
+        pub struct $for_ty(crate::intern::ArenaStr);
 
         impl std::cmp::PartialEq for $for_ty {
             fn eq(&self, other: &Self) -> bool {
-                std::ptr::eq(self.0.as_ptr(), other.0.as_ptr())
+                self.0.hash_ptr() == other.0.hash_ptr()
             }
         }
 
         impl std::cmp::Eq for $for_ty {}
 
+        impl std::cmp::PartialOrd for $for_ty {
+            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+                Some(self.cmp(other))
+            }
+        }
+
+        impl std::cmp::Ord for $for_ty {
+            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+                self.0.as_str().cmp(other.0.as_str())
+            }
+        }
+
         impl std::hash::Hash for $for_ty {
             fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-                state.write_usize(self.0.as_ptr() as usize);
+                state.write_usize(self.0.hash_ptr());
             }
         }
 
@@ -55,13 +70,13 @@ macro_rules! intern {
 
         impl std::cmp::PartialEq<str> for $for_ty {
             fn eq(&self, other: &str) -> bool {
-                self.0 == other
+                self.0.as_str() == other
             }
         }
 
         impl std::fmt::Display for $for_ty {
             fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-                write!(f, "{}", self.0)
+                write!(f, "{}", self.0.as_str())
             }
         }
 
@@ -74,28 +89,84 @@ macro_rules! intern {
         impl std::ops::Deref for $for_ty {
             type Target = str;
             fn deref(&self) -> &str {
-                self.0
+                self.0.as_str()
             }
         }
 
         impl crate::intern::InternString for $for_ty {
-            fn to_interned(v: &'static str) -> $for_ty {
+            unsafe fn to_interned(v: crate::intern::ArenaStr) -> $for_ty {
                 $for_ty(v)
             }
         }
     };
 }
 
 lazy_static::lazy_static! {
-    static ref INTERNED: Mutex<(HashSet<&'static str>, Bump)>
+    static ref INTERNED: Mutex<(HashSet<ArenaStr>, Bump)>
         = Mutex::new((HashSet::new(), Bump::new()));
 }
 
 pub fn intern<T: InternString>(value: &str) -> T {
     let mut guard = INTERNED.lock().unwrap();
 
     let (ref mut set, ref arena) = &mut *guard;
-    T::to_interned(set.get_or_insert_with(value, |_| -> &'static str {
-        unsafe { std::mem::transmute::<&str, &'static str>(arena.alloc_str(value)) }
-    }))
+    unsafe {
+        T::to_interned(*set.get_or_insert_with(value, |_| -> ArenaStr {
+            let ptr = arena.alloc_layout(
+                Layout::from_size_align(std::mem::size_of::<usize>() + value.len(), 1).unwrap(),
+            );
+            let start_at = ptr.as_ptr();
+            ptr::write(start_at as *mut _, value.len().to_ne_bytes());
+            let bytes = start_at.add(std::mem::size_of::<usize>());
+            ptr::copy_nonoverlapping(value.as_ptr(), bytes, value.len());
+
+            ArenaStr(start_at as *const u8)
+        }))
+    }
+}
+
+#[derive(serde::Serialize, Copy, Clone, PartialEq, Eq)]
+#[serde(into = "&'static str")]
+pub struct ArenaStr(*const u8);
+
+impl Into<&'static str> for ArenaStr {
+    fn into(self) -> &'static str {
+        self.as_str()
+    }
+}
+
+unsafe impl Send for ArenaStr {}
+unsafe impl Sync for ArenaStr {}
+
+impl ArenaStr {
+    pub fn as_str(self) -> &'static str {
+        unsafe {
+            let mut ptr = self.0;
+            let length = usize::from_ne_bytes(ptr::read(ptr as *const _));
+            ptr = ptr.add(std::mem::size_of::<usize>());
+            std::str::from_utf8_unchecked(std::slice::from_raw_parts(ptr, length))
+        }
+    }
+
+    pub fn hash_ptr(self) -> usize {
+        self.0 as usize
+    }
+}
+
+impl fmt::Debug for ArenaStr {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt::Debug::fmt(self.as_str(), f)
+    }
+}
+
+impl std::hash::Hash for ArenaStr {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.as_str().hash(state);
+    }
+}
+
+impl std::borrow::Borrow<str> for ArenaStr {
+    fn borrow(&self) -> &str {
+        self.as_str()
+    }
 }