Make AllocId decoding thread-safe

Author: Zoxc

src/librustc/mir/interpret/mod.rs

@@ -17,15 +17,17 @@ use mir;
 use hir::def_id::DefId;
 use ty::{self, TyCtxt, Instance};
 use ty::layout::{self, Align, HasDataLayout, Size};
+use ty::codec::{TyEncoder, TyDecoder};
 use middle::region;
 use std::iter;
 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::{Decodable, Encodable};
 use rustc_data_structures::sorted_map::SortedMap;
 use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::sync::{HashMapExt, LockGuard};
 use byteorder::{WriteBytesExt, ReadBytesExt, LittleEndian, BigEndian};
 
 #[derive(Clone, Debug, PartialEq, RustcEncodable, RustcDecodable)]
@@ -162,25 +164,37 @@ impl ::rustc_serialize::UseSpecializedDecodable for AllocId {}
 #[derive(RustcDecodable, RustcEncodable)]
 enum AllocKind {
     Alloc,
+    AllocAtPos,
     Fn,
     Static,
 }
 
 pub fn specialized_encode_alloc_id<
     'a, 'tcx,
-    E: Encoder,
+    E: TyEncoder,
+    M,
 >(
     encoder: &mut E,
     tcx: TyCtxt<'a, 'tcx, 'tcx>,
+    cache: M,
     alloc_id: AllocId,
-) -> Result<(), E::Error> {
+) -> Result<(), E::Error>
+where
+    M: Fn(&mut E) -> &mut FxHashMap<AllocId, usize>
+{
     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) => {
+            if let Some(alloc_pos) = cache(encoder).get(&alloc_id).cloned() {
+                AllocKind::AllocAtPos.encode(encoder)?;
+                return encoder.emit_usize(alloc_pos);
+            }
+            let pos = encoder.position();
             trace!("encoding {:?} with {:#?}", alloc_id, alloc);
             AllocKind::Alloc.encode(encoder)?;
             alloc.encode(encoder)?;
+            cache(encoder).insert_same(alloc_id, pos);
         }
         AllocType::Function(fn_instance) => {
             trace!("encoding {:?} with {:#?}", alloc_id, fn_instance);
@@ -199,23 +213,32 @@ pub fn specialized_encode_alloc_id<
 
 pub fn specialized_decode_alloc_id<
     'a, 'tcx,
-    D: Decoder,
-    CACHE: FnOnce(&mut D, AllocId),
+    D: TyDecoder<'a, 'tcx>,
+    CACHE: FnOnce(&mut D) -> LockGuard<'_, FxHashMap<usize, AllocId>>,
 >(
     decoder: &mut D,
     tcx: TyCtxt<'a, 'tcx, 'tcx>,
     cache: CACHE,
 ) -> Result<AllocId, D::Error> {
     match AllocKind::decode(decoder)? {
+        AllocKind::AllocAtPos => {
+            let pos = decoder.read_usize()?;
+            if let Some(alloc_id) = cache(decoder).get(&pos).cloned() {
+                return Ok(alloc_id);
+            }
+            decoder.with_position(pos, AllocId::decode)
+        },
         AllocKind::Alloc => {
-            let alloc_id = tcx.alloc_map.lock().reserve();
-            trace!("creating alloc id {:?}", alloc_id);
+            let pos = decoder.position();
             // insert early to allow recursive allocs
-            cache(decoder, alloc_id);
+            let alloc_id = *cache(decoder).entry(pos)
+                                          .or_insert_with(|| tcx.alloc_map.lock().reserve());
+            trace!("creating alloc id {:?}", 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);
+            // This may overwrite with the same allocation
+            tcx.alloc_map.lock().set_id_same_memory(alloc_id, allocation);
 
             Ok(alloc_id)
         },
@@ -225,14 +248,12 @@ pub fn specialized_decode_alloc_id<
             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)
         },
     }
@@ -333,6 +354,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);
         }
     }
+
+    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)]

src/librustc/ty/maps/on_disk_cache.rs

@@ -23,7 +23,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 +76,8 @@ 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>>,
+    interpret_alloc_cache: Lock<FxHashMap<usize, interpret::AllocId>>,
 }
 
 // This type is used only for (de-)serialization.
@@ -91,8 +87,6 @@ struct Footer {
     prev_cnums: Vec<(u32, String, CrateDisambiguator)>,
     query_result_index: EncodedQueryResultIndex,
     diagnostics_index: EncodedQueryResultIndex,
-    // the location of all allocations
-    interpret_alloc_index: Vec<AbsoluteBytePos>,
 }
 
 type EncodedQueryResultIndex = Vec<(SerializedDepNodeIndex, AbsoluteBytePos)>;
@@ -149,8 +143,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()),
+            interpret_alloc_cache: Lock::new(FxHashMap::default()),
         }
     }
 
@@ -166,8 +159,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()),
+            interpret_alloc_cache: Lock::new(FxHashMap::default()),
         }
     }
 
@@ -201,7 +193,6 @@ impl<'sess> OnDiskCache<'sess> {
                 predicate_shorthands: FxHashMap(),
                 expn_info_shorthands: FxHashMap(),
                 interpret_allocs: FxHashMap(),
-                interpret_allocs_inverse: Vec::new(),
                 codemap: CachingCodemapView::new(tcx.sess.codemap()),
                 file_to_file_index,
             };
@@ -279,31 +270,6 @@ impl<'sess> OnDiskCache<'sess> {
                 diagnostics_index
             };
 
-            let interpret_alloc_index = {
-                let mut interpret_alloc_index = Vec::new();
-                let mut n = 0;
-                loop {
-                    let new_n = encoder.interpret_allocs_inverse.len();
-                    // if we have found new ids, serialize those, too
-                    if n == new_n {
-                        // otherwise, abort
-                        break;
-                    }
-                    for idx in n..new_n {
-                        let id = encoder.interpret_allocs_inverse[idx];
-                        let pos = AbsoluteBytePos::new(encoder.position());
-                        interpret_alloc_index.push(pos);
-                        interpret::specialized_encode_alloc_id(
-                            &mut encoder,
-                            tcx,
-                            id,
-                        )?;
-                    }
-                    n = new_n;
-                }
-                interpret_alloc_index
-            };
-
             let sorted_cnums = sorted_cnums_including_local_crate(tcx);
             let prev_cnums: Vec<_> = sorted_cnums.iter().map(|&cnum| {
                 let crate_name = tcx.original_crate_name(cnum).as_str().to_string();
@@ -318,7 +284,6 @@ impl<'sess> OnDiskCache<'sess> {
                 prev_cnums,
                 query_result_index,
                 diagnostics_index,
-                interpret_alloc_index,
             })?;
 
             // Encode the position of the footer as the last 8 bytes of the
@@ -424,7 +389,6 @@ 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,
         };
 
@@ -487,9 +451,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],
+    interpret_alloc_cache: &'x Lock<FxHashMap<usize, interpret::AllocId>>,
 }
 
 impl<'a, 'tcx, 'x> CacheDecoder<'a, 'tcx, 'x> {
@@ -613,31 +575,14 @@ 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());
-                },
-            )
-        })
+        interpret::specialized_decode_alloc_id(
+            self,
+            tcx,
+            |this| this.interpret_alloc_cache.lock(),
+        )
     }
 }
+
 impl<'a, 'tcx, 'x> SpecializedDecoder<Span> for CacheDecoder<'a, 'tcx, 'x> {
     fn specialized_decode(&mut self) -> Result<Span, Self::Error> {
         let tag: u8 = Decodable::decode(self)?;
@@ -800,7 +745,6 @@ struct CacheEncoder<'enc, 'a, 'tcx, E>
     predicate_shorthands: FxHashMap<ty::Predicate<'tcx>, usize>,
     expn_info_shorthands: FxHashMap<Mark, AbsoluteBytePos>,
     interpret_allocs: FxHashMap<interpret::AllocId, usize>,
-    interpret_allocs_inverse: Vec<interpret::AllocId>,
     codemap: CachingCodemapView<'tcx>,
     file_to_file_index: FxHashMap<*const FileMap, FileMapIndex>,
 }
@@ -837,18 +781,13 @@ impl<'enc, 'a, 'tcx, E> SpecializedEncoder<interpret::AllocId> for CacheEncoder<
     where E: 'enc + ty_codec::TyEncoder
 {
     fn specialized_encode(&mut self, alloc_id: &interpret::AllocId) -> Result<(), Self::Error> {
-        use std::collections::hash_map::Entry;
-        let index = match self.interpret_allocs.entry(*alloc_id) {
-            Entry::Occupied(e) => *e.get(),
-            Entry::Vacant(e) => {
-                let idx = self.interpret_allocs_inverse.len();
-                self.interpret_allocs_inverse.push(*alloc_id);
-                e.insert(idx);
-                idx
-            },
-        };
-
-        index.encode(self)
+        let tcx = self.tcx;
+        interpret::specialized_encode_alloc_id(
+            self,
+            tcx,
+            |this| &mut this.interpret_allocs,
+            *alloc_id,
+        )
     }
 }
 

src/librustc_metadata/decoder.rs

@@ -13,7 +13,7 @@
 use cstore::{self, CrateMetadata, MetadataBlob, NativeLibrary, ForeignModule};
 use schema::*;
 
-use rustc_data_structures::sync::{Lrc, ReadGuard};
+use rustc_data_structures::sync::{Lrc, ReadGuard, Lock};
 use rustc::hir::map::{DefKey, DefPath, DefPathData, DefPathHash,
                       DisambiguatedDefPathData};
 use rustc::hir;
@@ -56,10 +56,7 @@ 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>>,
+    interpret_alloc_cache: Lock<FxHashMap<usize, interpret::AllocId>>,
 }
 
 /// Abstract over the various ways one can create metadata decoders.
@@ -78,8 +75,7 @@ pub trait Metadata<'a, 'tcx>: Copy {
             tcx,
             last_filemap_index: 0,
             lazy_state: LazyState::NoNode,
-            interpret_alloc_cache: FxHashMap::default(),
-            interpret_alloc_index: None,
+            interpret_alloc_cache: Lock::new(FxHashMap::default()),
         }
     }
 }
@@ -178,17 +174,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> {
@@ -300,21 +285,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);
-        }
-        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());
-                },
-            )
-        })
+        interpret::specialized_decode_alloc_id(
+            self,
+            tcx,
+            |this| this.interpret_alloc_cache.lock()
+        )
     }
 }