Auto merge of #136401 - Mark-Simulacrum:lockfree-as-str, r=<try>

Lazy-chunk Symbol interner

This fixes the unsoundness in Symbol::as_str, while also making it a bit faster (replacing locking with a single atomic load).

Author: bors

Cargo.lock

@@ -3492,6 +3492,7 @@ dependencies = [
  "either",
  "elsa",
  "ena",
+ "hashbrown 0.15.2",
  "indexmap",
  "jobserver",
  "libc",
@@ -4411,6 +4412,8 @@ version = "0.0.0"
 dependencies = [
  "blake3",
  "derive-where",
+ "elsa",
+ "hashbrown 0.15.2",
  "indexmap",
  "itoa",
  "md-5",

compiler/rustc_data_structures/Cargo.toml

@@ -10,6 +10,7 @@ bitflags = "2.4.1"
 either = "1.0"
 elsa = "1.11.0"
 ena = "0.14.3"
+hashbrown = { version = "0.15.2", default-features = false }
 indexmap = "2.4.0"
 jobserver_crate = { version = "0.1.28", package = "jobserver" }
 measureme = "11"

compiler/rustc_data_structures/src/marker.rs

@@ -70,6 +70,7 @@ impl_dyn_send!(
     [std::sync::LazyLock<T, F> where T: DynSend, F: DynSend]
     [std::collections::HashSet<K, S> where K: DynSend, S: DynSend]
     [std::collections::HashMap<K, V, S> where K: DynSend, V: DynSend, S: DynSend]
+    [hashbrown::HashTable<T> where T: DynSend]
     [std::collections::BTreeMap<K, V, A> where K: DynSend, V: DynSend, A: std::alloc::Allocator + Clone + DynSend]
     [Vec<T, A> where T: DynSend, A: std::alloc::Allocator + DynSend]
     [Box<T, A> where T: ?Sized + DynSend, A: std::alloc::Allocator + DynSend]
@@ -144,6 +145,7 @@ impl_dyn_sync!(
     [std::sync::LazyLock<T, F> where T: DynSend + DynSync, F: DynSend]
     [std::collections::HashSet<K, S> where K: DynSync, S: DynSync]
     [std::collections::HashMap<K, V, S> where K: DynSync, V: DynSync, S: DynSync]
+    [hashbrown::HashTable<T> where T: DynSync]
     [std::collections::BTreeMap<K, V, A> where K: DynSync, V: DynSync, A: std::alloc::Allocator + Clone + DynSync]
     [Vec<T, A> where T: DynSync, A: std::alloc::Allocator + DynSync]
     [Box<T, A> where T: ?Sized + DynSync, A: std::alloc::Allocator + DynSync]

compiler/rustc_span/Cargo.toml

@@ -7,6 +7,8 @@ edition = "2024"
 # tidy-alphabetical-start
 blake3 = "1.5.2"
 derive-where = "1.2.7"
+elsa = "1.11.0"
+hashbrown = { version = "0.15.2", default-features = false }
 indexmap = { version = "2.0.0" }
 itoa = "1.0"
 md5 = { package = "md-5", version = "0.10.0" }

compiler/rustc_span/src/lib.rs

@@ -27,11 +27,14 @@
 #![feature(let_chains)]
 #![feature(map_try_insert)]
 #![feature(negative_impls)]
+#![feature(new_zeroed_alloc)]
 #![feature(read_buf)]
 #![feature(round_char_boundary)]
 #![feature(rustc_attrs)]
 #![feature(rustdoc_internals)]
 #![feature(slice_as_chunks)]
+#![feature(str_from_raw_parts)]
+#![feature(sync_unsafe_cell)]
 #![warn(unreachable_pub)]
 // tidy-alphabetical-end
 

compiler/rustc_span/src/symbol.rs

@@ -2,11 +2,13 @@
 //! allows bidirectional lookup; i.e., given a value, one can easily find the
 //! type, and vice versa.
 
-use std::hash::{Hash, Hasher};
+use std::cell::SyncUnsafeCell;
+use std::hash::{BuildHasher, BuildHasherDefault, Hash, Hasher};
+use std::sync::LazyLock;
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::{fmt, str};
 
-use rustc_arena::DroplessArena;
-use rustc_data_structures::fx::FxIndexSet;
+use rustc_data_structures::fx::FxHasher;
 use rustc_data_structures::stable_hasher::{
     HashStable, StableCompare, StableHasher, ToStableHashKey,
 };
@@ -2497,18 +2499,9 @@ impl Symbol {
         with_session_globals(|session_globals| session_globals.symbol_interner.intern(string))
     }
 
-    /// Access the underlying string. This is a slowish operation because it
-    /// requires locking the symbol interner.
-    ///
-    /// Note that the lifetime of the return value is a lie. It's not the same
-    /// as `&self`, but actually tied to the lifetime of the underlying
-    /// interner. Interners are long-lived, and there are very few of them, and
-    /// this function is typically used for short-lived things, so in practice
-    /// it works out ok.
+    /// Access the underlying string.
     pub fn as_str(&self) -> &str {
-        with_session_globals(|session_globals| unsafe {
-            std::mem::transmute::<&str, &str>(session_globals.symbol_interner.get(*self))
-        })
+        with_session_globals(|session_globals| session_globals.symbol_interner.get(*self))
     }
 
     pub fn as_u32(self) -> u32 {
@@ -2563,53 +2556,181 @@ impl StableCompare for Symbol {
     }
 }
 
-pub(crate) struct Interner(Lock<InternerInner>);
+// This is never de-initialized and stores interned &str in static storage.
+// Each str is stored length-prefixed (u32), and we allow for random-access indexing with a u32
+// index by direct lookup in the arena. Indices <2^16 are stored in a separate structure (they are
+// pre-allocated at dense addresses so we can't use the same lockless O(1) hack for them).
+static GLOBAL_ARENA: LazyLock<StringArena> = LazyLock::new(|| StringArena::new());
 
-// The `&'static str`s in this type actually point into the arena.
-//
-// This type is private to prevent accidentally constructing more than one
-// `Interner` on the same thread, which makes it easy to mix up `Symbol`s
-// between `Interner`s.
-struct InternerInner {
-    arena: DroplessArena,
-    strings: FxIndexSet<&'static str>,
+const CHUNK_SIZE: usize = 4 * 1024 * 1024;
+const CHUNKS: usize = (u32::MAX as usize).div_ceil(CHUNK_SIZE);
+
+struct StringChunk {
+    array: LazyLock<Box<[SyncUnsafeCell<u8>; CHUNK_SIZE]>>,
 }
 
-impl Interner {
-    fn prefill(init: &[&'static str]) -> Self {
-        Interner(Lock::new(InternerInner {
-            arena: Default::default(),
-            strings: init.iter().copied().collect(),
-        }))
+impl Default for StringChunk {
+    fn default() -> Self {
+        Self {
+            array: LazyLock::new(|| unsafe {
+                // SAFETY: Zero-init'd UnsafeCell<u8> is initialized and has no other invariants to
+                // worry about.
+                Box::new_zeroed().assume_init()
+            }),
+        }
     }
+}
 
-    #[inline]
-    fn intern(&self, string: &str) -> Symbol {
-        let mut inner = self.0.lock();
-        if let Some(idx) = inner.strings.get_index_of(string) {
-            return Symbol::new(idx as u32);
+struct StringArena {
+    chunks: [StringChunk; CHUNKS],
+    next_start: AtomicU32,
+    interned: elsa::sync::LockFreeFrozenVec<InternedString>,
+}
+
+#[derive(Copy, Clone)]
+struct InternedString {
+    start: u32,
+    length: u32,
+}
+
+impl StringArena {
+    fn new() -> Self {
+        StringArena {
+            chunks: std::array::from_fn(|_| StringChunk::default()),
+            next_start: AtomicU32::new(0),
+            interned: Default::default(),
+        }
+    }
+
+    fn next(previous: u32, length: u32) -> u32 {
+        let end = previous.checked_add(length).unwrap();
+        if previous / CHUNK_SIZE as u32 == end / CHUNK_SIZE as u32 {
+            end
+        } else {
+            // If we don't fit in the previous chunk, bump to the start of the next chunk, and set
+            // length to the end.
+            previous.next_multiple_of(CHUNK_SIZE as u32) + length
+        }
+    }
+
+    /// Copy the passed &str into the arena. Returns an index that can be passed to `get` to
+    /// retrieve the &str.
+    ///
+    /// u32 is guaranteed to be at least u16::MAX.
+    fn alloc(&self, s: &str) -> u32 {
+        let len = u32::try_from(s.len()).unwrap();
+        assert!(len < CHUNK_SIZE as u32);
+
+        let previous = self
+            .next_start
+            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |previous| {
+                Some(Self::next(previous, len))
+            })
+            .unwrap();
+        let end = Self::next(previous, len);
+        let start = end - len;
+
+        let chunk = LazyLock::force(&self.chunks[start as usize / CHUNK_SIZE].array);
+        let offset = start as usize % CHUNK_SIZE;
+
+        // SAFETY:
+        //
+        // * `next_start` only increases, and always uniquely allocates `len` bytes. No one can read
+        //   this memory yet as we haven't pushed yet to `interned`.
+        // * all chunks are zero-init'd at allocation: no uninitialized memory here.
+        let dst = unsafe {
+            std::slice::from_raw_parts_mut(
+                chunk.as_ptr().cast::<u8>().add(offset).cast_mut(),
+                s.len(),
+            )
+        };
+        dst.copy_from_slice(s.as_bytes());
+
+        let idx = self.interned.push(InternedString { start, length: len });
+
+        idx.try_into().unwrap()
+    }
+
+    /// Get the allocated string at the passed index.
+    ///
+    /// Note that this **does not** check that the passed index is actually an index returned by
+    /// `alloc`.
+    fn get(&self, idx: u32) -> &str {
+        let interned = self.interned.get(idx as usize).unwrap_or_else(|| {
+            panic!("non-interned symbol index: {idx}");
+        });
+
+        let chunk = LazyLock::force(&self.chunks[interned.start as usize / CHUNK_SIZE].array);
+        let offset = interned.start as usize % CHUNK_SIZE;
+
+        // We write the string into this memory range prior to pushing into `interned`, so this is
+        // guaranteed UTF-8 and initialized. `next_start` is strictly increasing so we never write
+        // twice.
+        unsafe {
+            std::str::from_raw_parts(
+                chunk.as_ptr().add(offset).cast::<u8>(),
+                interned.length as usize,
+            )
         }
+    }
+}
 
-        let string: &str = inner.arena.alloc_str(string);
+pub(crate) struct Interner(&'static [&'static str], Lock<InternerInner>);
 
-        // SAFETY: we can extend the arena allocation to `'static` because we
-        // only access these while the arena is still alive.
-        let string: &'static str = unsafe { &*(string as *const str) };
+struct InternerInner {
+    strings: hashbrown::HashTable<Symbol>,
+}
 
-        // This second hash table lookup can be avoided by using `RawEntryMut`,
-        // but this code path isn't hot enough for it to be worth it. See
-        // #91445 for details.
-        let (idx, is_new) = inner.strings.insert_full(string);
-        debug_assert!(is_new); // due to the get_index_of check above
+impl Interner {
+    fn prefill(init: &'static [&'static str]) -> Self {
+        assert!(init.len() < u16::MAX as usize);
+        let mut strings = hashbrown::HashTable::new();
+
+        for (idx, s) in init.iter().copied().enumerate() {
+            let mut hasher = FxHasher::default();
+            s.hash(&mut hasher);
+            let hash = hasher.finish();
+            strings.insert_unique(hash, Symbol::new(idx as u32), |val| {
+                // has to be from `init` because we haven't yet inserted anything except those.
+                BuildHasherDefault::<FxHasher>::default().hash_one(init[val.0.index()])
+            });
+        }
 
-        Symbol::new(idx as u32)
+        Interner(init, Lock::new(InternerInner { strings }))
+    }
+
+    #[inline]
+    fn intern(&self, string: &str) -> Symbol {
+        let hash = BuildHasherDefault::<FxHasher>::default().hash_one(string);
+        let mut inner = self.1.lock();
+        match inner.strings.find_entry(hash, |v| self.get(*v) == string) {
+            Ok(e) => return *e.get(),
+            Err(e) => {
+                let idx = GLOBAL_ARENA.alloc(string);
+                // Reserve 2^16 u32 indices -- these will be used for pre-filled interning where we
+                // have a dense SymbolIndex space. We could make this exact but it doesn't really
+                // matter in practice, we won't run out of symbol space.
+                let idx = u32::from(u16::MAX).checked_add(idx).unwrap();
+                let res = Symbol::new(idx as u32);
+
+                e.into_table().insert_unique(hash, res, |val| {
+                    BuildHasherDefault::<FxHasher>::default().hash_one(self.get(*val))
+                });
+
+                res
+            }
+        }
     }
 
     /// Get the symbol as a string.
     ///
     /// [`Symbol::as_str()`] should be used in preference to this function.
-    fn get(&self, symbol: Symbol) -> &str {
-        self.0.lock().strings.get_index(symbol.0.as_usize()).unwrap()
+    fn get(&self, symbol: Symbol) -> &'static str {
+        if let Some(interned) = symbol.0.as_u32().checked_sub(u32::from(u16::MAX)) {
+            GLOBAL_ARENA.get(interned)
+        } else {
+            self.0[symbol.0.index()]
+        }
     }
 }
 

compiler/rustc_span/src/symbol/tests.rs

@@ -4,15 +4,17 @@ use crate::create_default_session_globals_then;
 #[test]
 fn interner_tests() {
     let i = Interner::prefill(&[]);
+    let dog = i.intern("dog");
     // first one is zero:
-    assert_eq!(i.intern("dog"), Symbol::new(0));
+    assert_eq!(i.intern("dog"), dog);
     // re-use gets the same entry:
-    assert_eq!(i.intern("dog"), Symbol::new(0));
+    assert_eq!(i.intern("dog"), dog);
     // different string gets a different #:
-    assert_eq!(i.intern("cat"), Symbol::new(1));
-    assert_eq!(i.intern("cat"), Symbol::new(1));
+    let cat = i.intern("cat");
+    assert_eq!(i.intern("cat"), cat);
+    assert_eq!(i.intern("cat"), cat);
     // dog is still at zero
-    assert_eq!(i.intern("dog"), Symbol::new(0));
+    assert_eq!(i.intern("dog"), dog);
 }
 
 #[test]