Skip to content

Commit 114a000

Browse files
ZoxcZoxc
committed
Add DropArena and use it to allocate types with few allocations
1 parent 85d8879 commit 114a000

File tree

4 files changed

+150
-25
lines changed

4 files changed

+150
-25
lines changed

‎src/librustc/arena.rs‎

Lines changed: 140 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,10 @@
11
use arena::{TypedArena, DroplessArena};
22
use std::mem;
3+
use std::ptr;
4+
use std::slice;
5+
use std::cell::RefCell;
6+
use std::marker::PhantomData;
7+
use smallvec::SmallVec;
38

49
#[macro_export]
510
macro_rules! arena_types {
@@ -9,29 +14,55 @@ macro_rules! arena_types {
914
rustc::hir::def_id::DefId,
1015
rustc::ty::subst::SubstsRef<$tcx>
1116
)>,
17+
[few] mir_keys: rustc::util::nodemap::DefIdSet,
1218
[decode] specialization_graph: rustc::traits::specialization_graph::Graph,
1319
], $tcx);
1420
)
1521
}
1622

23+
macro_rules! arena_for_type {
24+
([][$ty:ty]) => {
25+
TypedArena<$ty>
26+
};
27+
([few $(, $attrs:ident)*][$ty:ty]) => {
28+
PhantomData<$ty>
29+
};
30+
([$ignore:ident $(, $attrs:ident)*]$args:tt) => {
31+
arena_for_type!([$($attrs),*]$args)
32+
};
33+
}
34+
1735
macro_rules! declare_arena {
1836
([], [$($a:tt $name:ident: $ty:ty,)*], $tcx:lifetime) => {
1937
#[derive(Default)]
2038
pub struct Arena<$tcx> {
2139
dropless: DroplessArena,
22-
$($name: TypedArena<$ty>,)*
40+
drop: DropArena,
41+
$($name: arena_for_type!($a[$ty]),)*
2342
}
2443
}
2544
}
2645

46+
macro_rules! which_arena_for_type {
47+
([][$arena:expr]) => {
48+
Some($arena)
49+
};
50+
([few$(, $attrs:ident)*][$arena:expr]) => {
51+
None
52+
};
53+
([$ignore:ident$(, $attrs:ident)*]$args:tt) => {
54+
which_arena_for_type!([$($attrs),*]$args)
55+
};
56+
}
57+
2758
macro_rules! impl_arena_allocatable {
2859
([], [$($a:tt $name:ident: $ty:ty,)*], $tcx:lifetime) => {
2960
$(
3061
impl ArenaAllocatable for $ty {}
31-
impl<$tcx> ArenaField<$tcx> for $ty {
62+
unsafe impl<$tcx> ArenaField<$tcx> for $ty {
3263
#[inline]
33-
fn arena<'a>(arena: &'a Arena<$tcx>) -> &'a TypedArena<Self> {
34-
&arena.$name
64+
fn arena<'a>(_arena: &'a Arena<$tcx>) -> Option<&'a TypedArena<Self>> {
65+
which_arena_for_type!($a[&_arena.$name])
3566
}
3667
}
3768
)*
@@ -46,39 +77,130 @@ pub trait ArenaAllocatable {}
4677

4778
impl<T: Copy> ArenaAllocatable for T {}
4879

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

54-
impl<'tcx, T> ArenaField<'tcx> for T {
86+
unsafe impl<'tcx, T> ArenaField<'tcx> for T {
5587
#[inline]
56-
default fn arena<'a>(_: &'a Arena<'tcx>) -> &'a TypedArena<Self> {
88+
default fn arena<'a>(_: &'a Arena<'tcx>) -> Option<&'a TypedArena<Self>> {
5789
panic!()
5890
}
5991
}
6092

6193
impl<'tcx> Arena<'tcx> {
6294
#[inline]
6395
pub fn alloc<T: ArenaAllocatable>(&self, value: T) -> &mut T {
64-
if mem::needs_drop::<T>() {
65-
<T as ArenaField<'tcx>>::arena(self).alloc(value)
66-
} else {
67-
self.dropless.alloc(value)
96+
if !mem::needs_drop::<T>() {
97+
return self.dropless.alloc(value);
98+
}
99+
match <T as ArenaField<'tcx>>::arena(self) {
100+
Some(arena) => arena.alloc(value),
101+
None => unsafe { self.drop.alloc(value) },
68102
}
69103
}
70104

71105
pub fn alloc_from_iter<
72106
T: ArenaAllocatable,
73107
I: IntoIterator<Item = T>
74108
>(
75-
&self,
109+
&'a self,
76110
iter: I
77-
) -> &mut [T] {
78-
if mem::needs_drop::<T>() {
79-
<T as ArenaField<'tcx>>::arena(self).alloc_from_iter(iter)
80-
} else {
81-
self.dropless.alloc_from_iter(iter)
111+
) -> &'a mut [T] {
112+
if !mem::needs_drop::<T>() {
113+
return self.dropless.alloc_from_iter(iter);
114+
}
115+
match <T as ArenaField<'tcx>>::arena(self) {
116+
Some(arena) => arena.alloc_from_iter(iter),
117+
None => unsafe { self.drop.alloc_from_iter(iter) },
118+
}
119+
}
120+
}
121+
122+
/// Calls the destructor for an object when dropped.
123+
struct DropType {
124+
drop_fn: unsafe fn(*mut u8),
125+
obj: *mut u8,
126+
}
127+
128+
unsafe fn drop_for_type<T>(to_drop: *mut u8) {
129+
std::ptr::drop_in_place(to_drop as *mut T)
130+
}
131+
132+
impl Drop for DropType {
133+
fn drop(&mut self) {
134+
unsafe {
135+
(self.drop_fn)(self.obj)
136+
}
137+
}
138+
}
139+
140+
/// An arena which can be used to allocate any type.
141+
/// Allocating in this arena is unsafe since the type system
142+
/// doesn't know which types it contains. In order to
143+
/// allocate safetly, you must store a PhantomData<T>
144+
/// alongside this arena for each type T you allocate.
145+
#[derive(Default)]
146+
struct DropArena {
147+
/// A list of destructors to run when the arena drops.
148+
/// Ordered so `destructors` gets dropped before the arena
149+
/// since its destructor can reference memory in the arena.
150+
destructors: RefCell<Vec<DropType>>,
151+
arena: DroplessArena,
152+
}
153+
154+
impl DropArena {
155+
#[inline]
156+
unsafe fn alloc<T>(&self, object: T) -> &mut T {
157+
let mem = self.arena.alloc_raw(
158+
mem::size_of::<T>(),
159+
mem::align_of::<T>()
160+
) as *mut _ as *mut T;
161+
// Write into uninitialized memory.
162+
ptr::write(mem, object);
163+
let result = &mut *mem;
164+
// Record the destructor after doing the allocation as that may panic
165+
// and would cause `object`'s destuctor to run twice if it was recorded before
166+
self.destructors.borrow_mut().push(DropType {
167+
drop_fn: drop_for_type::<T>,
168+
obj: result as *mut T as *mut u8,
169+
});
170+
result
171+
}
172+
173+
#[inline]
174+
unsafe fn alloc_from_iter<T, I: IntoIterator<Item = T>>(&self, iter: I) -> &mut [T] {
175+
let mut vec: SmallVec<[_; 8]> = iter.into_iter().collect();
176+
if vec.is_empty() {
177+
return &mut [];
82178
}
179+
let len = vec.len();
180+
181+
let start_ptr = self.arena.alloc_raw(
182+
len.checked_mul(mem::size_of::<T>()).unwrap(),
183+
mem::align_of::<T>()
184+
) as *mut _ as *mut T;
185+
186+
let mut destructors = self.destructors.borrow_mut();
187+
// Reserve space for the destructors so we can't panic while adding them
188+
destructors.reserve(len);
189+
190+
// Move the content to the arena by copying it and then forgetting
191+
// the content of the SmallVec
192+
vec.as_ptr().copy_to_nonoverlapping(start_ptr, len);
193+
mem::forget(vec.drain());
194+
195+
// Record the destructors after doing the allocation as that may panic
196+
// and would cause `object`'s destuctor to run twice if it was recorded before
197+
for i in 0..len {
198+
destructors.push(DropType {
199+
drop_fn: drop_for_type::<T>,
200+
obj: start_ptr.offset(i as isize) as *mut u8,
201+
});
202+
}
203+
204+
slice::from_raw_parts_mut(start_ptr, len)
83205
}
84206
}

‎src/librustc/query/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ rustc_queries! {
8484
/// Set of all the `DefId`s in this crate that have MIR associated with
8585
/// them. This includes all the body owners, but also things like struct
8686
/// constructors.
87-
query mir_keys(_: CrateNum) -> Lrc<DefIdSet> {
87+
query mir_keys(_: CrateNum) -> &'tcx DefIdSet {
8888
desc { "getting a list of all mir_keys" }
8989
}
9090

‎src/librustc/ty/codec.rs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -296,7 +296,9 @@ macro_rules! __impl_decoder_methods {
296296

297297
#[macro_export]
298298
macro_rules! impl_arena_allocatable_decoder {
299-
([$DecoderName:ident [$($typaram:tt),*]], [[decode] $name:ident: $ty:ty], $tcx:lifetime) => {
299+
([]$args:tt) => {};
300+
([decode $(, $attrs:ident)*]
301+
[[$DecoderName:ident [$($typaram:tt),*]], [$name:ident: $ty:ty], $tcx:lifetime]) => {
300302
impl<$($typaram),*> SpecializedDecoder<&$tcx $ty> for $DecoderName<$($typaram),*> {
301303
#[inline]
302304
fn specialized_decode(&mut self) -> Result<&$tcx $ty, Self::Error> {
@@ -311,14 +313,16 @@ macro_rules! impl_arena_allocatable_decoder {
311313
}
312314
}
313315
};
314-
([$DecoderName:ident [$($typaram:tt),*]], [[] $name:ident: $ty:ty], $tcx:lifetime) => {};
316+
([$ignore:ident $(, $attrs:ident)*]$args:tt) => {
317+
impl_arena_allocatable_decoder!([$($attrs),*]$args);
318+
};
315319
}
316320

317321
#[macro_export]
318322
macro_rules! impl_arena_allocatable_decoders {
319323
($args:tt, [$($a:tt $name:ident: $ty:ty,)*], $tcx:lifetime) => {
320324
$(
321-
impl_arena_allocatable_decoder!($args, [$a $name: $ty], $tcx);
325+
impl_arena_allocatable_decoder!($a [$args, [$name: $ty], $tcx]);
322326
)*
323327
}
324328
}

‎src/librustc_mir/transform/mod.rs‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ use rustc::ty::steal::Steal;
88
use rustc::hir;
99
use rustc::hir::intravisit::{self, Visitor, NestedVisitorMap};
1010
use rustc::util::nodemap::DefIdSet;
11-
use rustc_data_structures::sync::Lrc;
1211
use std::borrow::Cow;
1312
use syntax::ast;
1413
use syntax_pos::Span;
@@ -59,7 +58,7 @@ fn is_mir_available<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, def_id: DefId) -> boo
5958
/// Finds the full set of `DefId`s within the current crate that have
6059
/// MIR associated with them.
6160
fn mir_keys<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, krate: CrateNum)
62-
-> Lrc<DefIdSet> {
61+
-> &'tcx DefIdSet {
6362
assert_eq!(krate, LOCAL_CRATE);
6463

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

97-
Lrc::new(set)
96+
tcx.arena.alloc(set)
9897
}
9998

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

0 commit comments

Comments
 (0)