Skip to content
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.

Sign up for GitHub

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

Jump to bottom

rustc_arena overhaul #116224

Merged
merged 9 commits into from
Oct 1, 2023
Merged

rustc_arena overhaul #116224

Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
79b11f0
rustc_arena: tweak some comments.
nnethercote Sep 19, 2023
a11f7e4
Remove some unnecessary lifetimes.
nnethercote Sep 19, 2023
0001edd
Inline and remove `DroplessArena::grow_and_alloc`.
nnethercote Sep 19, 2023
55de23e
Inline and remove `TypedArena::ensure_capacity`.
nnethercote Sep 19, 2023
51edc21
Remove `unsafe` from `TypedArena::alloc_raw_slice`.
nnethercote Sep 19, 2023
98d97b7
Use `Layout::new` consistently in `DroplessArena::alloc`.
nnethercote Sep 19, 2023
25407bc
Make `DroplessArena::alloc` call `DroplessArena::alloc_raw`.
nnethercote Sep 19, 2023
55a1a52
Reduce `grow_and_alloc_raw` to a single call site.
nnethercote Sep 19, 2023
bb5344a
Inline and remove `DroplessArena::alloc_raw_without_grow`.
nnethercote Sep 19, 2023
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
124 changes: 51 additions & 73 deletions compiler/rustc_arena/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,8 +172,8 @@ impl<T, const N: usize> IterExt<T> for std::array::IntoIter<T, N> {
return &mut [];
}
// Move the content to the arena by copying and then forgetting it.
let start_ptr = arena.alloc_raw_slice(len);
unsafe {
let start_ptr = arena.alloc_raw_slice(len);
self.as_slice().as_ptr().copy_to_nonoverlapping(start_ptr, len);
mem::forget(self);
slice::from_raw_parts_mut(start_ptr, len)
Expand All @@ -189,8 +189,8 @@ impl<T> IterExt<T> for Vec<T> {
return &mut [];
}
// Move the content to the arena by copying and then forgetting it.
let start_ptr = arena.alloc_raw_slice(len);
unsafe {
let start_ptr = arena.alloc_raw_slice(len);
self.as_ptr().copy_to_nonoverlapping(start_ptr, len);
self.set_len(0);
slice::from_raw_parts_mut(start_ptr, len)
Expand All @@ -206,8 +206,8 @@ impl<A: smallvec::Array> IterExt<A::Item> for SmallVec<A> {
return &mut [];
}
// Move the content to the arena by copying and then forgetting it.
let start_ptr = arena.alloc_raw_slice(len);
unsafe {
let start_ptr = arena.alloc_raw_slice(len);
self.as_ptr().copy_to_nonoverlapping(start_ptr, len);
self.set_len(0);
slice::from_raw_parts_mut(start_ptr, len)
Expand Down Expand Up @@ -250,25 +250,20 @@ impl<T> TypedArena<T> {
available_bytes >= additional_bytes
}

/// Ensures there's enough space in the current chunk to fit `len` objects.
#[inline]
fn ensure_capacity(&self, additional: usize) {
if !self.can_allocate(additional) {
self.grow(additional);
debug_assert!(self.can_allocate(additional));
}
}

#[inline]
unsafe fn alloc_raw_slice(&self, len: usize) -> *mut T {
fn alloc_raw_slice(&self, len: usize) -> *mut T {
assert!(mem::size_of::<T>() != 0);
assert!(len != 0);

self.ensure_capacity(len);
// Ensure the current chunk can fit `len` objects.
if !self.can_allocate(len) {
self.grow(len);
debug_assert!(self.can_allocate(len));
}

let start_ptr = self.ptr.get();
// SAFETY: `self.ensure_capacity` makes sure that there is enough space
// for `len` elements.
// SAFETY: `can_allocate`/`grow` ensures that there is enough space for
// `len` elements.
unsafe { self.ptr.set(start_ptr.add(len)) };
start_ptr
}
Expand Down Expand Up @@ -407,6 +402,8 @@ impl Default for DroplessArena {
#[inline]
fn default() -> DroplessArena {
DroplessArena {
// We set both `start` and `end` to 0 so that the first call to
// alloc() will trigger a grow().
start: Cell::new(ptr::null_mut()),
end: Cell::new(ptr::null_mut()),
chunks: Default::default(),
Expand All @@ -415,9 +412,11 @@ impl Default for DroplessArena {
}

impl DroplessArena {
#[inline(never)]
#[cold]
fn grow(&self, layout: Layout) {
// Add some padding so we can align `self.end` while
// stilling fitting in a `layout` allocation.
// still fitting in a `layout` allocation.
let additional = layout.size() + cmp::max(DROPLESS_ALIGNMENT, layout.align()) - 1;

unsafe {
Expand All @@ -441,7 +440,7 @@ impl DroplessArena {
let mut chunk = ArenaChunk::new(align_up(new_cap, PAGE));
self.start.set(chunk.start());

// Align the end to DROPLESS_ALIGNMENT
// Align the end to DROPLESS_ALIGNMENT.
let end = align_down(chunk.end().addr(), DROPLESS_ALIGNMENT);

// Make sure we don't go past `start`. This should not happen since the allocation
Expand All @@ -454,69 +453,48 @@ impl DroplessArena {
}
}

#[inline(never)]
#[cold]
fn grow_and_alloc_raw(&self, layout: Layout) -> *mut u8 {
self.grow(layout);
self.alloc_raw_without_grow(layout).unwrap()
}

#[inline(never)]
#[cold]
fn grow_and_alloc<T>(&self) -> *mut u8 {
self.grow_and_alloc_raw(Layout::new::<T>())
}

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

// Align allocated bytes so that `self.end` stays aligned to DROPLESS_ALIGNMENT
let bytes = align_up(layout.size(), DROPLESS_ALIGNMENT);

// Tell LLVM that `end` is aligned to DROPLESS_ALIGNMENT
unsafe { intrinsics::assume(end == align_down(end, DROPLESS_ALIGNMENT)) };

let new_end = align_down(end.checked_sub(bytes)?, layout.align());
if start <= new_end {
let new_end = old_end.with_addr(new_end);
// `new_end` is aligned to DROPLESS_ALIGNMENT as `align_down` preserves alignment
// as both `end` and `bytes` are already aligned to DROPLESS_ALIGNMENT.
self.end.set(new_end);
Some(new_end)
} else {
None
}
}

#[inline]
pub fn alloc_raw(&self, layout: Layout) -> *mut u8 {
assert!(layout.size() != 0);
if let Some(a) = self.alloc_raw_without_grow(layout) {
return a;

// This loop executes once or twice: if allocation fails the first
// time, the `grow` ensures it will succeed the second time.
loop {
Copy link
Contributor

Choose a reason for hiding this comment

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

I'm not a huge fan of this loop. If there's some logic error, it may run more than twice and I'd prefer a panic in that case. The loop variant before #108693 is a bit better as it avoids the large loop body.

Copy link
Contributor Author

@nnethercote nnethercote Sep 29, 2023 •
edited
Loading

Choose a reason for hiding this comment

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

I'm not a huge fan of this loop. If there's some logic error, it may run more than twice and I'd prefer a panic in that case.

Interesting that the loop was used before #108693, I didn't know about that. If that's how it was before, and I changed it back here, that suggests it's somehow a natural way to do it. I'm not worried about a logic error, because (a) the code used to be like that, and (b) this code is hot and so naturally gets high levels of testing.

The loop variant before #108693 is a bit better as it avoids the large loop body.

The "large" loop body is less than 30 lines, which isn't very large.

The current code has all these little functions, several with a single call site, including:

  • grow
  • grow_and_alloc_raw
  • grow_and_alloc
  • alloc_raw_without_grow
  • alloc_raw
  • alloc

The new code has just:

  • grow
  • alloc_raw
  • alloc

I find this much easier to understand. When I read the current code I had trouble keeping straight all these different functions, I had to write them all down and take notes to understand them.

let start = self.start.get().addr();
let old_end = self.end.get();
let end = old_end.addr();

// Align allocated bytes so that `self.end` stays aligned to
// DROPLESS_ALIGNMENT.
let bytes = align_up(layout.size(), DROPLESS_ALIGNMENT);

// Tell LLVM that `end` is aligned to DROPLESS_ALIGNMENT.
unsafe { intrinsics::assume(end == align_down(end, DROPLESS_ALIGNMENT)) };

if let Some(sub) = end.checked_sub(bytes) {
let new_end = align_down(sub, layout.align());
if start <= new_end {
let new_end = old_end.with_addr(new_end);
// `new_end` is aligned to DROPLESS_ALIGNMENT as `align_down`
// preserves alignment as both `end` and `bytes` are already
// aligned to DROPLESS_ALIGNMENT.
self.end.set(new_end);
return new_end;
}
}

// No free space left. Allocate a new chunk to satisfy the request.
// On failure the grow will panic or abort.
self.grow(layout);
}
// No free space left. Allocate a new chunk to satisfy the request.
// On failure the grow will panic or abort.
self.grow_and_alloc_raw(layout)
}

#[inline]
pub fn alloc<T>(&self, object: T) -> &mut T {
assert!(!mem::needs_drop::<T>());
assert!(mem::size_of::<T>() != 0);

let mem = if let Some(a) = self.alloc_raw_without_grow(Layout::for_value::<T>(&object)) {
a
} else {
// No free space left. Allocate a new chunk to satisfy the request.
// On failure the grow will panic or abort.
self.grow_and_alloc::<T>()
} as *mut T;
let mem = self.alloc_raw(Layout::new::<T>()) as *mut T;

unsafe {
// Write into uninitialized memory.
Expand Down Expand Up @@ -713,10 +691,10 @@ pub macro declare_arena([$($a:tt $name:ident: $ty:ty,)*]) {
}

#[allow(clippy::mut_from_ref)]
pub fn alloc_from_iter<'a, T: ArenaAllocatable<'tcx, C>, C>(
&'a self,
pub fn alloc_from_iter<T: ArenaAllocatable<'tcx, C>, C>(
&self,
iter: impl ::std::iter::IntoIterator<Item = T>,
) -> &'a mut [T] {
) -> &mut [T] {
T::allocate_from_iter(self, iter)
}
}
Expand Down
Loading