From 56a99fad5188d127f8a62315e66203eeb2d8ab10 Mon Sep 17 00:00:00 2001 From: Nicholas Nethercote Date: Fri, 29 Sep 2023 11:31:22 +1000 Subject: [PATCH 1/2] Make `TypedArena::alloc_from_iter` like `DroplessArena::alloc_from_iter`. `TypedArena` and `DroplessArena` both have an `alloc_from_impl` function, but they are implemented in completely different ways. - `TypedArena` uses the `IterExt` trait, which by default (a) collects the iterator's elements into a temporary `SmallVec`, (b) allocates memory, and (c) bulk-copies (and consumes) the elements into the arena memory. There are also specializations for `array::IntoIter`/`Vec`/`SmallVec` that avoid copying into a temporary `SmallVec`. - `DroplessArena` has two ways of doing it: - For known-sized iterators, it allocates memory and then copies items in one at a time. - Otherwise, it works much like the default case for `TypedArena`, collecting into a temporary `SmallVec` and then bulk-copying. Performance-wise, they are very similar. `DroplessArena` is a little more general, because it automatically works with any known-sized iterator type, while `TypedArena`'s approach requires specialization on a per-type basis. This commit removes `IterExt` and moves `DroplessArena::{write,alloc}_from_iter` so they can be used from `TypedArena`. This reduces the amount of code in this file. --- compiler/rustc_arena/src/lib.rs | 197 +++++++++++--------------------- 1 file changed, 64 insertions(+), 133 deletions(-) diff --git a/compiler/rustc_arena/src/lib.rs b/compiler/rustc_arena/src/lib.rs index 06eb8a185be96..d125fb644afcb 100644 --- a/compiler/rustc_arena/src/lib.rs +++ b/compiler/rustc_arena/src/lib.rs @@ -144,77 +144,6 @@ impl Default for TypedArena { } } -trait IterExt { - fn alloc_from_iter(self, arena: &TypedArena) -> &mut [T]; -} - -impl IterExt for I -where - I: IntoIterator, -{ - // This default collects into a `SmallVec` and then allocates by copying - // from it. The specializations below for types like `Vec` are more - // efficient, copying directly without the intermediate collecting step. - // This default could be made more efficient, like - // `DroplessArena::alloc_from_iter`, but it's not hot enough to bother. - #[inline] - default fn alloc_from_iter(self, arena: &TypedArena) -> &mut [T] { - let vec: SmallVec<[_; 8]> = self.into_iter().collect(); - vec.alloc_from_iter(arena) - } -} - -impl IterExt for std::array::IntoIter { - #[inline] - fn alloc_from_iter(self, arena: &TypedArena) -> &mut [T] { - let len = self.len(); - if len == 0 { - return &mut []; - } - // Move the content to the arena by copying and then forgetting it. - let start_ptr = arena.alloc_raw_slice(len); - unsafe { - self.as_slice().as_ptr().copy_to_nonoverlapping(start_ptr, len); - mem::forget(self); - slice::from_raw_parts_mut(start_ptr, len) - } - } -} - -impl IterExt for Vec { - #[inline] - fn alloc_from_iter(mut self, arena: &TypedArena) -> &mut [T] { - let len = self.len(); - if len == 0 { - return &mut []; - } - // Move the content to the arena by copying and then forgetting it. - let start_ptr = arena.alloc_raw_slice(len); - unsafe { - self.as_ptr().copy_to_nonoverlapping(start_ptr, len); - self.set_len(0); - slice::from_raw_parts_mut(start_ptr, len) - } - } -} - -impl IterExt for SmallVec { - #[inline] - fn alloc_from_iter(mut self, arena: &TypedArena) -> &mut [A::Item] { - let len = self.len(); - if len == 0 { - return &mut []; - } - // Move the content to the arena by copying and then forgetting it. - let start_ptr = arena.alloc_raw_slice(len); - unsafe { - self.as_ptr().copy_to_nonoverlapping(start_ptr, len); - self.set_len(0); - slice::from_raw_parts_mut(start_ptr, len) - } - } -} - impl TypedArena { /// Allocates an object in the `TypedArena`, returning a reference to it. #[inline] @@ -270,8 +199,7 @@ impl TypedArena { #[inline] pub fn alloc_from_iter>(&self, iter: I) -> &mut [T] { - assert!(mem::size_of::() != 0); - iter.alloc_from_iter(self) + alloc_from_iter(iter, |len| self.alloc_raw_slice(len)) } /// Grows the arena. @@ -527,76 +455,79 @@ impl DroplessArena { } } - /// # Safety - /// - /// The caller must ensure that `mem` is valid for writes up to - /// `size_of::() * len`. - #[inline] - unsafe fn write_from_iter>( - &self, - mut iter: I, - len: usize, - mem: *mut T, - ) -> &mut [T] { - let mut i = 0; - // Use a manual loop since LLVM manages to optimize it better for - // slice iterators - loop { - // SAFETY: The caller must ensure that `mem` is valid for writes up to - // `size_of::() * len`. - unsafe { - match iter.next() { - Some(value) if i < len => mem.add(i).write(value), - Some(_) | None => { - // We only return as many items as the iterator gave us, even - // though it was supposed to give us `len` - return slice::from_raw_parts_mut(mem, i); - } - } - } - i += 1; - } + pub fn alloc_from_iter>(&self, iter: I) -> &mut [T] { + alloc_from_iter(iter, |len| self.alloc_raw(Layout::array::(len).unwrap()) as *mut T) } +} - #[inline] - pub fn alloc_from_iter>(&self, iter: I) -> &mut [T] { - let iter = iter.into_iter(); - assert!(mem::size_of::() != 0); - assert!(!mem::needs_drop::()); +#[inline] +pub fn alloc_from_iter<'a, T, I, F>(iter: I, alloc_raw: F) -> &'a mut [T] +where + I: IntoIterator, + F: FnOnce(usize) -> *mut T, +{ + let iter = iter.into_iter(); + assert!(mem::size_of::() != 0); - let size_hint = iter.size_hint(); + match iter.size_hint() { + (min, Some(max)) if min == max => { + // We know the exact number of elements the iterator will produce here + let len = min; - match size_hint { - (min, Some(max)) if min == max => { - // We know the exact number of elements the iterator will produce here - let len = min; + if len == 0 { + return &mut []; + } - if len == 0 { + let mem = alloc_raw(len); + unsafe { write_from_iter(iter, len, mem) } + } + (_, _) => { + outline(move || -> &mut [T] { + let mut vec: SmallVec<[_; 8]> = iter.collect(); + if vec.is_empty() { return &mut []; } + // Move the content to the arena by copying it and then forgetting + // the content of the SmallVec + unsafe { + let len = vec.len(); + let start_ptr = alloc_raw(len); + vec.as_ptr().copy_to_nonoverlapping(start_ptr, len); + vec.set_len(0); + slice::from_raw_parts_mut(start_ptr, len) + } + }) + } + } +} - let mem = self.alloc_raw(Layout::array::(len).unwrap()) as *mut T; - unsafe { self.write_from_iter(iter, len, mem) } - } - (_, _) => { - outline(move || -> &mut [T] { - let mut vec: SmallVec<[_; 8]> = iter.collect(); - if vec.is_empty() { - return &mut []; - } - // Move the content to the arena by copying it and then forgetting - // the content of the SmallVec - unsafe { - let len = vec.len(); - let start_ptr = - self.alloc_raw(Layout::for_value::<[T]>(vec.as_slice())) as *mut T; - vec.as_ptr().copy_to_nonoverlapping(start_ptr, len); - vec.set_len(0); - slice::from_raw_parts_mut(start_ptr, len) - } - }) +/// # Safety +/// +/// The caller must ensure that `mem` is valid for writes up to +/// `size_of::() * len`. +#[inline] +unsafe fn write_from_iter<'a, T, I: Iterator>( + mut iter: I, + len: usize, + mem: *mut T, +) -> &'a mut [T] { + let mut i = 0; + // Use a manual loop since LLVM manages to optimize it better for + // slice iterators + loop { + // SAFETY: The caller must ensure that `mem` is valid for writes up to + // `size_of::() * len`. + unsafe { + match iter.next() { + Some(value) if i < len => mem.add(i).write(value), + Some(_) | None => { + // We only return as many items as the iterator gave us, even + // though it was supposed to give us `len` + return slice::from_raw_parts_mut(mem, i); + } } } + i += 1; } } From 2c40e8276e24319fc139fbf186b03697025fe7c4 Mon Sep 17 00:00:00 2001 From: Nicholas Nethercote Date: Mon, 2 Oct 2023 15:01:36 +1100 Subject: [PATCH 2/2] Move `TypedArena` declaration next to its impls. --- compiler/rustc_arena/src/lib.rs | 34 ++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/compiler/rustc_arena/src/lib.rs b/compiler/rustc_arena/src/lib.rs index d125fb644afcb..01209e5832d36 100644 --- a/compiler/rustc_arena/src/lib.rs +++ b/compiler/rustc_arena/src/lib.rs @@ -44,23 +44,6 @@ fn outline R, R>(f: F) -> R { f() } -/// An arena that can hold objects of only one type. -pub struct TypedArena { - /// A pointer to the next object to be allocated. - ptr: Cell<*mut T>, - - /// A pointer to the end of the allocated area. When this pointer is - /// reached, a new chunk is allocated. - end: Cell<*mut T>, - - /// A vector of arena chunks. - chunks: RefCell>>, - - /// Marker indicating that dropping the arena causes its owned - /// instances of `T` to be dropped. - _own: PhantomData, -} - struct ArenaChunk { /// The raw storage for the arena chunk. storage: NonNull<[MaybeUninit]>, @@ -130,6 +113,23 @@ impl ArenaChunk { const PAGE: usize = 4096; const HUGE_PAGE: usize = 2 * 1024 * 1024; +/// An arena that can hold objects of only one type. +pub struct TypedArena { + /// A pointer to the next object to be allocated. + ptr: Cell<*mut T>, + + /// A pointer to the end of the allocated area. When this pointer is + /// reached, a new chunk is allocated. + end: Cell<*mut T>, + + /// A vector of arena chunks. + chunks: RefCell>>, + + /// Marker indicating that dropping the arena causes its owned + /// instances of `T` to be dropped. + _own: PhantomData, +} + impl Default for TypedArena { /// Creates a new `TypedArena`. fn default() -> TypedArena {