abstract mono_hash_map through a trait, only miri actually needs the fancy one

RalfJung · RalfJung · commit 83667d64a2ae · 2018-10-10T10:08:17.000+02:00
diff --git a/src/librustc_mir/const_eval.rs b/src/librustc_mir/const_eval.rs
@@ -12,7 +12,9 @@
 
 use std::fmt;
 use std::error::Error;
-use std::borrow::Cow;
+use std::borrow::{Borrow, Cow};
+use std::hash::Hash;
+use std::collections::hash_map::Entry;
 
 use rustc::hir::{self, def_id::DefId};
 use rustc::mir::interpret::ConstEvalErr;
@@ -21,13 +23,14 @@ use rustc::ty::{self, TyCtxt, Instance, query::TyCtxtAt};
 use rustc::ty::layout::{self, LayoutOf, TyLayout};
 use rustc::ty::subst::Subst;
 use rustc_data_structures::indexed_vec::IndexVec;
+use rustc_data_structures::fx::FxHashMap;
 
 use syntax::ast::Mutability;
 use syntax::source_map::{Span, DUMMY_SP};
 
 use rustc::mir::interpret::{
     EvalResult, EvalError, EvalErrorKind, GlobalId,
-    Scalar, Allocation, ConstValue,
+    Scalar, Allocation, AllocId, ConstValue,
 };
 use interpret::{self,
     Place, PlaceTy, MemPlace, OpTy, Operand, Value,
@@ -265,6 +268,67 @@ impl<'a, 'mir, 'tcx> CompileTimeInterpreter<'a, 'mir, 'tcx> {
     }
 }
 
+impl<K: Hash + Eq, V> interpret::AllocMap<K, V> for FxHashMap<K, V> {
+    #[inline(always)]
+    fn contains_key<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> bool
+        where K: Borrow<Q>
+    {
+        FxHashMap::contains_key(self, k)
+    }
+
+    #[inline(always)]
+    fn insert(&mut self, k: K, v: V) -> Option<V>
+    {
+        FxHashMap::insert(self, k, v)
+    }
+
+    #[inline(always)]
+    fn remove<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> Option<V>
+        where K: Borrow<Q>
+    {
+        FxHashMap::remove(self, k)
+    }
+
+    #[inline(always)]
+    fn filter_map_collect<T>(&self, mut f: impl FnMut(&K, &V) -> Option<T>) -> Vec<T> {
+        self.iter()
+            .filter_map(move |(k, v)| f(k, &*v))
+            .collect()
+    }
+
+    #[inline(always)]
+    fn get_or<E>(
+        &self,
+        k: K,
+        vacant: impl FnOnce() -> Result<V, E>
+    ) -> Result<&V, E>
+    {
+        match self.get(&k) {
+            Some(v) => Ok(v),
+            None => {
+                vacant()?;
+                bug!("The CTFE machine shouldn't ever need to extend the alloc_map when reading")
+            }
+        }
+    }
+
+    #[inline(always)]
+    fn get_mut_or<E>(
+        &mut self,
+        k: K,
+        vacant: impl FnOnce() -> Result<V, E>
+    ) -> Result<&mut V, E>
+    {
+        match self.entry(k) {
+            Entry::Occupied(e) => Ok(e.into_mut()),
+            Entry::Vacant(e) => {
+                let v = vacant()?;
+                Ok(e.insert(v))
+            }
+        }
+    }
+}
+
 type CompileTimeEvalContext<'a, 'mir, 'tcx> =
     EvalContext<'a, 'mir, 'tcx, CompileTimeInterpreter<'a, 'mir, 'tcx>>;
 
@@ -275,6 +339,8 @@ impl<'a, 'mir, 'tcx> interpret::Machine<'a, 'mir, 'tcx>
     type MemoryKinds = !;
     type PointerTag = ();
 
+    type MemoryMap = FxHashMap<AllocId, (MemoryKind<!>, Allocation<()>)>;
+
     const STATIC_KIND: Option<!> = None; // no copying of statics allowed
     const ENFORCE_VALIDITY: bool = false; // for now, we don't
 
diff --git a/src/librustc_mir/interpret/machine.rs b/src/librustc_mir/interpret/machine.rs
@@ -12,27 +12,68 @@
 //! This separation exists to ensure that no fancy miri features like
 //! interpreting common C functions leak into CTFE.
 
-use std::borrow::Cow;
+use std::borrow::{Borrow, Cow};
 use std::hash::Hash;
 
 use rustc::hir::def_id::DefId;
-use rustc::mir::interpret::{Allocation, EvalResult, Scalar};
+use rustc::mir::interpret::{Allocation, AllocId, EvalResult, Scalar};
 use rustc::mir;
 use rustc::ty::{self, layout::TyLayout, query::TyCtxtAt};
 
-use super::{EvalContext, PlaceTy, OpTy};
+use super::{EvalContext, PlaceTy, OpTy, MemoryKind};
+
+/// The functionality needed by memory to manage its allocations
+pub trait AllocMap<K: Hash + Eq, V> {
+    /// Test if the map contains the given key.
+    /// Deliberately takes `&mut` because that is sufficient, and some implementations
+    /// can be more efficient then (using `RefCell::get_mut`).
+    fn contains_key<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> bool
+        where K: Borrow<Q>;
+
+    /// Insert new entry into the map.
+    fn insert(&mut self, k: K, v: V) -> Option<V>;
+
+    /// Remove entry from the map.
+    fn remove<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> Option<V>
+        where K: Borrow<Q>;
+
+    /// Return data based the keys and values in the map.
+    fn filter_map_collect<T>(&self, f: impl FnMut(&K, &V) -> Option<T>) -> Vec<T>;
+
+    /// Return a reference to entry `k`.  If no such entry exists, call
+    /// `vacant` and either forward its error, or add its result to the map
+    /// and return a reference to *that*.
+    fn get_or<E>(
+        &self,
+        k: K,
+        vacant: impl FnOnce() -> Result<V, E>
+    ) -> Result<&V, E>;
+
+    /// Return a mutable reference to entry `k`.  If no such entry exists, call
+    /// `vacant` and either forward its error, or add its result to the map
+    /// and return a reference to *that*.
+    fn get_mut_or<E>(
+        &mut self,
+        k: K,
+        vacant: impl FnOnce() -> Result<V, E>
+    ) -> Result<&mut V, E>;
+}
 
 /// Methods of this trait signifies a point where CTFE evaluation would fail
 /// and some use case dependent behaviour can instead be applied.
-/// FIXME: We should be able to get rid of the 'a here if we can get rid of the 'a in
-/// `snapshot::EvalSnapshot`.
 pub trait Machine<'a, 'mir, 'tcx>: Sized {
     /// Additional data that can be accessed via the Memory
     type MemoryData;
 
     /// Additional memory kinds a machine wishes to distinguish from the builtin ones
     type MemoryKinds: ::std::fmt::Debug + Copy + Eq;
 
+    /// Memory's allocation map
+    type MemoryMap:
+        AllocMap<AllocId, (MemoryKind<Self::MemoryKinds>, Allocation<Self::PointerTag>)> +
+        Default +
+        Clone;
+
     /// Tag tracked alongside every pointer.  This is inert for now, in preparation for
     /// a future implementation of "Stacked Borrows"
     /// <https://www.ralfj.de/blog/2018/08/07/stacked-borrows.html>.
diff --git a/src/librustc_mir/interpret/memory.rs b/src/librustc_mir/interpret/memory.rs
@@ -16,22 +16,23 @@
 //! integer.  It is crucial that these operations call `check_align` *before*
 //! short-circuiting the empty case!
 
-use std::collections::hash_map::Entry;
 use std::collections::VecDeque;
 use std::ptr;
 use std::borrow::Cow;
 
 use rustc::ty::{self, Instance, ParamEnv, query::TyCtxtAt};
 use rustc::ty::layout::{self, Align, TargetDataLayout, Size, HasDataLayout};
-use rustc::mir::interpret::{Pointer, AllocId, Allocation, ConstValue, GlobalId,
-                            EvalResult, Scalar, EvalErrorKind, AllocType, PointerArithmetic,
-                            truncate};
+use rustc::mir::interpret::{
+    Pointer, AllocId, Allocation, ConstValue, GlobalId,
+    EvalResult, Scalar, EvalErrorKind, AllocType, PointerArithmetic,
+    truncate
+};
 pub use rustc::mir::interpret::{write_target_uint, read_target_uint};
 use rustc_data_structures::fx::{FxHashSet, FxHashMap};
 
 use syntax::ast::Mutability;
 
-use super::{Machine, MonoHashMap, ScalarMaybeUndef};
+use super::{Machine, AllocMap, ScalarMaybeUndef};
 
 #[derive(Debug, PartialEq, Eq, Copy, Clone, Hash)]
 pub enum MemoryKind<T> {
@@ -52,13 +53,13 @@ pub struct Memory<'a, 'mir, 'tcx: 'a + 'mir, M: Machine<'a, 'mir, 'tcx>> {
     /// Allocations local to this instance of the miri engine.  The kind
     /// helps ensure that the same mechanism is used for allocation and
     /// deallocation.  When an allocation is not found here, it is a
-    /// static and looked up in the `tcx` for read access.  If this machine
-    /// does pointer provenance tracking, the type of alloctions in `tcx`
-    /// and here do not match, so we have a `MonoHashMap` to be able to
-    /// put the "mapped" allocation into `alloc_map` even on a read access.
+    /// static and looked up in the `tcx` for read access.  Some machines may
+    /// have to mutate this map even on a read-only access to a static (because
+    /// they do pointer provenance tracking and the allocations in `tcx` have
+    /// the wrong type), so we let the machine override this type.
     /// Either way, if the machine allows writing to a static, doing so will
     /// create a copy of the static allocation here.
-    alloc_map: MonoHashMap<AllocId, (MemoryKind<M::MemoryKinds>, Allocation<M::PointerTag>)>,
+    alloc_map: M::MemoryMap,
 
     /// To be able to compare pointers with NULL, and to check alignment for accesses
     /// to ZSTs (where pointers may dangle), we keep track of the size even for allocations
@@ -106,7 +107,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
     pub fn new(tcx: TyCtxtAt<'a, 'tcx, 'tcx>, data: M::MemoryData) -> Self {
         Memory {
             data,
-            alloc_map: MonoHashMap::default(),
+            alloc_map: Default::default(),
             dead_alloc_map: FxHashMap::default(),
             tcx,
         }
@@ -419,30 +420,32 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
         &mut self,
         id: AllocId,
     ) -> EvalResult<'tcx, &mut Allocation<M::PointerTag>> {
-        Ok(match self.alloc_map.entry(id) {
-            // Normal alloc?
-            Entry::Occupied(alloc) => {
-                let alloc = &mut alloc.into_mut().1;
-                if alloc.mutability == Mutability::Immutable {
-                    return err!(ModifiedConstantMemory);
-                }
-                alloc
+        let tcx = self.tcx;
+        let a = self.alloc_map.get_mut_or(id, || {
+            // Need to make a copy, even if `get_static_alloc` is able
+            // to give us a cheap reference.
+            let alloc = Self::get_static_alloc(tcx, id)?;
+            if alloc.mutability == Mutability::Immutable {
+                return err!(ModifiedConstantMemory);
             }
-            // Static.
-            Entry::Vacant(entry) => {
-                // Need to make a copy, even if `get_static_alloc` is able
-                // to give us a cheap reference.
-                let alloc = Self::get_static_alloc(self.tcx, id)?;
-                if alloc.mutability == Mutability::Immutable {
+            let kind = M::STATIC_KIND.expect(
+                "I got an owned allocation that I have to copy but the machine does \
+                    not expect that to happen"
+            );
+            Ok((MemoryKind::Machine(kind), alloc.into_owned()))
+        });
+        // Unpack the error type manually because type inference doesn't
+        // work otherwise (and we cannot help it because `impl Trait`)
+        match a {
+            Err(e) => Err(e),
+            Ok(a) => {
+                let a = &mut a.1;
+                if a.mutability == Mutability::Immutable {
                     return err!(ModifiedConstantMemory);
                 }
-                let kind = M::STATIC_KIND.expect(
-                    "I got an owned allocation that I have to copy but the machine does \
-                        not expect that to happen"
-                );
-                &mut entry.insert(Box::new((MemoryKind::Machine(kind), alloc.into_owned()))).1
+                Ok(a)
             }
-        })
+        }
     }
 
     pub fn get_fn(&self, ptr: Pointer<M::PointerTag>) -> EvalResult<'tcx, Instance<'tcx>> {
@@ -534,8 +537,8 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
             let msg = format!("Alloc {:<5} ", format!("{}:", id));
 
             // normal alloc?
-            match self.alloc_map.get(&id) {
-                Some((kind, alloc)) => {
+            match self.alloc_map.get_or(id, || Err(())) {
+                Ok((kind, alloc)) => {
                     let extra = match kind {
                         MemoryKind::Stack => " (stack)".to_owned(),
                         MemoryKind::Vtable => " (vtable)".to_owned(),
@@ -546,7 +549,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
                         msg, alloc, extra
                     );
                 },
-                None => {
+                Err(()) => {
                     // static alloc?
                     match self.tcx.alloc_map.lock().get(id) {
                         Some(AllocType::Memory(alloc)) => {
@@ -664,7 +667,11 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> Memory<'a, 'mir, 'tcx, M> {
 }
 
 /// Interning (for CTFE)
-impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx, PointerTag=()>> Memory<'a, 'mir, 'tcx, M> {
+impl<'a, 'mir, 'tcx, M> Memory<'a, 'mir, 'tcx, M>
+where
+    M: Machine<'a, 'mir, 'tcx, PointerTag=()>,
+    M::MemoryMap: AllocMap<AllocId, (MemoryKind<M::MemoryKinds>, Allocation<()>)>,
+{
     /// mark an allocation as static and initialized, either mutable or not
     pub fn intern_static(
         &mut self,
diff --git a/src/librustc_mir/interpret/mod.rs b/src/librustc_mir/interpret/mod.rs
@@ -23,7 +23,6 @@ mod terminator;
 mod traits;
 mod validity;
 mod intrinsics;
-mod mono_hash_map;
 
 pub use self::eval_context::{
     EvalContext, Frame, StackPopCleanup, LocalValue,
@@ -33,10 +32,8 @@ pub use self::place::{Place, PlaceTy, MemPlace, MPlaceTy};
 
 pub use self::memory::{Memory, MemoryKind};
 
-pub use self::machine::Machine;
+pub use self::machine::{Machine, AllocMap};
 
 pub use self::operand::{ScalarMaybeUndef, Value, ValTy, Operand, OpTy};
 
 pub use self::validity::RefTracking;
-
-pub use self::mono_hash_map::MonoHashMap;
diff --git a/src/librustc_mir/interpret/mono_hash_map.rs b/src/librustc_mir/interpret/mono_hash_map.rs
diff --git a/src/librustc_mir/interpret/place.rs b/src/librustc_mir/interpret/place.rs
