Add DropArena and use it to allocate types with few allocations

4 files changed, 150 insertions, 25 deletions

diff --git a/src/librustc/arena.rs b/src/librustc/arena.rs
index a4c3326c848..cafa973ea8b 100644
--- a/src/librustc/arena.rs
+++ b/src/librustc/arena.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,14 +77,15 @@ 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!()
     }
 }
@@ -61,10 +93,12 @@ impl<'tcx, T> ArenaField<'tcx> for T {
 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) },
         }
     }
 
@@ -72,13 +106,101 @@ impl<'tcx> Arena<'tcx> {
         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)
     }
 }
diff --git a/src/librustc/query/mod.rs b/src/librustc/query/mod.rs
index a62a2939544..329add5edd0 100644
--- a/src/librustc/query/mod.rs
+++ b/src/librustc/query/mod.rs
@@ -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" }
         }
 
diff --git a/src/librustc/ty/codec.rs b/src/librustc/ty/codec.rs
index 3046d53086c..a76cc3dfdec 100644
--- a/src/librustc/ty/codec.rs
+++ b/src/librustc/ty/codec.rs
@@ -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]);
         )*
     }
 }
diff --git a/src/librustc_mir/transform/mod.rs b/src/librustc_mir/transform/mod.rs
index 8df0d72407b..27cb87f5dca 100644
--- a/src/librustc_mir/transform/mod.rs
+++ b/src/librustc_mir/transform/mod.rs
@@ -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>> {