Auto merge of #51060 - michaelwoerister:thread-safe-consts, r=Zoxc

Make const decoding thread-safe.

This is an alternative to https://github.com/rust-lang/rust/pull/50957. It's a proof of concept (e.g. it doesn't adapt metadata decoding, just the incr. comp. cache) but I think it turned out nice. It's rather simple and does not require passing around a bunch of weird closures, like we currently do.

If you (@Zoxc & @oli-obk) think this approach is good then I'm happy to finish and clean this up.

Note: The current version just spins when it encounters an in-progress decoding. I don't have a strong preference for this approach. Decoding concurrently is equally fine by me (or maybe even better because it doesn't require poisoning).

r? @Zoxc

8 files changed, 449 insertions, 115 deletions

diff --git a/src/librustc/lib.rs b/src/librustc/lib.rs
index da6340b5f61..10e8905054d 100644
--- a/src/librustc/lib.rs
+++ b/src/librustc/lib.rs
@@ -68,6 +68,7 @@
 #![feature(trace_macros)]
 #![feature(trusted_len)]
 #![feature(catch_expr)]
+#![feature(integer_atomics)]
 #![feature(test)]
 #![feature(in_band_lifetimes)]
 #![feature(macro_at_most_once_rep)]
diff --git a/src/librustc/mir/interpret/mod.rs b/src/librustc/mir/interpret/mod.rs
index b41652469ae..6bd5814799a 100644
--- a/src/librustc/mir/interpret/mod.rs
+++ b/src/librustc/mir/interpret/mod.rs
@@ -23,10 +23,15 @@ use std::io;
 use std::ops::{Deref, DerefMut};
 use std::hash::Hash;
 use syntax::ast::Mutability;
-use rustc_serialize::{Encoder, Decoder, Decodable, Encodable};
+use rustc_serialize::{Encoder, Decodable, Encodable};
 use rustc_data_structures::sorted_map::SortedMap;
 use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::sync::{Lock as Mutex, HashMapExt};
+use rustc_data_structures::tiny_list::TinyList;
 use byteorder::{WriteBytesExt, ReadBytesExt, LittleEndian, BigEndian};
+use ty::codec::TyDecoder;
+use std::sync::atomic::{AtomicU32, Ordering};
+use std::num::NonZeroU32;
 
 #[derive(Clone, Debug, PartialEq, RustcEncodable, RustcDecodable)]
 pub enum Lock {
@@ -204,44 +209,163 @@ pub fn specialized_encode_alloc_id<
     Ok(())
 }
 
-pub fn specialized_decode_alloc_id<
-    'a, 'tcx,
-    D: Decoder,
-    CACHE: FnOnce(&mut D, AllocId),
->(
-    decoder: &mut D,
-    tcx: TyCtxt<'a, 'tcx, 'tcx>,
-    cache: CACHE,
-) -> Result<AllocId, D::Error> {
-    match AllocKind::decode(decoder)? {
-        AllocKind::Alloc => {
-            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 = <&'tcx Allocation as Decodable>::decode(decoder)?;
-            trace!("decoded alloc {:?} {:#?}", 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.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.alloc_map.lock().intern_static(did);
-            cache(decoder, alloc_id);
-            Ok(alloc_id)
-        },
+// Used to avoid infinite recursion when decoding cyclic allocations.
+type DecodingSessionId = NonZeroU32;
+
+#[derive(Clone)]
+enum State {
+    Empty,
+    InProgressNonAlloc(TinyList<DecodingSessionId>),
+    InProgress(TinyList<DecodingSessionId>, AllocId),
+    Done(AllocId),
+}
+
+pub struct AllocDecodingState {
+    // For each AllocId we keep track of which decoding state it's currently in.
+    decoding_state: Vec<Mutex<State>>,
+    // The offsets of each allocation in the data stream.
+    data_offsets: Vec<u32>,
+}
+
+impl AllocDecodingState {
+
+    pub fn new_decoding_session(&self) -> AllocDecodingSession {
+        static DECODER_SESSION_ID: AtomicU32 = AtomicU32::new(0);
+        let counter = DECODER_SESSION_ID.fetch_add(1, Ordering::SeqCst);
+
+        // Make sure this is never zero
+        let session_id = DecodingSessionId::new((counter & 0x7FFFFFFF) + 1).unwrap();
+
+        AllocDecodingSession {
+            state: self,
+            session_id,
+        }
+    }
+
+    pub fn new(data_offsets: Vec<u32>) -> AllocDecodingState {
+        let decoding_state: Vec<_> = ::std::iter::repeat(Mutex::new(State::Empty))
+            .take(data_offsets.len())
+            .collect();
+
+        AllocDecodingState {
+            decoding_state: decoding_state,
+            data_offsets,
+        }
+    }
+}
+
+#[derive(Copy, Clone)]
+pub struct AllocDecodingSession<'s> {
+    state: &'s AllocDecodingState,
+    session_id: DecodingSessionId,
+}
+
+impl<'s> AllocDecodingSession<'s> {
+
+    // Decodes an AllocId in a thread-safe way.
+    pub fn decode_alloc_id<'a, 'tcx, D>(&self,
+                                        decoder: &mut D)
+                                        -> Result<AllocId, D::Error>
+        where D: TyDecoder<'a, 'tcx>,
+              'tcx: 'a,
+    {
+        // Read the index of the allocation
+        let idx = decoder.read_u32()? as usize;
+        let pos = self.state.data_offsets[idx] as usize;
+
+        // Decode the AllocKind now so that we know if we have to reserve an
+        // AllocId.
+        let (alloc_kind, pos) = decoder.with_position(pos, |decoder| {
+            let alloc_kind = AllocKind::decode(decoder)?;
+            Ok((alloc_kind, decoder.position()))
+        })?;
+
+        // Check the decoding state, see if it's already decoded or if we should
+        // decode it here.
+        let alloc_id = {
+            let mut entry = self.state.decoding_state[idx].lock();
+
+            match *entry {
+                State::Done(alloc_id) => {
+                    return Ok(alloc_id);
+                }
+                ref mut entry @ State::Empty => {
+                    // We are allowed to decode
+                    match alloc_kind {
+                        AllocKind::Alloc => {
+                            // If this is an allocation, we need to reserve an
+                            // AllocId so we can decode cyclic graphs.
+                            let alloc_id = decoder.tcx().alloc_map.lock().reserve();
+                            *entry = State::InProgress(
+                                TinyList::new_single(self.session_id),
+                                alloc_id);
+                            Some(alloc_id)
+                        },
+                        AllocKind::Fn | AllocKind::Static => {
+                            // Fns and statics cannot be cyclic and their AllocId
+                            // is determined later by interning
+                            *entry = State::InProgressNonAlloc(
+                                TinyList::new_single(self.session_id));
+                            None
+                        }
+                    }
+                }
+                State::InProgressNonAlloc(ref mut sessions) => {
+                    if sessions.contains(&self.session_id) {
+                        bug!("This should be unreachable")
+                    } else {
+                        // Start decoding concurrently
+                        sessions.insert(self.session_id);
+                        None
+                    }
+                }
+                State::InProgress(ref mut sessions, alloc_id) => {
+                    if sessions.contains(&self.session_id) {
+                        // Don't recurse.
+                        return Ok(alloc_id)
+                    } else {
+                        // Start decoding concurrently
+                        sessions.insert(self.session_id);
+                        Some(alloc_id)
+                    }
+                }
+            }
+        };
+
+        // Now decode the actual data
+        let alloc_id = decoder.with_position(pos, |decoder| {
+            match alloc_kind {
+                AllocKind::Alloc => {
+                    let allocation = <&'tcx Allocation as Decodable>::decode(decoder)?;
+                    // We already have a reserved AllocId.
+                    let alloc_id = alloc_id.unwrap();
+                    trace!("decoded alloc {:?} {:#?}", alloc_id, allocation);
+                    decoder.tcx().alloc_map.lock().set_id_same_memory(alloc_id, allocation);
+                    Ok(alloc_id)
+                },
+                AllocKind::Fn => {
+                    assert!(alloc_id.is_none());
+                    trace!("creating fn alloc id");
+                    let instance = ty::Instance::decode(decoder)?;
+                    trace!("decoded fn alloc instance: {:?}", instance);
+                    let alloc_id = decoder.tcx().alloc_map.lock().create_fn_alloc(instance);
+                    Ok(alloc_id)
+                },
+                AllocKind::Static => {
+                    assert!(alloc_id.is_none());
+                    trace!("creating extern static alloc id at");
+                    let did = DefId::decode(decoder)?;
+                    let alloc_id = decoder.tcx().alloc_map.lock().intern_static(did);
+                    Ok(alloc_id)
+                }
+            }
+        })?;
+
+        self.state.decoding_state[idx].with_lock(|entry| {
+            *entry = State::Done(alloc_id);
+        });
+
+        Ok(alloc_id)
     }
 }
 
@@ -340,6 +464,10 @@ impl<'tcx, M: fmt::Debug + Eq + Hash + Clone> AllocMap<'tcx, M> {
             bug!("tried to set allocation id {}, but it was already existing as {:#?}", id, old);
         }
     }
+
+    pub fn set_id_same_memory(&mut self, id: AllocId, mem: M) {
+       self.id_to_type.insert_same(id, AllocType::Memory(mem));
+    }
 }
 
 #[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Ord, Hash, RustcEncodable, RustcDecodable)]
diff --git a/src/librustc/ty/maps/on_disk_cache.rs b/src/librustc/ty/maps/on_disk_cache.rs
index 9ca90a06c4e..cd317ff6cdb 100644
--- a/src/librustc/ty/maps/on_disk_cache.rs
+++ b/src/librustc/ty/maps/on_disk_cache.rs
@@ -16,6 +16,7 @@ use hir::def_id::{CrateNum, DefIndex, DefId, LocalDefId,
 use hir::map::definitions::DefPathHash;
 use ich::{CachingCodemapView, Fingerprint};
 use mir::{self, interpret};
+use mir::interpret::{AllocDecodingSession, AllocDecodingState};
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sync::{Lrc, Lock, HashMapExt, Once};
 use rustc_data_structures::indexed_vec::{IndexVec, Idx};
@@ -23,7 +24,6 @@ use rustc_serialize::{Decodable, Decoder, Encodable, Encoder, opaque,
                       SpecializedDecoder, SpecializedEncoder,
                       UseSpecializedDecodable, UseSpecializedEncodable};
 use session::{CrateDisambiguator, Session};
-use std::cell::RefCell;
 use std::mem;
 use syntax::ast::NodeId;
 use syntax::codemap::{CodeMap, StableFilemapId};
@@ -77,11 +77,7 @@ pub struct OnDiskCache<'sess> {
     // `serialized_data`.
     prev_diagnostics_index: FxHashMap<SerializedDepNodeIndex, AbsoluteBytePos>,
 
-    // Alloc indices to memory location map
-    prev_interpret_alloc_index: Vec<AbsoluteBytePos>,
-
-    /// Deserialization: A cache to ensure we don't read allocations twice
-    interpret_alloc_cache: RefCell<FxHashMap<usize, interpret::AllocId>>,
+    alloc_decoding_state: AllocDecodingState,
 }
 
 // This type is used only for (de-)serialization.
@@ -92,7 +88,7 @@ struct Footer {
     query_result_index: EncodedQueryResultIndex,
     diagnostics_index: EncodedQueryResultIndex,
     // the location of all allocations
-    interpret_alloc_index: Vec<AbsoluteBytePos>,
+    interpret_alloc_index: Vec<u32>,
 }
 
 type EncodedQueryResultIndex = Vec<(SerializedDepNodeIndex, AbsoluteBytePos)>;
@@ -149,8 +145,7 @@ impl<'sess> OnDiskCache<'sess> {
             query_result_index: footer.query_result_index.into_iter().collect(),
             prev_diagnostics_index: footer.diagnostics_index.into_iter().collect(),
             synthetic_expansion_infos: Lock::new(FxHashMap()),
-            prev_interpret_alloc_index: footer.interpret_alloc_index,
-            interpret_alloc_cache: RefCell::new(FxHashMap::default()),
+            alloc_decoding_state: AllocDecodingState::new(footer.interpret_alloc_index),
         }
     }
 
@@ -166,8 +161,7 @@ impl<'sess> OnDiskCache<'sess> {
             query_result_index: FxHashMap(),
             prev_diagnostics_index: FxHashMap(),
             synthetic_expansion_infos: Lock::new(FxHashMap()),
-            prev_interpret_alloc_index: Vec::new(),
-            interpret_alloc_cache: RefCell::new(FxHashMap::default()),
+            alloc_decoding_state: AllocDecodingState::new(Vec::new()),
         }
     }
 
@@ -291,7 +285,7 @@ impl<'sess> OnDiskCache<'sess> {
                     }
                     for idx in n..new_n {
                         let id = encoder.interpret_allocs_inverse[idx];
-                        let pos = AbsoluteBytePos::new(encoder.position());
+                        let pos = encoder.position() as u32;
                         interpret_alloc_index.push(pos);
                         interpret::specialized_encode_alloc_id(
                             &mut encoder,
@@ -424,8 +418,7 @@ impl<'sess> OnDiskCache<'sess> {
             file_index_to_file: &self.file_index_to_file,
             file_index_to_stable_id: &self.file_index_to_stable_id,
             synthetic_expansion_infos: &self.synthetic_expansion_infos,
-            prev_interpret_alloc_index: &self.prev_interpret_alloc_index,
-            interpret_alloc_cache: &self.interpret_alloc_cache,
+            alloc_decoding_session: self.alloc_decoding_state.new_decoding_session(),
         };
 
         match decode_tagged(&mut decoder, dep_node_index) {
@@ -487,9 +480,7 @@ struct CacheDecoder<'a, 'tcx: 'a, 'x> {
     synthetic_expansion_infos: &'x Lock<FxHashMap<AbsoluteBytePos, SyntaxContext>>,
     file_index_to_file: &'x Lock<FxHashMap<FileMapIndex, Lrc<FileMap>>>,
     file_index_to_stable_id: &'x FxHashMap<FileMapIndex, StableFilemapId>,
-    interpret_alloc_cache: &'x RefCell<FxHashMap<usize, interpret::AllocId>>,
-    /// maps from index in the cache file to location in the cache file
-    prev_interpret_alloc_index: &'x [AbsoluteBytePos],
+    alloc_decoding_session: AllocDecodingSession<'x>,
 }
 
 impl<'a, 'tcx, 'x> CacheDecoder<'a, 'tcx, 'x> {
@@ -612,30 +603,8 @@ implement_ty_decoder!( CacheDecoder<'a, 'tcx, 'x> );
 
 impl<'a, 'tcx, 'x> SpecializedDecoder<interpret::AllocId> for CacheDecoder<'a, 'tcx, 'x> {
     fn specialized_decode(&mut self) -> Result<interpret::AllocId, Self::Error> {
-        let tcx = self.tcx;
-        let idx = usize::decode(self)?;
-        trace!("loading index {}", idx);
-
-        if let Some(cached) = self.interpret_alloc_cache.borrow().get(&idx).cloned() {
-            trace!("loading alloc id {:?} from alloc_cache", cached);
-            return Ok(cached);
-        }
-        let pos = self.prev_interpret_alloc_index[idx].to_usize();
-        trace!("loading position {}", pos);
-        self.with_position(pos, |this| {
-            interpret::specialized_decode_alloc_id(
-                this,
-                tcx,
-                |this, alloc_id| {
-                    trace!("caching idx {} for alloc id {} at position {}", idx, alloc_id, pos);
-                    assert!(this
-                        .interpret_alloc_cache
-                        .borrow_mut()
-                        .insert(idx, alloc_id)
-                        .is_none());
-                },
-            )
-        })
+        let alloc_decoding_session = self.alloc_decoding_session;
+        alloc_decoding_session.decode_alloc_id(self)
     }
 }
 impl<'a, 'tcx, 'x> SpecializedDecoder<Span> for CacheDecoder<'a, 'tcx, 'x> {
diff --git a/src/librustc_data_structures/lib.rs b/src/librustc_data_structures/lib.rs
index e2a80acbd12..23a920739b9 100644
--- a/src/librustc_data_structures/lib.rs
+++ b/src/librustc_data_structures/lib.rs
@@ -74,6 +74,7 @@ pub mod control_flow_graph;
 pub mod flock;
 pub mod sync;
 pub mod owning_ref;
+pub mod tiny_list;
 pub mod sorted_map;
 
 pub struct OnDrop<F: Fn()>(pub F);
diff --git a/src/librustc_data_structures/tiny_list.rs b/src/librustc_data_structures/tiny_list.rs
new file mode 100644
index 00000000000..5b1b2aadec5
--- /dev/null
+++ b/src/librustc_data_structures/tiny_list.rs
@@ -0,0 +1,251 @@
+// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+
+//! A singly-linked list.
+//!
+//! Using this data structure only makes sense under very specific
+//! circumstances:
+//!
+//! - If you have a list that rarely stores more than one element, then this
+//!   data-structure can store the element without allocating and only uses as
+//!   much space as a `Option<(T, usize)>`. If T can double as the `Option`
+//!   discriminant, it will even only be as large as `T, usize`.
+//!
+//! If you expect to store more than 1 element in the common case, steer clear
+//! and use a `Vec<T>`, `Box<[T]>`, or a `SmallVec<T>`.
+
+use std::mem;
+
+#[derive(Clone, Hash, Debug, PartialEq)]
+pub struct TinyList<T: PartialEq> {
+    head: Option<Element<T>>
+}
+
+impl<T: PartialEq> TinyList<T> {
+
+    #[inline]
+    pub fn new() -> TinyList<T> {
+        TinyList {
+            head: None
+        }
+    }
+
+    #[inline]
+    pub fn new_single(data: T) -> TinyList<T> {
+        TinyList {
+            head: Some(Element {
+                data,
+                next: None,
+            })
+        }
+    }
+
+    #[inline]
+    pub fn insert(&mut self, data: T) {
+        let current_head = mem::replace(&mut self.head, None);
+
+        if let Some(current_head) = current_head {
+            let current_head = Box::new(current_head);
+            self.head = Some(Element {
+                data,
+                next: Some(current_head)
+            });
+        } else {
+            self.head = Some(Element {
+                data,
+                next: None,
+            })
+        }
+    }
+
+    #[inline]
+    pub fn remove(&mut self, data: &T) -> bool {
+        let remove_head = if let Some(ref mut head) = self.head {
+            if head.data == *data {
+                Some(mem::replace(&mut head.next, None))
+            } else {
+                None
+            }
+        } else {
+            return false
+        };
+
+        if let Some(remove_head) = remove_head {
+            if let Some(next) = remove_head {
+                self.head = Some(*next);
+            } else {
+                self.head = None;
+            }
+            return true
+        }
+
+        self.head.as_mut().unwrap().remove_next(data)
+    }
+
+    #[inline]
+    pub fn contains(&self, data: &T) -> bool {
+        if let Some(ref head) = self.head {
+            head.contains(data)
+        } else {
+            false
+        }
+    }
+
+    #[inline]
+    pub fn len(&self) -> usize {
+        if let Some(ref head) = self.head {
+            head.len()
+        } else {
+            0
+        }
+    }
+}
+
+#[derive(Clone, Hash, Debug, PartialEq)]
+struct Element<T: PartialEq> {
+    data: T,
+    next: Option<Box<Element<T>>>,
+}
+
+impl<T: PartialEq> Element<T> {
+
+    fn remove_next(&mut self, data: &T) -> bool {
+        let new_next = if let Some(ref mut next) = self.next {
+            if next.data != *data {
+                return next.remove_next(data)
+            } else {
+                mem::replace(&mut next.next, None)
+            }
+        } else {
+            return false
+        };
+
+        self.next = new_next;
+        return true
+    }
+
+    fn len(&self) -> usize {
+        if let Some(ref next) = self.next {
+            1 + next.len()
+        } else {
+            1
+        }
+    }
+
+    fn contains(&self, data: &T) -> bool {
+        if self.data == *data {
+            return true
+        }
+
+        if let Some(ref next) = self.next {
+            next.contains(data)
+        } else {
+            false
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_contains_and_insert() {
+        fn do_insert(i : u32) -> bool {
+            i % 2 == 0
+        }
+
+        let mut list = TinyList::new();
+
+        for i in 0 .. 10 {
+            for j in 0 .. i {
+                if do_insert(j) {
+                    assert!(list.contains(&j));
+                } else {
+                    assert!(!list.contains(&j));
+                }
+            }
+
+            assert!(!list.contains(&i));
+
+            if do_insert(i) {
+                list.insert(i);
+                assert!(list.contains(&i));
+            }
+        }
+    }
+
+    #[test]
+    fn test_remove_first() {
+        let mut list = TinyList::new();
+        list.insert(1);
+        list.insert(2);
+        list.insert(3);
+        list.insert(4);
+        assert_eq!(list.len(), 4);
+
+        assert!(list.remove(&4));
+        assert!(!list.contains(&4));
+
+        assert_eq!(list.len(), 3);
+        assert!(list.contains(&1));
+        assert!(list.contains(&2));
+        assert!(list.contains(&3));
+    }
+
+    #[test]
+    fn test_remove_last() {
+        let mut list = TinyList::new();
+        list.insert(1);
+        list.insert(2);
+        list.insert(3);
+        list.insert(4);
+        assert_eq!(list.len(), 4);
+
+        assert!(list.remove(&1));
+        assert!(!list.contains(&1));
+
+        assert_eq!(list.len(), 3);
+        assert!(list.contains(&2));
+        assert!(list.contains(&3));
+        assert!(list.contains(&4));
+    }
+
+    #[test]
+    fn test_remove_middle() {
+        let mut list = TinyList::new();
+        list.insert(1);
+        list.insert(2);
+        list.insert(3);
+        list.insert(4);
+        assert_eq!(list.len(), 4);
+
+        assert!(list.remove(&2));
+        assert!(!list.contains(&2));
+
+        assert_eq!(list.len(), 3);
+        assert!(list.contains(&1));
+        assert!(list.contains(&3));
+        assert!(list.contains(&4));
+    }
+
+    #[test]
+    fn test_remove_single() {
+        let mut list = TinyList::new();
+        list.insert(1);
+        assert_eq!(list.len(), 1);
+
+        assert!(list.remove(&1));
+        assert!(!list.contains(&1));
+
+        assert_eq!(list.len(), 0);
+    }
+}
diff --git a/src/librustc_metadata/creader.rs b/src/librustc_metadata/creader.rs
index 2467d5cf97c..e41b3f5f53b 100644
--- a/src/librustc_metadata/creader.rs
+++ b/src/librustc_metadata/creader.rs
@@ -19,6 +19,7 @@ use rustc::hir::def_id::{CrateNum, CRATE_DEF_INDEX};
 use rustc::hir::svh::Svh;
 use rustc::middle::allocator::AllocatorKind;
 use rustc::middle::cstore::DepKind;
+use rustc::mir::interpret::AllocDecodingState;
 use rustc::session::{Session, CrateDisambiguator};
 use rustc::session::config::{Sanitizer, self};
 use rustc_target::spec::{PanicStrategy, TargetTriple};
@@ -222,6 +223,9 @@ impl<'a> CrateLoader<'a> {
             crate_root.def_path_table.decode((&metadata, self.sess))
         });
 
+        let interpret_alloc_index: Vec<u32> = crate_root.interpret_alloc_index
+                                                        .decode(&metadata)
+                                                        .collect();
         let trait_impls = crate_root
             .impls
             .decode((&metadata, self.sess))
@@ -242,6 +246,7 @@ impl<'a> CrateLoader<'a> {
             cnum,
             dependencies: Lock::new(dependencies),
             codemap_import_info: RwLock::new(vec![]),
+            alloc_decoding_state: AllocDecodingState::new(interpret_alloc_index),
             dep_kind: Lock::new(dep_kind),
             source: cstore::CrateSource {
                 dylib,
diff --git a/src/librustc_metadata/cstore.rs b/src/librustc_metadata/cstore.rs
index 763563eabe0..2bc5f607486 100644
--- a/src/librustc_metadata/cstore.rs
+++ b/src/librustc_metadata/cstore.rs
@@ -12,10 +12,10 @@
 // crates and libraries
 
 use schema;
-
 use rustc::hir::def_id::{CrateNum, DefIndex};
 use rustc::hir::map::definitions::DefPathTable;
 use rustc::middle::cstore::{DepKind, ExternCrate, MetadataLoader};
+use rustc::mir::interpret::AllocDecodingState;
 use rustc_data_structures::indexed_vec::IndexVec;
 use rustc::util::nodemap::{FxHashMap, NodeMap};
 
@@ -66,6 +66,9 @@ pub struct CrateMetadata {
     pub dependencies: Lock<Vec<CrateNum>>,
     pub codemap_import_info: RwLock<Vec<ImportedFileMap>>,
 
+    /// Used for decoding interpret::AllocIds in a cached & thread-safe manner.
+    pub alloc_decoding_state: AllocDecodingState,
+
     pub root: schema::CrateRoot,
 
     /// For each public item in this crate, we encode a key.  When the
diff --git a/src/librustc_metadata/decoder.rs b/src/librustc_metadata/decoder.rs
index 11c653895fc..69e873bb95d 100644
--- a/src/librustc_metadata/decoder.rs
+++ b/src/librustc_metadata/decoder.rs
@@ -25,12 +25,12 @@ use rustc::hir::def_id::{CrateNum, DefId, DefIndex,
 use rustc::ich::Fingerprint;
 use rustc::middle::lang_items;
 use rustc::mir::{self, interpret};
+use rustc::mir::interpret::AllocDecodingSession;
 use rustc::session::Session;
 use rustc::ty::{self, Ty, TyCtxt};
 use rustc::ty::codec::TyDecoder;
 use rustc::mir::Mir;
 use rustc::util::captures::Captures;
-use rustc::util::nodemap::FxHashMap;
 
 use std::io;
 use std::mem;
@@ -55,11 +55,8 @@ pub struct DecodeContext<'a, 'tcx: 'a> {
 
     lazy_state: LazyState,
 
-    // interpreter allocation cache
-    interpret_alloc_cache: FxHashMap<usize, interpret::AllocId>,
-
-    // Read from the LazySeq CrateRoot::inpterpret_alloc_index on demand
-    interpret_alloc_index: Option<Vec<u32>>,
+    // Used for decoding interpret::AllocIds in a cached & thread-safe manner.
+    alloc_decoding_session: Option<AllocDecodingSession<'a>>,
 }
 
 /// Abstract over the various ways one can create metadata decoders.
@@ -78,8 +75,9 @@ pub trait Metadata<'a, 'tcx>: Copy {
             tcx,
             last_filemap_index: 0,
             lazy_state: LazyState::NoNode,
-            interpret_alloc_cache: FxHashMap::default(),
-            interpret_alloc_index: None,
+            alloc_decoding_session: self.cdata().map(|cdata| {
+                cdata.alloc_decoding_state.new_decoding_session()
+            }),
         }
     }
 }
@@ -178,17 +176,6 @@ impl<'a, 'tcx> DecodeContext<'a, 'tcx> {
         self.lazy_state = LazyState::Previous(position + min_size);
         Ok(position)
     }
-
-    fn interpret_alloc(&mut self, idx: usize) -> usize {
-        if let Some(index) = self.interpret_alloc_index.as_mut() {
-            return index[idx] as usize;
-        }
-        let cdata = self.cdata();
-        let index: Vec<u32> = cdata.root.interpret_alloc_index.decode(cdata).collect();
-        let pos = index[idx];
-        self.interpret_alloc_index = Some(index);
-        pos as usize
-    }
 }
 
 impl<'a, 'tcx: 'a> TyDecoder<'a, 'tcx> for DecodeContext<'a, 'tcx> {
@@ -299,22 +286,11 @@ impl<'a, 'tcx> SpecializedDecoder<LocalDefId> for DecodeContext<'a, 'tcx> {
 
 impl<'a, 'tcx> SpecializedDecoder<interpret::AllocId> for DecodeContext<'a, 'tcx> {
     fn specialized_decode(&mut self) -> Result<interpret::AllocId, Self::Error> {
-        let tcx = self.tcx.unwrap();
-        let idx = usize::decode(self)?;
-
-        if let Some(cached) = self.interpret_alloc_cache.get(&idx).cloned() {
-            return Ok(cached);
+        if let Some(alloc_decoding_session) = self.alloc_decoding_session {
+            alloc_decoding_session.decode_alloc_id(self)
+        } else {
+            bug!("Attempting to decode interpret::AllocId without CrateMetadata")
         }
-        let pos = self.interpret_alloc(idx);
-        self.with_position(pos, |this| {
-            interpret::specialized_decode_alloc_id(
-                this,
-                tcx,
-                |this, alloc_id| {
-                    assert!(this.interpret_alloc_cache.insert(idx, alloc_id).is_none());
-                },
-            )
-        })
     }
 }