Skip to content

proc_macro/bridge: stop using a remote object handle for proc_macro Ident and Literal #98189

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jul 19, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
402 changes: 173 additions & 229 deletions compiler/rustc_expand/src/proc_macro_server.rs

Large diffs are not rendered by default.

113 changes: 113 additions & 0 deletions library/proc_macro/src/bridge/arena.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,113 @@
//! A minimal arena allocator inspired by `rustc_arena::DroplessArena`.
//!
//! This is unfortunately a minimal re-implementation rather than a dependency
//! as it is difficult to depend on crates from within `proc_macro`, due to it
//! being built at the same time as `std`.

use std::cell::{Cell, RefCell};
use std::cmp;
use std::mem::MaybeUninit;
use std::ops::Range;
use std::ptr;
use std::slice;
use std::str;

// The arenas start with PAGE-sized chunks, and then each new chunk is twice as
// big as its predecessor, up until we reach HUGE_PAGE-sized chunks, whereupon
// we stop growing. This scales well, from arenas that are barely used up to
// arenas that are used for 100s of MiBs. Note also that the chosen sizes match
// the usual sizes of pages and huge pages on Linux.
const PAGE: usize = 4096;
const HUGE_PAGE: usize = 2 * 1024 * 1024;

/// A minimal arena allocator inspired by `rustc_arena::DroplessArena`.
///
/// This is unfortunately a complete re-implementation rather than a dependency
/// as it is difficult to depend on crates from within `proc_macro`, due to it
/// being built at the same time as `std`.
///
/// This arena doesn't have support for allocating anything other than byte
/// slices, as that is all that is necessary.
pub(crate) struct Arena {
start: Cell<*mut MaybeUninit<u8>>,
end: Cell<*mut MaybeUninit<u8>>,
chunks: RefCell<Vec<Box<[MaybeUninit<u8>]>>>,
}

impl Arena {
pub(crate) fn new() -> Self {
Arena {
start: Cell::new(ptr::null_mut()),
end: Cell::new(ptr::null_mut()),
chunks: RefCell::new(Vec::new()),
}
}

/// Add a new chunk with at least `additional` free bytes.
#[inline(never)]
#[cold]
fn grow(&self, additional: usize) {
let mut chunks = self.chunks.borrow_mut();
let mut new_cap;
if let Some(last_chunk) = chunks.last_mut() {
// If the previous chunk's len is less than HUGE_PAGE
// bytes, then this chunk will be least double the previous
// chunk's size.
new_cap = last_chunk.len().min(HUGE_PAGE / 2);
new_cap *= 2;
} else {
new_cap = PAGE;
}
// Also ensure that this chunk can fit `additional`.
new_cap = cmp::max(additional, new_cap);

let mut chunk = Box::new_uninit_slice(new_cap);
let Range { start, end } = chunk.as_mut_ptr_range();
self.start.set(start);
self.end.set(end);
chunks.push(chunk);
}

/// Allocates a byte slice with specified size from the current memory
/// chunk. Returns `None` if there is no free space left to satisfy the
/// request.
fn alloc_raw_without_grow(&self, bytes: usize) -> Option<&mut [MaybeUninit<u8>]> {
let start = self.start.get().addr();
let old_end = self.end.get();
let end = old_end.addr();

let new_end = end.checked_sub(bytes)?;
if start <= new_end {
let new_end = old_end.with_addr(new_end);
self.end.set(new_end);
// SAFETY: `bytes` bytes starting at `new_end` were just reserved.
Some(unsafe { slice::from_raw_parts_mut(new_end, bytes) })
} else {
None
}
}

fn alloc_raw(&self, bytes: usize) -> &mut [MaybeUninit<u8>] {
if bytes == 0 {
return &mut [];
}

loop {
if let Some(a) = self.alloc_raw_without_grow(bytes) {
break a;
}
// No free space left. Allocate a new chunk to satisfy the request.
// On failure the grow will panic or abort.
self.grow(bytes);
}
}

pub(crate) fn alloc_str<'a>(&'a self, string: &str) -> &'a mut str {
let alloc = self.alloc_raw(string.len());
let bytes = MaybeUninit::write_slice(alloc, string.as_bytes());

// SAFETY: we convert from `&str` to `&[u8]`, clone it into the arena,
// and immediately convert the clone back to `&str`.
unsafe { str::from_utf8_unchecked_mut(bytes) }
}
}
30 changes: 9 additions & 21 deletions library/proc_macro/src/bridge/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -175,13 +175,11 @@ define_handles! {
'owned:
FreeFunctions,
TokenStream,
Literal,
SourceFile,
MultiSpan,
Diagnostic,

'interned:
Ident,
Span,
}

Expand All @@ -197,25 +195,6 @@ impl Clone for TokenStream {
}
}

impl Clone for Literal {
fn clone(&self) -> Self {
self.clone()
}
}

impl fmt::Debug for Literal {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Literal")
// format the kind without quotes, as in `kind: Float`
.field("kind", &format_args!("{}", &self.debug_kind()))
.field("symbol", &self.symbol())
// format `Some("...")` on one line even in {:#?} mode
.field("suffix", &format_args!("{:?}", &self.suffix()))
.field("span", &self.span())
.finish()
}
}

impl Clone for SourceFile {
fn clone(&self) -> Self {
self.clone()
Expand All @@ -242,6 +221,8 @@ impl fmt::Debug for Span {
}
}

pub(crate) use super::symbol::Symbol;

macro_rules! define_client_side {
($($name:ident {
$(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) $(-> $ret_ty:ty)*;)*
Expand Down Expand Up @@ -405,6 +386,9 @@ fn run_client<A: for<'a, 's> DecodeMut<'a, 's, ()>, R: Encode<()>>(
panic::catch_unwind(panic::AssertUnwindSafe(|| {
maybe_install_panic_hook(force_show_panics);

// Make sure the symbol store is empty before decoding inputs.
Symbol::invalidate_all();

let reader = &mut &buf[..];
let (globals, input) = <(ExpnGlobals<Span>, A)>::decode(reader, &mut ());

Expand Down Expand Up @@ -438,6 +422,10 @@ fn run_client<A: for<'a, 's> DecodeMut<'a, 's, ()>, R: Encode<()>>(
buf.clear();
Err::<(), _>(e).encode(&mut buf, &mut ());
});

// Now that a response has been serialized, invalidate all symbols
// registered with the interner.
Symbol::invalidate_all();
buf
}

Expand Down
117 changes: 117 additions & 0 deletions library/proc_macro/src/bridge/fxhash.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,117 @@
//! This is a copy of the `rustc_hash` crate, adapted to work as a module.
//!
//! If in the future it becomes more reasonable to add dependencies to
//! `proc_macro`, this module should be removed and replaced with a dependency
//! on the `rustc_hash` crate.

use std::collections::HashMap;
use std::convert::TryInto;
use std::default::Default;
use std::hash::BuildHasherDefault;
use std::hash::Hasher;
use std::mem::size_of;
use std::ops::BitXor;

/// Type alias for a hashmap using the `fx` hash algorithm.
pub type FxHashMap<K, V> = HashMap<K, V, BuildHasherDefault<FxHasher>>;

/// A speedy hash algorithm for use within rustc. The hashmap in liballoc
/// by default uses SipHash which isn't quite as speedy as we want. In the
/// compiler we're not really worried about DOS attempts, so we use a fast
/// non-cryptographic hash.
///
/// This is the same as the algorithm used by Firefox -- which is a homespun
/// one not based on any widely-known algorithm -- though modified to produce
/// 64-bit hash values instead of 32-bit hash values. It consistently
/// out-performs an FNV-based hash within rustc itself -- the collision rate is
/// similar or slightly worse than FNV, but the speed of the hash function
/// itself is much higher because it works on up to 8 bytes at a time.
pub struct FxHasher {
hash: usize,
}

#[cfg(target_pointer_width = "32")]
const K: usize = 0x9e3779b9;
#[cfg(target_pointer_width = "64")]
const K: usize = 0x517cc1b727220a95;

impl Default for FxHasher {
#[inline]
fn default() -> FxHasher {
FxHasher { hash: 0 }
}
}

impl FxHasher {
#[inline]
fn add_to_hash(&mut self, i: usize) {
self.hash = self.hash.rotate_left(5).bitxor(i).wrapping_mul(K);
}
}

impl Hasher for FxHasher {
#[inline]
fn write(&mut self, mut bytes: &[u8]) {
#[cfg(target_pointer_width = "32")]
let read_usize = |bytes: &[u8]| u32::from_ne_bytes(bytes[..4].try_into().unwrap());
#[cfg(target_pointer_width = "64")]
let read_usize = |bytes: &[u8]| u64::from_ne_bytes(bytes[..8].try_into().unwrap());

let mut hash = FxHasher { hash: self.hash };
assert!(size_of::<usize>() <= 8);
while bytes.len() >= size_of::<usize>() {
hash.add_to_hash(read_usize(bytes) as usize);
bytes = &bytes[size_of::<usize>()..];
}
if (size_of::<usize>() > 4) && (bytes.len() >= 4) {
hash.add_to_hash(u32::from_ne_bytes(bytes[..4].try_into().unwrap()) as usize);
bytes = &bytes[4..];
}
if (size_of::<usize>() > 2) && bytes.len() >= 2 {
hash.add_to_hash(u16::from_ne_bytes(bytes[..2].try_into().unwrap()) as usize);
bytes = &bytes[2..];
}
if (size_of::<usize>() > 1) && bytes.len() >= 1 {
hash.add_to_hash(bytes[0] as usize);
}
self.hash = hash.hash;
}

#[inline]
fn write_u8(&mut self, i: u8) {
self.add_to_hash(i as usize);
}

#[inline]
fn write_u16(&mut self, i: u16) {
self.add_to_hash(i as usize);
}

#[inline]
fn write_u32(&mut self, i: u32) {
self.add_to_hash(i as usize);
}

#[cfg(target_pointer_width = "32")]
#[inline]
fn write_u64(&mut self, i: u64) {
self.add_to_hash(i as usize);
self.add_to_hash((i >> 32) as usize);
}

#[cfg(target_pointer_width = "64")]
#[inline]
fn write_u64(&mut self, i: u64) {
self.add_to_hash(i as usize);
}

#[inline]
fn write_usize(&mut self, i: usize) {
self.add_to_hash(i);
}

#[inline]
fn finish(&self) -> u64 {
self.hash as u64
}
}
26 changes: 6 additions & 20 deletions library/proc_macro/src/bridge/handle.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
//! Server-side handles and storage for per-handle data.

use std::collections::{BTreeMap, HashMap};
use std::hash::{BuildHasher, Hash};
use std::collections::BTreeMap;
use std::hash::Hash;
use std::num::NonZeroU32;
use std::ops::{Index, IndexMut};
use std::sync::atomic::{AtomicUsize, Ordering};

use super::fxhash::FxHashMap;

pub(super) type Handle = NonZeroU32;

/// A store that associates values of type `T` with numeric handles. A value can
Expand Down Expand Up @@ -51,31 +53,15 @@ impl<T> IndexMut<Handle> for OwnedStore<T> {
}
}

// HACK(eddyb) deterministic `std::collections::hash_map::RandomState` replacement
// that doesn't require adding any dependencies to `proc_macro` (like `rustc-hash`).
#[derive(Clone)]
struct NonRandomState;

impl BuildHasher for NonRandomState {
type Hasher = std::collections::hash_map::DefaultHasher;
#[inline]
fn build_hasher(&self) -> Self::Hasher {
Self::Hasher::new()
}
}

/// Like `OwnedStore`, but avoids storing any value more than once.
pub(super) struct InternedStore<T: 'static> {
owned: OwnedStore<T>,
interner: HashMap<T, Handle, NonRandomState>,
interner: FxHashMap<T, Handle>,
Comment on lines -70 to +59
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two thoughts here:

  • any chance this is significant enough on its own to get tested in a separate PR?
  • just realized that the "base index" trick could work here to allow a simple Vec backing storage - I guess the handle counters would need to become thread-local in that case?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure how valuable it would be in a separate PR given that the types being hashed are presumably simpler without this change, but I suppose we could test it if we wanted to. I've kept that part in a separate commit so it can easily be split off.

I think the simple Vec storage might only work for InternedStore, and it'd require a thread-local counter and only one server per-thread as well, but that's probably reasonable. It seems unlikely we'd get into a situation where a single client thread is shared by multiple different server threads, which could cause a setup like that to break for server-side interning.

The Vec approach with a non-atomic counter feels the most reliable for client-side unique name generation.

}

impl<T: Copy + Eq + Hash> InternedStore<T> {
pub(super) fn new(counter: &'static AtomicUsize) -> Self {
InternedStore {
owned: OwnedStore::new(counter),
interner: HashMap::with_hasher(NonRandomState),
}
InternedStore { owned: OwnedStore::new(counter), interner: FxHashMap::default() }
}

pub(super) fn alloc(&mut self, x: T) -> Handle {
Expand Down
Loading