Skip to content
Permalink
Browse files

Add DropArena and use it to allocate types with few allocations

  • Loading branch information...
Zoxc committed Mar 30, 2019
1 parent 4ccb9ae commit 43e33ea1fff581f74478b8181614b0a424dee2b0
Showing with 150 additions and 25 deletions.
  1. +140 −18 src/librustc/arena.rs
  2. +1 −1 src/librustc/query/mod.rs
  3. +7 −3 src/librustc/ty/codec.rs
  4. +2 −3 src/librustc_mir/transform/mod.rs
@@ -1,5 +1,10 @@
use arena::{TypedArena, DroplessArena};
use std::mem;
use std::ptr;
use std::slice;
use std::cell::RefCell;
use std::marker::PhantomData;
use smallvec::SmallVec;

#[macro_export]
macro_rules! arena_types {
@@ -9,29 +14,55 @@ macro_rules! arena_types {
rustc::hir::def_id::DefId,
rustc::ty::subst::SubstsRef<$tcx>
)>,
[few] mir_keys: rustc::util::nodemap::DefIdSet,
[decode] specialization_graph: rustc::traits::specialization_graph::Graph,
], $tcx);
)
}

macro_rules! arena_for_type {
([][$ty:ty]) => {
TypedArena<$ty>
};
([few $(, $attrs:ident)*][$ty:ty]) => {
PhantomData<$ty>
};
([$ignore:ident $(, $attrs:ident)*]$args:tt) => {
arena_for_type!([$($attrs),*]$args)
};
}

macro_rules! declare_arena {
([], [$($a:tt $name:ident: $ty:ty,)*], $tcx:lifetime) => {
#[derive(Default)]
pub struct Arena<$tcx> {
dropless: DroplessArena,
$($name: TypedArena<$ty>,)*
drop: DropArena,
$($name: arena_for_type!($a[$ty]),)*
}
}
}

macro_rules! which_arena_for_type {
([][$arena:expr]) => {
Some($arena)
};
([few$(, $attrs:ident)*][$arena:expr]) => {
None
};
([$ignore:ident$(, $attrs:ident)*]$args:tt) => {
which_arena_for_type!([$($attrs),*]$args)
};
}

macro_rules! impl_arena_allocatable {
([], [$($a:tt $name:ident: $ty:ty,)*], $tcx:lifetime) => {
$(
impl ArenaAllocatable for $ty {}
impl<$tcx> ArenaField<$tcx> for $ty {
unsafe impl<$tcx> ArenaField<$tcx> for $ty {
#[inline]
fn arena<'a>(arena: &'a Arena<$tcx>) -> &'a TypedArena<Self> {
&arena.$name
fn arena<'a>(_arena: &'a Arena<$tcx>) -> Option<&'a TypedArena<Self>> {
which_arena_for_type!($a[&_arena.$name])
}
}
)*
@@ -46,39 +77,130 @@ pub trait ArenaAllocatable {}

impl<T: Copy> ArenaAllocatable for T {}

pub trait ArenaField<'tcx>: Sized {
pub unsafe trait ArenaField<'tcx>: Sized {
/// Returns a specific arena to allocate from.
fn arena<'a>(arena: &'a Arena<'tcx>) -> &'a TypedArena<Self>;
/// If None is returned, the DropArena will be used.
fn arena<'a>(arena: &'a Arena<'tcx>) -> Option<&'a TypedArena<Self>>;
}

impl<'tcx, T> ArenaField<'tcx> for T {
unsafe impl<'tcx, T> ArenaField<'tcx> for T {
#[inline]
default fn arena<'a>(_: &'a Arena<'tcx>) -> &'a TypedArena<Self> {
default fn arena<'a>(_: &'a Arena<'tcx>) -> Option<&'a TypedArena<Self>> {
panic!()
}
}

impl<'tcx> Arena<'tcx> {
#[inline]
pub fn alloc<T: ArenaAllocatable>(&self, value: T) -> &mut T {
if mem::needs_drop::<T>() {
<T as ArenaField<'tcx>>::arena(self).alloc(value)
} else {
self.dropless.alloc(value)
if !mem::needs_drop::<T>() {
return self.dropless.alloc(value);
}
match <T as ArenaField<'tcx>>::arena(self) {
Some(arena) => arena.alloc(value),
None => unsafe { self.drop.alloc(value) },
}
}

pub fn alloc_from_iter<
T: ArenaAllocatable,
I: IntoIterator<Item = T>
>(
&self,
&'a self,
iter: I
) -> &mut [T] {
if mem::needs_drop::<T>() {
<T as ArenaField<'tcx>>::arena(self).alloc_from_iter(iter)
} else {
self.dropless.alloc_from_iter(iter)
) -> &'a mut [T] {
if !mem::needs_drop::<T>() {
return self.dropless.alloc_from_iter(iter);
}
match <T as ArenaField<'tcx>>::arena(self) {
Some(arena) => arena.alloc_from_iter(iter),
None => unsafe { self.drop.alloc_from_iter(iter) },
}
}
}

/// Calls the destructor for an object when dropped.
struct DropType {
drop_fn: unsafe fn(*mut u8),
obj: *mut u8,
}

unsafe fn drop_for_type<T>(to_drop: *mut u8) {
std::ptr::drop_in_place(to_drop as *mut T)
}

impl Drop for DropType {
fn drop(&mut self) {
unsafe {
(self.drop_fn)(self.obj)
}
}
}

/// An arena which can be used to allocate any type.
/// Allocating in this arena is unsafe since the type system
/// doesn't know which types it contains. In order to
/// allocate safetly, you must store a PhantomData<T>
/// alongside this arena for each type T you allocate.
#[derive(Default)]
struct DropArena {
/// A list of destructors to run when the arena drops.
/// Ordered so `destructors` gets dropped before the arena
/// since its destructor can reference memory in the arena.
destructors: RefCell<Vec<DropType>>,
arena: DroplessArena,
}

impl DropArena {
#[inline]
unsafe fn alloc<T>(&self, object: T) -> &mut T {
let mem = self.arena.alloc_raw(
mem::size_of::<T>(),
mem::align_of::<T>()
) as *mut _ as *mut T;
// Write into uninitialized memory.
ptr::write(mem, object);
let result = &mut *mem;
// Record the destructor after doing the allocation as that may panic
// and would cause `object`'s destuctor to run twice if it was recorded before
self.destructors.borrow_mut().push(DropType {
drop_fn: drop_for_type::<T>,
obj: result as *mut T as *mut u8,
});
result
}

#[inline]
unsafe fn alloc_from_iter<T, I: IntoIterator<Item = T>>(&self, iter: I) -> &mut [T] {
let mut vec: SmallVec<[_; 8]> = iter.into_iter().collect();
if vec.is_empty() {
return &mut [];
}
let len = vec.len();

let start_ptr = self.arena.alloc_raw(
len.checked_mul(mem::size_of::<T>()).unwrap(),
mem::align_of::<T>()
) as *mut _ as *mut T;

let mut destructors = self.destructors.borrow_mut();
// Reserve space for the destructors so we can't panic while adding them
destructors.reserve(len);

// Move the content to the arena by copying it and then forgetting
// the content of the SmallVec
vec.as_ptr().copy_to_nonoverlapping(start_ptr, len);
mem::forget(vec.drain());

// Record the destructors after doing the allocation as that may panic
// and would cause `object`'s destuctor to run twice if it was recorded before
for i in 0..len {
destructors.push(DropType {
drop_fn: drop_for_type::<T>,
obj: start_ptr.offset(i as isize) as *mut u8,
});
}

slice::from_raw_parts_mut(start_ptr, len)
}
}
@@ -84,7 +84,7 @@ rustc_queries! {
/// Set of all the `DefId`s in this crate that have MIR associated with
/// them. This includes all the body owners, but also things like struct
/// constructors.
query mir_keys(_: CrateNum) -> Lrc<DefIdSet> {
query mir_keys(_: CrateNum) -> &'tcx DefIdSet {
desc { "getting a list of all mir_keys" }
}

@@ -296,7 +296,9 @@ macro_rules! __impl_decoder_methods {

#[macro_export]
macro_rules! impl_arena_allocatable_decoder {
([$DecoderName:ident [$($typaram:tt),*]], [[decode] $name:ident: $ty:ty], $tcx:lifetime) => {
([]$args:tt) => {};
([decode $(, $attrs:ident)*]
[[$DecoderName:ident [$($typaram:tt),*]], [$name:ident: $ty:ty], $tcx:lifetime]) => {
impl<$($typaram),*> SpecializedDecoder<&$tcx $ty> for $DecoderName<$($typaram),*> {
#[inline]
fn specialized_decode(&mut self) -> Result<&$tcx $ty, Self::Error> {
@@ -311,14 +313,16 @@ macro_rules! impl_arena_allocatable_decoder {
}
}
};
([$DecoderName:ident [$($typaram:tt),*]], [[] $name:ident: $ty:ty], $tcx:lifetime) => {};
([$ignore:ident $(, $attrs:ident)*]$args:tt) => {
impl_arena_allocatable_decoder!([$($attrs),*]$args);
};
}

#[macro_export]
macro_rules! impl_arena_allocatable_decoders {
($args:tt, [$($a:tt $name:ident: $ty:ty,)*], $tcx:lifetime) => {
$(
impl_arena_allocatable_decoder!($args, [$a $name: $ty], $tcx);
impl_arena_allocatable_decoder!($a [$args, [$name: $ty], $tcx]);
)*
}
}
@@ -8,7 +8,6 @@ use rustc::ty::steal::Steal;
use rustc::hir;
use rustc::hir::intravisit::{self, Visitor, NestedVisitorMap};
use rustc::util::nodemap::DefIdSet;
use rustc_data_structures::sync::Lrc;
use std::borrow::Cow;
use syntax::ast;
use syntax_pos::Span;
@@ -59,7 +58,7 @@ fn is_mir_available<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, def_id: DefId) -> boo
/// Finds the full set of `DefId`s within the current crate that have
/// MIR associated with them.
fn mir_keys<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, krate: CrateNum)
-> Lrc<DefIdSet> {
-> &'tcx DefIdSet {
assert_eq!(krate, LOCAL_CRATE);

let mut set = DefIdSet::default();
@@ -94,7 +93,7 @@ fn mir_keys<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, krate: CrateNum)
set: &mut set,
}.as_deep_visitor());

Lrc::new(set)
tcx.arena.alloc(set)
}

fn mir_built<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, def_id: DefId) -> &'tcx Steal<Mir<'tcx>> {

0 comments on commit 43e33ea

Please sign in to comment.
You can’t perform that action at this time.