Introduce AllocType which indicates what AllocIds point to

src/librustc/ich/impls_ty.rs

@@ -420,17 +420,6 @@ impl_stable_hash_for!(struct mir::interpret::MemoryPointer {
     offset
 });
 
-enum AllocDiscriminant {
-    Alloc,
-    Static,
-    Function,
-}
-impl_stable_hash_for!(enum self::AllocDiscriminant {
-    Alloc,
-    Static,
-    Function
-});
-
 impl<'a> HashStable<StableHashingContext<'a>> for mir::interpret::AllocId {
     fn hash_stable<W: StableHasherResult>(
         &self,
@@ -440,30 +429,29 @@ impl<'a> HashStable<StableHashingContext<'a>> for mir::interpret::AllocId {
         ty::tls::with_opt(|tcx| {
             trace!("hashing {:?}", *self);
             let tcx = tcx.expect("can't hash AllocIds during hir lowering");
-            if let Some(def_id) = tcx.interpret_interner.get_static(*self) {
-                AllocDiscriminant::Static.hash_stable(hcx, hasher);
-                trace!("hashing {:?} as static {:?}", *self, def_id);
-                def_id.hash_stable(hcx, hasher);
-            } else if let Some(alloc) = tcx.interpret_interner.get_alloc(*self) {
-                AllocDiscriminant::Alloc.hash_stable(hcx, hasher);
-                if hcx.alloc_id_recursion_tracker.insert(*self) {
-                    trace!("hashing {:?} as alloc {:#?}", *self, alloc);
-                    alloc.hash_stable(hcx, hasher);
-                    assert!(hcx.alloc_id_recursion_tracker.remove(self));
-                } else {
-                    trace!("skipping hashing of {:?} due to recursion", *self);
-                }
-            } else if let Some(inst) = tcx.interpret_interner.get_fn(*self) {
-                trace!("hashing {:?} as fn {:#?}", *self, inst);
-                AllocDiscriminant::Function.hash_stable(hcx, hasher);
-                inst.hash_stable(hcx, hasher);
-            } else {
-                bug!("no allocation for {}", self);
-            }
+            let alloc_kind = tcx.alloc_map.lock().get(*self).expect("no value for AllocId");
+            alloc_kind.hash_stable(hcx, hasher);
         });
     }
 }
 
+impl<'a, 'gcx, M: HashStable<StableHashingContext<'a>>> HashStable<StableHashingContext<'a>>
+for mir::interpret::AllocType<'gcx, M> {
+    fn hash_stable<W: StableHasherResult>(&self,
+                                          hcx: &mut StableHashingContext<'a>,
+                                          hasher: &mut StableHasher<W>) {
+        use mir::interpret::AllocType::*;
+
+        mem::discriminant(self).hash_stable(hcx, hasher);
+
+        match *self {
+            Function(instance) => instance.hash_stable(hcx, hasher),
+            Static(def_id) => def_id.hash_stable(hcx, hasher),
+            Memory(ref mem) => mem.hash_stable(hcx, hasher),
+        }
+    }
+}
+
 impl<'a> HashStable<StableHashingContext<'a>> for mir::interpret::Allocation {
     fn hash_stable<W: StableHasherResult>(
         &self,

src/librustc/mir/interpret/mod.rs

@@ -16,13 +16,15 @@ use std::collections::BTreeMap;
 use std::fmt;
 use mir;
 use hir::def_id::DefId;
-use ty::{self, TyCtxt};
+use ty::{self, TyCtxt, Instance};
 use ty::layout::{self, Align, HasDataLayout, Size};
 use middle::region;
 use std::iter;
 use std::io;
+use std::hash::Hash;
 use syntax::ast::Mutability;
 use rustc_serialize::{Encoder, Decoder, Decodable, Encodable};
+use rustc_data_structures::fx::FxHashMap;
 use byteorder::{WriteBytesExt, ReadBytesExt, LittleEndian, BigEndian};
 
 #[derive(Clone, Debug, PartialEq, RustcEncodable, RustcDecodable)]
@@ -150,7 +152,7 @@ impl<'tcx> MemoryPointer {
 }
 
 
-#[derive(Copy, Clone, Default, Eq, Hash, Ord, PartialEq, PartialOrd, Debug)]
+#[derive(Copy, Clone, Eq, Hash, Ord, PartialEq, PartialOrd, Debug)]
 pub struct AllocId(pub u64);
 
 impl ::rustc_serialize::UseSpecializedEncodable for AllocId {}
@@ -171,20 +173,25 @@ pub fn specialized_encode_alloc_id<
     tcx: TyCtxt<'a, 'tcx, 'tcx>,
     alloc_id: AllocId,
 ) -> Result<(), E::Error> {
-    if let Some(alloc) = tcx.interpret_interner.get_alloc(alloc_id) {
-        trace!("encoding {:?} with {:#?}", alloc_id, alloc);
-        AllocKind::Alloc.encode(encoder)?;
-        alloc.encode(encoder)?;
-    } else if let Some(fn_instance) = tcx.interpret_interner.get_fn(alloc_id) {
-        trace!("encoding {:?} with {:#?}", alloc_id, fn_instance);
-        AllocKind::Fn.encode(encoder)?;
-        fn_instance.encode(encoder)?;
-    } else if let Some(did) = tcx.interpret_interner.get_static(alloc_id) {
-        // referring to statics doesn't need to know about their allocations, just about its DefId
-        AllocKind::Static.encode(encoder)?;
-        did.encode(encoder)?;
-    } else {
-        bug!("alloc id without corresponding allocation: {}", alloc_id);
+    let alloc_type: AllocType<'tcx, &'tcx Allocation> =
+        tcx.alloc_map.lock().get(alloc_id).expect("no value for AllocId");
+    match alloc_type {
+        AllocType::Memory(alloc) => {
+            trace!("encoding {:?} with {:#?}", alloc_id, alloc);
+            AllocKind::Alloc.encode(encoder)?;
+            alloc.encode(encoder)?;
+        }
+        AllocType::Function(fn_instance) => {
+            trace!("encoding {:?} with {:#?}", alloc_id, fn_instance);
+            AllocKind::Fn.encode(encoder)?;
+            fn_instance.encode(encoder)?;
+        }
+        AllocType::Static(did) => {
+            // referring to statics doesn't need to know about their allocations,
+            // just about its DefId
+            AllocKind::Static.encode(encoder)?;
+            did.encode(encoder)?;
+        }
     }
     Ok(())
 }
@@ -200,31 +207,30 @@ pub fn specialized_decode_alloc_id<
 ) -> Result<AllocId, D::Error> {
     match AllocKind::decode(decoder)? {
         AllocKind::Alloc => {
-            let alloc_id = tcx.interpret_interner.reserve();
+            let alloc_id = tcx.alloc_map.lock().reserve();
             trace!("creating alloc id {:?}", alloc_id);
             // insert early to allow recursive allocs
             cache(decoder, alloc_id);
 
-            let allocation = Allocation::decode(decoder)?;
+            let allocation = <&'tcx Allocation as Decodable>::decode(decoder)?;
             trace!("decoded alloc {:?} {:#?}", alloc_id, allocation);
-            let allocation = tcx.intern_const_alloc(allocation);
-            tcx.interpret_interner.intern_at_reserved(alloc_id, allocation);
+            tcx.alloc_map.lock().set_id_memory(alloc_id, allocation);
 
             Ok(alloc_id)
         },
         AllocKind::Fn => {
             trace!("creating fn alloc id");
             let instance = ty::Instance::decode(decoder)?;
             trace!("decoded fn alloc instance: {:?}", instance);
-            let id = tcx.interpret_interner.create_fn_alloc(instance);
+            let id = tcx.alloc_map.lock().create_fn_alloc(instance);
             trace!("created fn alloc id: {:?}", id);
             cache(decoder, id);
             Ok(id)
         },
         AllocKind::Static => {
             trace!("creating extern static alloc id at");
             let did = DefId::decode(decoder)?;
-            let alloc_id = tcx.interpret_interner.cache_static(did);
+            let alloc_id = tcx.alloc_map.lock().intern_static(did);
             cache(decoder, alloc_id);
             Ok(alloc_id)
         },
@@ -237,6 +243,97 @@ impl fmt::Display for AllocId {
     }
 }
 
+#[derive(Debug, Clone, Eq, PartialEq, Hash, RustcDecodable, RustcEncodable)]
+pub enum AllocType<'tcx, M> {
+    /// The alloc id is used as a function pointer
+    Function(Instance<'tcx>),
+    /// The alloc id points to a static variable
+    Static(DefId),
+    /// The alloc id points to memory
+    Memory(M)
+}
+
+pub struct AllocMap<'tcx, M> {
+    /// Lets you know what an AllocId refers to
+    id_to_type: FxHashMap<AllocId, AllocType<'tcx, M>>,
+
+    /// Used to ensure that functions and statics only get one associated AllocId
+    type_interner: FxHashMap<AllocType<'tcx, M>, AllocId>,
+
+    /// The AllocId to assign to the next requested id.
+    /// Always incremented, never gets smaller.
+    next_id: AllocId,
+}
+
+impl<'tcx, M: fmt::Debug + Eq + Hash + Clone> AllocMap<'tcx, M> {
+    pub fn new() -> Self {
+        AllocMap {
+            id_to_type: FxHashMap(),
+            type_interner: FxHashMap(),
+            next_id: AllocId(0),
+        }
+    }
+
+    /// obtains a new allocation ID that can be referenced but does not
+    /// yet have an allocation backing it.
+    pub fn reserve(
+        &mut self,
+    ) -> AllocId {
+        let next = self.next_id;
+        self.next_id.0 = self.next_id.0
+            .checked_add(1)
+            .expect("You overflowed a u64 by incrementing by 1... \
+                     You've just earned yourself a free drink if we ever meet. \
+                     Seriously, how did you do that?!");
+        next
+    }
+
+    fn intern(&mut self, alloc_type: AllocType<'tcx, M>) -> AllocId {
+        if let Some(&alloc_id) = self.type_interner.get(&alloc_type) {
+            return alloc_id;
+        }
+        let id = self.reserve();
+        debug!("creating alloc_type {:?} with id {}", alloc_type, id);
+        self.id_to_type.insert(id, alloc_type.clone());
+        self.type_interner.insert(alloc_type, id);
+        id
+    }
+
+    // FIXME: Check if functions have identity. If not, we should not intern these,
+    // but instead create a new id per use.
+    // Alternatively we could just make comparing function pointers an error.
+    pub fn create_fn_alloc(&mut self, instance: Instance<'tcx>) -> AllocId {
+        self.intern(AllocType::Function(instance))
+    }
+
+    pub fn get(&self, id: AllocId) -> Option<AllocType<'tcx, M>> {
+        self.id_to_type.get(&id).cloned()
+    }
+
+    pub fn unwrap_memory(&self, id: AllocId) -> M {
+        match self.get(id) {
+            Some(AllocType::Memory(mem)) => mem,
+            _ => bug!("expected allocation id {} to point to memory", id),
+        }
+    }
+
+    pub fn intern_static(&mut self, static_id: DefId) -> AllocId {
+        self.intern(AllocType::Static(static_id))
+    }
+
+    pub fn allocate(&mut self, mem: M) -> AllocId {
+        let id = self.reserve();
+        self.set_id_memory(id, mem);
+        id
+    }
+
+    pub fn set_id_memory(&mut self, id: AllocId, mem: M) {
+        if let Some(old) = self.id_to_type.insert(id, AllocType::Memory(mem)) {
+            bug!("tried to set allocation id {}, but it was already existing as {:#?}", id, old);
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Hash, RustcEncodable, RustcDecodable)]
 pub struct Allocation {
     /// The actual bytes of the allocation.

src/librustc/mir/mod.rs

@@ -1908,17 +1908,15 @@ pub fn print_miri_value<W: Write>(value: Value, ty: Ty, f: &mut W) -> fmt::Resul
         (Value::ByValPair(PrimVal::Ptr(ptr), PrimVal::Bytes(len)),
          &TyRef(_, &ty::TyS { sty: TyStr, .. }, _)) => {
             ty::tls::with(|tcx| {
-                let alloc = tcx
-                    .interpret_interner
-                    .get_alloc(ptr.alloc_id);
-                if let Some(alloc) = alloc {
-                    assert_eq!(len as usize as u128, len);
-                    let slice = &alloc.bytes[(ptr.offset.bytes() as usize)..][..(len as usize)];
-                    let s = ::std::str::from_utf8(slice)
-                        .expect("non utf8 str from miri");
-                    write!(f, "{:?}", s)
-                } else {
-                    write!(f, "pointer to erroneous constant {:?}, {:?}", ptr, len)
+                match tcx.alloc_map.lock().get(ptr.alloc_id) {
+                    Some(interpret::AllocType::Memory(alloc)) => {
+                        assert_eq!(len as usize as u128, len);
+                        let slice = &alloc.bytes[(ptr.offset.bytes() as usize)..][..(len as usize)];
+                        let s = ::std::str::from_utf8(slice)
+                            .expect("non utf8 str from miri");
+                        write!(f, "{:?}", s)
+                    }
+                    _ => write!(f, "pointer to erroneous constant {:?}, {:?}", ptr, len),
                 }
             })
         },