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.

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

Remove the TypedArena::alloc_from_iter specialization. #116370

Merged
merged 2 commits into from
Oct 5, 2023
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.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
135 changes: 45 additions & 90 deletions compiler/rustc_arena/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
#![feature(dropck_eyepatch)]
#![feature(new_uninit)]
#![feature(maybe_uninit_slice)]
#![feature(min_specialization)]
#![feature(decl_macro)]
#![feature(pointer_byte_offsets)]
#![feature(rustc_attrs)]
Expand Down Expand Up @@ -44,23 +43,6 @@ fn outline<F: FnOnce() -> R, R>(f: F) -> R {
f()
}

/// An arena that can hold objects of only one type.
pub struct TypedArena<T> {
/// 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<Vec<ArenaChunk<T>>>,

/// Marker indicating that dropping the arena causes its owned
/// instances of `T` to be dropped.
_own: PhantomData<T>,
}

struct ArenaChunk<T = u8> {
/// The raw storage for the arena chunk.
storage: NonNull<[MaybeUninit<T>]>,
Expand Down Expand Up @@ -130,6 +112,23 @@ impl<T> ArenaChunk<T> {
const PAGE: usize = 4096;
const HUGE_PAGE: usize = 2 * 1024 * 1024;

/// An arena that can hold objects of only one type.
pub struct TypedArena<T> {
/// 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<Vec<ArenaChunk<T>>>,

/// Marker indicating that dropping the arena causes its owned
/// instances of `T` to be dropped.
_own: PhantomData<T>,
}

impl<T> Default for TypedArena<T> {
/// Creates a new `TypedArena`.
fn default() -> TypedArena<T> {
Expand All @@ -144,77 +143,6 @@ impl<T> Default for TypedArena<T> {
}
}

trait IterExt<T> {
fn alloc_from_iter(self, arena: &TypedArena<T>) -> &mut [T];
}

impl<I, T> IterExt<T> for I
where
I: IntoIterator<Item = T>,
{
// 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<T>) -> &mut [T] {
let vec: SmallVec<[_; 8]> = self.into_iter().collect();
vec.alloc_from_iter(arena)
}
}

impl<T, const N: usize> IterExt<T> for std::array::IntoIter<T, N> {
#[inline]
fn alloc_from_iter(self, arena: &TypedArena<T>) -> &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<T> IterExt<T> for Vec<T> {
#[inline]
fn alloc_from_iter(mut self, arena: &TypedArena<T>) -> &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<A: smallvec::Array> IterExt<A::Item> for SmallVec<A> {
#[inline]
fn alloc_from_iter(mut self, arena: &TypedArena<A::Item>) -> &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<T> TypedArena<T> {
/// Allocates an object in the `TypedArena`, returning a reference to it.
#[inline]
Expand Down Expand Up @@ -270,8 +198,35 @@ impl<T> TypedArena<T> {

#[inline]
pub fn alloc_from_iter<I: IntoIterator<Item = T>>(&self, iter: I) -> &mut [T] {
// This implementation is entirely separate to
// `DroplessIterator::alloc_from_iter`, even though conceptually they
// are the same.
//
// `DroplessIterator` (in the fast case) writes elements from the
// iterator one at a time into the allocated memory. That's easy
// because the elements don't implement `Drop`. But for `TypedArena`
// they do implement `Drop`, which means that if the iterator panics we
// could end up with some allocated-but-uninitialized elements, which
// will then cause UB in `TypedArena::drop`.
//
// Instead we use an approach where any iterator panic will occur
// before the memory is allocated. This function is much less hot than
// `DroplessArena::alloc_from_iter`, so it doesn't need to be
// hyper-optimized.
Copy link
Contributor

Choose a reason for hiding this comment

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

Do we need a fast path if !needs_drop::<T>()? That would call the drop-less codepath and eventually benefit from the opt.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Not sure I understand the question. Any T not implementing Drop will be allocated in the DroplessArena, as per this comment.

assert!(mem::size_of::<T>() != 0);
iter.alloc_from_iter(self)

let mut vec: SmallVec<[_; 8]> = iter.into_iter().collect();
if vec.is_empty() {
return &mut [];
}
// Move the content to the arena by copying and then forgetting it.
let len = vec.len();
let start_ptr = self.alloc_raw_slice(len);
unsafe {
vec.as_ptr().copy_to_nonoverlapping(start_ptr, len);
vec.set_len(0);
slice::from_raw_parts_mut(start_ptr, len)
}
}

/// Grows the arena.
Expand Down
47 changes: 18 additions & 29 deletions compiler/rustc_ast_lowering/src/format.rs
Original file line number Diff line number Diff line change
Expand Up @@ -410,15 +410,11 @@ fn expand_format_args<'hir>(
let format_options = use_format_options.then(|| {
// Generate:
// &[format_spec_0, format_spec_1, format_spec_2]
let elements: Vec<_> = fmt
.template
.iter()
.filter_map(|piece| {
let FormatArgsPiece::Placeholder(placeholder) = piece else { return None };
Some(make_format_spec(ctx, macsp, placeholder, &mut argmap))
})
.collect();
ctx.expr_array_ref(macsp, ctx.arena.alloc_from_iter(elements))
let elements = ctx.arena.alloc_from_iter(fmt.template.iter().filter_map(|piece| {
let FormatArgsPiece::Placeholder(placeholder) = piece else { return None };
Some(make_format_spec(ctx, macsp, placeholder, &mut argmap))
}));
ctx.expr_array_ref(macsp, elements)
});

let arguments = fmt.arguments.all_args();
Expand Down Expand Up @@ -477,10 +473,8 @@ fn expand_format_args<'hir>(
// <core::fmt::Argument>::new_debug(&arg2),
// …
// ]
let elements: Vec<_> = arguments
.iter()
.zip(argmap)
.map(|(arg, ((_, ty), placeholder_span))| {
let elements = ctx.arena.alloc_from_iter(arguments.iter().zip(argmap).map(
|(arg, ((_, ty), placeholder_span))| {
let placeholder_span =
placeholder_span.unwrap_or(arg.expr.span).with_ctxt(macsp.ctxt());
let arg_span = match arg.kind {
Expand All @@ -493,9 +487,9 @@ fn expand_format_args<'hir>(
hir::ExprKind::AddrOf(hir::BorrowKind::Ref, hir::Mutability::Not, arg),
));
make_argument(ctx, placeholder_span, ref_arg, ty)
})
.collect();
ctx.expr_array_ref(macsp, ctx.arena.alloc_from_iter(elements))
},
));
ctx.expr_array_ref(macsp, elements)
} else {
// Generate:
// &match (&arg0, &arg1, &…) {
Expand Down Expand Up @@ -528,19 +522,14 @@ fn expand_format_args<'hir>(
make_argument(ctx, placeholder_span, arg, ty)
},
));
let elements: Vec<_> = arguments
.iter()
.map(|arg| {
let arg_expr = ctx.lower_expr(&arg.expr);
ctx.expr(
arg.expr.span.with_ctxt(macsp.ctxt()),
hir::ExprKind::AddrOf(hir::BorrowKind::Ref, hir::Mutability::Not, arg_expr),
)
})
.collect();
let args_tuple = ctx
.arena
.alloc(ctx.expr(macsp, hir::ExprKind::Tup(ctx.arena.alloc_from_iter(elements))));
let elements = ctx.arena.alloc_from_iter(arguments.iter().map(|arg| {
let arg_expr = ctx.lower_expr(&arg.expr);
ctx.expr(
arg.expr.span.with_ctxt(macsp.ctxt()),
hir::ExprKind::AddrOf(hir::BorrowKind::Ref, hir::Mutability::Not, arg_expr),
)
}));
let args_tuple = ctx.arena.alloc(ctx.expr(macsp, hir::ExprKind::Tup(elements)));
let array = ctx.arena.alloc(ctx.expr(macsp, hir::ExprKind::Array(args)));
let match_arms = ctx.arena.alloc_from_iter([ctx.arm(args_pat, array)]);
let match_expr = ctx.arena.alloc(ctx.expr_match(
Expand Down
2 changes: 1 addition & 1 deletion compiler/rustc_hir_analysis/src/variance/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -192,5 +192,5 @@ fn variance_of_opaque(tcx: TyCtxt<'_>, item_def_id: LocalDefId) -> &[ty::Varianc
}
}
}
tcx.arena.alloc_from_iter(collector.variances.into_iter())
tcx.arena.alloc_from_iter(collector.variances)
}
14 changes: 8 additions & 6 deletions compiler/rustc_middle/src/ty/codec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -348,9 +348,10 @@ impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> Decodable<D> for ty::Const<'tcx> {

impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> RefDecodable<'tcx, D> for [ty::ValTree<'tcx>] {
fn decode(decoder: &mut D) -> &'tcx Self {
decoder.interner().arena.alloc_from_iter(
(0..decoder.read_usize()).map(|_| Decodable::decode(decoder)).collect::<Vec<_>>(),
)
decoder
.interner()
.arena
.alloc_from_iter((0..decoder.read_usize()).map(|_| Decodable::decode(decoder)))
}
}

Expand All @@ -368,9 +369,10 @@ impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> Decodable<D> for AdtDef<'tcx> {

impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> RefDecodable<'tcx, D> for [(ty::Clause<'tcx>, Span)] {
fn decode(decoder: &mut D) -> &'tcx Self {
decoder.interner().arena.alloc_from_iter(
(0..decoder.read_usize()).map(|_| Decodable::decode(decoder)).collect::<Vec<_>>(),
)
decoder
.interner()
.arena
.alloc_from_iter((0..decoder.read_usize()).map(|_| Decodable::decode(decoder)))
}
}

Expand Down
2 changes: 1 addition & 1 deletion compiler/rustc_trait_selection/src/traits/vtable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -316,7 +316,7 @@ fn vtable_entries<'tcx>(
dump_vtable_entries(tcx, sp, trait_ref, &entries);
}

tcx.arena.alloc_from_iter(entries.into_iter())
tcx.arena.alloc_from_iter(entries)
}

/// Find slot base for trait methods within vtable entries of another trait
Expand Down
2 changes: 1 addition & 1 deletion compiler/rustc_ty_utils/src/consts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ pub(crate) fn destructure_const<'tcx>(
_ => bug!("cannot destructure constant {:?}", const_),
};

let fields = tcx.arena.alloc_from_iter(fields.into_iter());
let fields = tcx.arena.alloc_from_iter(fields);

ty::DestructuredConst { variant, fields }
}
Expand Down