Auto merge of #136401 - Mark-Simulacrum:lockfree-as-str, r= · rust-lang/rust@bad4b51 (original) (raw)

`@@ -2,11 +2,13 @@

2

`//! allows bidirectional lookup; i.e., given a value, one can easily find the

3

`//! type, and vice versa.

4

5

use std::hash::{Hash, Hasher};

5

use std::alloc::Layout;

6

use std::hash::{BuildHasher, BuildHasherDefault, Hash, Hasher};

7

use std::sync::Mutex;

8

use std::sync::atomic::{AtomicPtr, Ordering};

6

9

`use std::{fmt, str};

7

10

8

use rustc_arena::DroplessArena;

9

use rustc_data_structures::fx::FxIndexSet;

11

use rustc_data_structures::fx::FxHasher;

10

12

`use rustc_data_structures::stable_hasher::{

11

13

`HashStable, StableCompare, StableHasher, ToStableHashKey,

12

14

`};

`@@ -2460,18 +2462,9 @@ impl Symbol {

2460

2462

`with_session_globals(|session_globals| session_globals.symbol_interner.intern(string))

2461

2463

2462

2464

2463

/// Access the underlying string. This is a slowish operation because it

2464

/// requires locking the symbol interner.

2465

///

2466

/// Note that the lifetime of the return value is a lie. It's not the same

2467

`` -

/// as &self, but actually tied to the lifetime of the underlying

2468

/// interner. Interners are long-lived, and there are very few of them, and

2469

/// this function is typically used for short-lived things, so in practice

2470

/// it works out ok.

2465

/// Access the underlying string.

2471

2466

`pub fn as_str(&self) -> &str {

2472

with_session_globals(|session_globals| unsafe {

2473

std::mem::transmute::<&str, &str>(session_globals.symbol_interner.get(*self))

2474

})

2467

with_session_globals(|session_globals| session_globals.symbol_interner.get(*self))

2475

2468

2476

2469

2477

2470

`pub fn as_u32(self) -> u32 {

`@@ -2526,53 +2519,137 @@ impl StableCompare for Symbol {

2526

2519

2527

2520

2528

2521

2529

pub(crate) struct Interner(Lock);

2522

// This is never de-initialized and stores interned &str in static storage.

2523

// Each str is stored length-prefixed (u32), and we allow for random-access indexing with a u32

2524

// index by direct lookup in the arena. Indices <2^16 are stored in a separate structure (they are

2525

// pre-allocated at dense addresses so we can't use the same lockless O(1) hack for them).

2526

static GLOBAL_ARENA: std::sync::LazyLock =

2527

std::sync::LazyLock::new(|| StringArena::new());

2528

+

2529

struct StringArena {

2530

base: *mut u8,

2531

end: *mut u8,

2532

write_at: AtomicPtr,

2533

`` +

// this guards writes to write_at, but not reads, which proceed lock-free. write_at is only moved

2534

// forward so this ends up being safe.

2535

writer: std::sync::Mutex<()>,

2536

}

2537

+

2538

unsafe impl Sync for StringArena {}

2539

unsafe impl Send for StringArena {}

2540

+

2541

impl StringArena {

2542

fn new() -> Self {

2543

unsafe {

2544

let layout =

2545

Layout::from_size_align(u32::MAX as usize, std::mem::align_of::()).unwrap();

2546

let allocation = std::alloc::alloc_zeroed(layout);

2547

if allocation.is_null() {

2548

std::alloc::handle_alloc_error(layout)

2549

}

2550

StringArena {

2551

base: allocation,

2552

end: allocation.add(layout.size()),

2553

// Reserve 2^16 u32 indices -- these will be used for pre-filled interning where we

2554

// have a dense SymbolIndex space. We could make this exact but it doesn't really

2555

// matter for this initial test anyway.

2556

write_at: AtomicPtr::new(allocation.add(u16::MAX as usize)),

2557

writer: Mutex::new(()),

2558

}

2559

}

2560

}

2561

+

2562

fn alloc(&self, s: &str) -> u32 {

2563

unsafe {

2564

let _guard = self.writer.lock().unwrap();

2565

// Allocate a chunk of the region, and fill it with the &str's length and bytes.

2566

let dst = self.write_at.load(Ordering::Acquire);

2567

// Assert we're in-bounds.

2568

assert!(

2569

dst.addr().checked_add(4).unwrap().checked_add(s.len()).unwrap() < self.end.addr()

2570

);

2571

dst.cast::().write_unaligned(s.len().try_into().unwrap());

2572

dst.add(4).copy_from_nonoverlapping(s.as_ptr(), s.len());

2573

+

2574

// Semantically this releases the memory for readers.

2575

self.write_at.store(dst.add(4 + s.len()), Ordering::Release);

2576

+

2577

dst.byte_offset_from(self.base).try_into().unwrap()

2578

}

2579

}

2580

+

2581

fn get(&self, idx: u32) -> &str {

2582

`` +

// SAFETY: write_at is only updated after writing to a given region, so we never have

2583

// concurrent writes to this memory. It's initialized at allocation (alloc_zeroed) and we

2584

// only write initialized bytes to it, so there's also never any uninit bytes in the region.

2585

let region = unsafe {

2586

let end = self.write_at.load(Ordering::Acquire);

2587

std::slice::from_raw_parts(self.base, end.offset_from(self.base) as usize)

2588

};

2589

+

2590

let len = u32::from_ne_bytes(region[idx as usize..idx as usize + 4].try_into().unwrap());

2591

let data = &region[idx as usize + 4..][..len as usize];

2592

+

2593

`` +

// SAFETY: We (in theory) could be passed a random idx into the memory region, so we need

2594

// to re-check that the region is actually UTF-8 before returning. If it is, then it is safe

2595

// to return &str, though it might not be a useful &str due to having near-random

2596

// contents.

2597

std::str::from_utf8(data).unwrap()

2598

}

2599

}

2600

+

2601

pub(crate) struct Interner(&'static [&'static str], Lock);

2530

2602

2531

`` -

// The &'static strs in this type actually point into the arena.

2532

2533

// This type is private to prevent accidentally constructing more than one

2534

`` -

// Interner on the same thread, which makes it easy to mix up Symbols

2535

`` -

// between Interners.

2536

2603

`struct InternerInner {

2537

arena: DroplessArena,

2538

strings: FxIndexSet<&'static str>,

2604

strings: hashbrown::HashTable,

2539

2605

2540

2606

2541

2607

`impl Interner {

2542

fn prefill(init: &[&'static str]) -> Self {

2543

Interner(Lock::new(InternerInner {

2544

arena: Default::default(),

2545

strings: init.iter().copied().collect(),

2546

}))

2608

fn prefill(init: &'static [&'static str]) -> Self {

2609

assert!(init.len() < u16::MAX as usize);

2610

let mut strings = hashbrown::HashTable::new();

2611

+

2612

for (idx, s) in init.iter().copied().enumerate() {

2613

let mut hasher = FxHasher::default();

2614

s.hash(&mut hasher);

2615

let hash = hasher.finish();

2616

strings.insert_unique(hash, Symbol::new(idx as u32), |val| {

2617

`` +

// has to be from init because we haven't yet inserted anything except those.

2618

BuildHasherDefault::::default().hash_one(init[val.0.index()])

2619

});

2620

}

2621

+

2622

Interner(init, Lock::new(InternerInner { strings }))

2547

2623

2548

2624

2549

2625

`#[inline]

2550

2626

`fn intern(&self, string: &str) -> Symbol {

2551

let mut inner = self.0.lock();

2552

if let Some(idx) = inner.strings.get_index_of(string) {

2553

return Symbol::new(idx as u32);

2627

let hash = BuildHasherDefault::::default().hash_one(string);

2628

let mut inner = self.1.lock();

2629

match inner.strings.find_entry(hash, |v| self.get(*v) == string) {

2630

Ok(e) => return *e.get(),

2631

Err(e) => {

2632

let idx = GLOBAL_ARENA.alloc(string);

2633

let res = Symbol::new(idx as u32);

2634

+

2635

e.into_table().insert_unique(hash, res, |val| {

2636

BuildHasherDefault::::default().hash_one(self.get(*val))

2637

});

2638

+

2639

res

2640

}

2554

2641

2555

-

2556

let string: &str = inner.arena.alloc_str(string);

2557

-

2558

`` -

// SAFETY: we can extend the arena allocation to 'static because we

2559

// only access these while the arena is still alive.

2560

let string: &'static str = unsafe { &*(string as *const str) };

2561

-

2562

`` -

// This second hash table lookup can be avoided by using RawEntryMut,

2563

// but this code path isn't hot enough for it to be worth it. See

2564

// #91445 for details.

2565

let (idx, is_new) = inner.strings.insert_full(string);

2566

debug_assert!(is_new); // due to the get_index_of check above

2567

-

2568

Symbol::new(idx as u32)

2569

2642

2570

2643

2571

2644

`/// Get the symbol as a string.

2572

2645

`///

2573

2646

`` /// [Symbol::as_str()] should be used in preference to this function.

2574

fn get(&self, symbol: Symbol) -> &str {

2575

self.0.lock().strings.get_index(symbol.0.as_usize()).unwrap()

2647

fn get(&self, symbol: Symbol) -> &'static str {

2648

if symbol.0.index() < u16::MAX as usize {

2649

self.0[symbol.0.index()]

2650

} else {

2651

GLOBAL_ARENA.get(symbol.0.as_u32())

2652

}

2576

2653

2577

2654

2578

2655