Remove slow HashSet during miri stack frame creation

src/librustc_mir/interpret/eval_context.rs

@@ -1,22 +1,23 @@
-use std::collections::HashSet;
 use std::fmt::Write;
 
 use rustc::hir::def_id::DefId;
+use rustc::hir::def::Def;
 use rustc::hir::map::definitions::DefPathData;
 use rustc::middle::const_val::{ConstVal, ErrKind};
 use rustc::mir;
 use rustc::ty::layout::{self, Size, Align, HasDataLayout, LayoutOf, TyLayout};
 use rustc::ty::subst::{Subst, Substs};
 use rustc::ty::{self, Ty, TyCtxt};
 use rustc::ty::maps::TyCtxtAt;
-use rustc_data_structures::indexed_vec::Idx;
+use rustc_data_structures::indexed_vec::{IndexVec, Idx};
 use rustc::middle::const_val::FrameInfo;
 use syntax::codemap::{self, Span};
 use syntax::ast::Mutability;
 use rustc::mir::interpret::{
     GlobalId, Value, Pointer, PrimVal, PrimValKind,
     EvalError, EvalResult, EvalErrorKind, MemoryPointer,
 };
+use std::mem;
 
 use super::{Place, PlaceExtra, Memory,
             HasMemory, MemoryKind,
@@ -71,12 +72,12 @@ pub struct Frame<'mir, 'tcx: 'mir> {
     pub return_place: Place,
 
     /// The list of locals for this stack frame, stored in order as
-    /// `[arguments..., variables..., temporaries...]`. The locals are stored as `Option<Value>`s.
+    /// `[return_ptr, arguments..., variables..., temporaries...]`. The locals are stored as `Option<Value>`s.
     /// `None` represents a local that is currently dead, while a live local
     /// can either directly contain `PrimVal` or refer to some part of an `Allocation`.
     ///
     /// Before being initialized, arguments are `Value::ByVal(PrimVal::Undef)` and other locals are `None`.
-    pub locals: Vec<Option<Value>>,
+    pub locals: IndexVec<mir::Local, Option<Value>>,
 
     ////////////////////////////////////////////////////////////////////////////////
     // Current position within the function
@@ -383,39 +384,29 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M
     ) -> EvalResult<'tcx> {
         ::log_settings::settings().indentation += 1;
 
-        /// Return the set of locals that have a storage annotation anywhere
-        fn collect_storage_annotations<'mir, 'tcx>(mir: &'mir mir::Mir<'tcx>) -> HashSet<mir::Local> {
-            use rustc::mir::StatementKind::*;
-
-            let mut set = HashSet::new();
-            for block in mir.basic_blocks() {
-                for stmt in block.statements.iter() {
-                    match stmt.kind {
-                        StorageLive(local) |
-                        StorageDead(local) => {
-                            set.insert(local);
+        let locals = if mir.local_decls.len() > 1 {
+            let mut locals = IndexVec::from_elem(Some(Value::ByVal(PrimVal::Undef)), &mir.local_decls);
+            match self.tcx.describe_def(instance.def_id()) {
+                // statics and constants don't have `Storage*` statements, no need to look for them
+                Some(Def::Static(..)) | Some(Def::Const(..)) | Some(Def::AssociatedConst(..)) => {},
+                _ => {
+                    trace!("push_stack_frame: {:?}: num_bbs: {}", span, mir.basic_blocks().len());
+                    for block in mir.basic_blocks() {
+                        for stmt in block.statements.iter() {
+                            use rustc::mir::StatementKind::{StorageDead, StorageLive};
+                            match stmt.kind {
+                                StorageLive(local) |
+                                StorageDead(local) => locals[local] = None,
+                                _ => {}
+                            }
                         }
-                        _ => {}
                     }
-                }
-            }
-            set
-        }
-
-        // Subtract 1 because `local_decls` includes the ReturnMemoryPointer, but we don't store a local
-        // `Value` for that.
-        let num_locals = mir.local_decls.len() - 1;
-
-        let locals = {
-            let annotated_locals = collect_storage_annotations(mir);
-            let mut locals = vec![None; num_locals];
-            for i in 0..num_locals {
-                let local = mir::Local::new(i + 1);
-                if !annotated_locals.contains(&local) {
-                    locals[i] = Some(Value::ByVal(PrimVal::Undef));
-                }
+                },
             }
             locals
+        } else {
+            // don't allocate at all for trivial constants
+            IndexVec::new()
         };
 
         self.stack.push(Frame {
@@ -973,8 +964,7 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M
     pub fn force_allocation(&mut self, place: Place) -> EvalResult<'tcx, Place> {
         let new_place = match place {
             Place::Local { frame, local } => {
-                // -1 since we don't store the return value
-                match self.stack[frame].locals[local.index() - 1] {
+                match self.stack[frame].locals[local] {
                     None => return err!(DeadLocal),
                     Some(Value::ByRef(ptr, align)) => {
                         Place::Ptr {
@@ -988,7 +978,7 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M
                         let ty = self.monomorphize(ty, self.stack[frame].instance.substs);
                         let layout = self.layout_of(ty)?;
                         let ptr = self.alloc_ptr(ty)?;
-                        self.stack[frame].locals[local.index() - 1] =
+                        self.stack[frame].locals[local] =
                             Some(Value::ByRef(ptr.into(), layout.align)); // it stays live
                         let place = Place::from_ptr(ptr, layout.align);
                         self.write_value(ValTy { value: val, ty }, place)?;
@@ -1702,13 +1692,11 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M
 
 impl<'mir, 'tcx> Frame<'mir, 'tcx> {
     pub fn get_local(&self, local: mir::Local) -> EvalResult<'tcx, Value> {
-        // Subtract 1 because we don't store a value for the ReturnPointer, the local with index 0.
-        self.locals[local.index() - 1].ok_or(EvalErrorKind::DeadLocal.into())
+        self.locals[local].ok_or(EvalErrorKind::DeadLocal.into())
     }
 
     fn set_local(&mut self, local: mir::Local, value: Value) -> EvalResult<'tcx> {
-        // Subtract 1 because we don't store a value for the ReturnPointer, the local with index 0.
-        match self.locals[local.index() - 1] {
+        match self.locals[local] {
             None => err!(DeadLocal),
             Some(ref mut local) => {
                 *local = value;
@@ -1717,20 +1705,17 @@ impl<'mir, 'tcx> Frame<'mir, 'tcx> {
         }
     }
 
-    pub fn storage_live(&mut self, local: mir::Local) -> EvalResult<'tcx, Option<Value>> {
+    pub fn storage_live(&mut self, local: mir::Local) -> Option<Value> {
         trace!("{:?} is now live", local);
 
-        let old = self.locals[local.index() - 1];
-        self.locals[local.index() - 1] = Some(Value::ByVal(PrimVal::Undef)); // StorageLive *always* kills the value that's currently stored
-        return Ok(old);
+        // StorageLive *always* kills the value that's currently stored
+        mem::replace(&mut self.locals[local], Some(Value::ByVal(PrimVal::Undef)))
     }
 
     /// Returns the old value of the local
-    pub fn storage_dead(&mut self, local: mir::Local) -> EvalResult<'tcx, Option<Value>> {
+    pub fn storage_dead(&mut self, local: mir::Local) -> Option<Value> {
         trace!("{:?} is now dead", local);
 
-        let old = self.locals[local.index() - 1];
-        self.locals[local.index() - 1] = None;
-        return Ok(old);
+        self.locals[local].take()
     }
 }

src/librustc_mir/interpret/memory.rs

@@ -1,12 +1,13 @@
 use byteorder::{ReadBytesExt, WriteBytesExt, LittleEndian, BigEndian};
-use std::collections::{btree_map, BTreeMap, HashMap, HashSet, VecDeque};
+use std::collections::{btree_map, BTreeMap, VecDeque};
 use std::{ptr, io};
 
 use rustc::ty::Instance;
 use rustc::ty::maps::TyCtxtAt;
 use rustc::ty::layout::{self, Align, TargetDataLayout};
 use syntax::ast::Mutability;
 
+use rustc_data_structures::fx::{FxHashSet, FxHashMap};
 use rustc::mir::interpret::{MemoryPointer, AllocId, Allocation, AccessKind, UndefMask, Value, Pointer,
                             EvalResult, PrimVal, EvalErrorKind};
 
@@ -33,15 +34,15 @@ pub struct Memory<'a, 'mir, 'tcx: 'a + 'mir, M: Machine<'mir, 'tcx>> {
     pub data: M::MemoryData,
 
     /// Helps guarantee that stack allocations aren't deallocated via `rust_deallocate`
-    alloc_kind: HashMap<AllocId, MemoryKind<M::MemoryKinds>>,
+    alloc_kind: FxHashMap<AllocId, MemoryKind<M::MemoryKinds>>,
 
     /// Actual memory allocations (arbitrary bytes, may contain pointers into other allocations).
-    alloc_map: HashMap<AllocId, Allocation>,
+    alloc_map: FxHashMap<AllocId, Allocation>,
 
     /// Actual memory allocations (arbitrary bytes, may contain pointers into other allocations).
     ///
     /// Stores statics while they are being processed, before they are interned and thus frozen
-    uninitialized_statics: HashMap<AllocId, Allocation>,
+    uninitialized_statics: FxHashMap<AllocId, Allocation>,
 
     /// The current stack frame.  Used to check accesses against locks.
     pub cur_frame: usize,
@@ -53,9 +54,9 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
     pub fn new(tcx: TyCtxtAt<'a, 'tcx, 'tcx>, data: M::MemoryData) -> Self {
         Memory {
             data,
-            alloc_kind: HashMap::new(),
-            alloc_map: HashMap::new(),
-            uninitialized_statics: HashMap::new(),
+            alloc_kind: FxHashMap::default(),
+            alloc_map: FxHashMap::default(),
+            uninitialized_statics: FxHashMap::default(),
             tcx,
             cur_frame: usize::max_value(),
         }
@@ -338,7 +339,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
         allocs.sort();
         allocs.dedup();
         let mut allocs_to_print = VecDeque::from(allocs);
-        let mut allocs_seen = HashSet::new();
+        let mut allocs_seen = FxHashSet::default();
 
         while let Some(id) = allocs_to_print.pop_front() {
             let mut msg = format!("Alloc {:<5} ", format!("{}:", id));

src/librustc_mir/interpret/step.rs

@@ -69,13 +69,13 @@ impl<'a, 'mir, 'tcx, M: Machine<'mir, 'tcx>> EvalContext<'a, 'mir, 'tcx, M> {
 
             // Mark locals as alive
             StorageLive(local) => {
-                let old_val = self.frame_mut().storage_live(local)?;
+                let old_val = self.frame_mut().storage_live(local);
                 self.deallocate_local(old_val)?;
             }
 
             // Mark locals as dead
             StorageDead(local) => {
-                let old_val = self.frame_mut().storage_dead(local)?;
+                let old_val = self.frame_mut().storage_dead(local);
                 self.deallocate_local(old_val)?;
             }