Encode the allocation size so we can skip ahead

Author: Zoxc

src/librustc/mir/interpret/mod.rs

@@ -12,7 +12,6 @@ pub use self::error::{EvalError, EvalResult, EvalErrorKind, AssertMessage};
 
 pub use self::value::{PrimVal, PrimValKind, Value, Pointer, ConstValue};
 
-use std::collections::hash_map::Entry;
 use std::fmt;
 use mir;
 use hir::def_id::DefId;
@@ -195,7 +194,21 @@ where
             assert!(cache(encoder).insert(alloc_id, pos).is_none());
             trace!("encoding {:?} with {:#?}", alloc_id, alloc);
             AllocKind::Alloc.encode(encoder)?;
+
+            // Write placeholder for size
+            let size_pos = encoder.position();
+            0usize.encode(encoder)?;
+
+            let start = encoder.position();
             alloc.encode(encoder)?;
+            let end = encoder.position();
+
+            // Overwrite the placeholder with the real size
+            let size: usize = end - start;
+            encoder.set_position(size_pos);
+            size.encode(encoder)?;
+            encoder.set_position(end);
+
         }
         AllocType::Function(fn_instance) => {
             trace!("encoding {:?} with {:#?}", alloc_id, fn_instance);
@@ -230,40 +243,23 @@ pub fn specialized_decode_alloc_id<
             decoder.with_position(real_pos, AllocId::decode)
         },
         AllocKind::Alloc => {
-            let alloc_id = {
-                let mut cache = global_cache(decoder);
-                let entry = cache.entry(pos);
-                match entry {
-                    Entry::Occupied(occupied) => {
-                        let id = occupied.get().0;
-
-                        // If the alloc id is fully loaded we just return here.
-                        if occupied.get().1 {
-                            return Ok(id)
-                        }
-
-                        // It was only partially loaded.
-                        // This may be loading further up the stack
-                        // or concurrently in another thread.
-                        id
-                    }
-                    Entry::Vacant(vacant) => {
-                        // Insert early to allow recursive allocs
-                        let id = tcx.alloc_map.lock().reserve();
-                        vacant.insert((id, false));
-                        id
-                    }
-                }
-            };
-
-            // Insert early to allow recursive allocs and ot indicate that the current
-            // session will eventually fully load this alloc id
-            if !local_cache(decoder).insert(alloc_id) {
+            // Read the size of the allocation.
+            // Used to skip ahead if we don't need to decode this.
+            let alloc_size = usize::decode(decoder)?;
+
+            let (alloc_id, fully_loaded) = *global_cache(decoder).entry(pos).or_insert_with(|| {
+                // Create an id which is not fully loaded
+                (tcx.alloc_map.lock().reserve(), false)
+            });
+            if fully_loaded || !local_cache(decoder).insert(alloc_id) {
                 // We have started decoding this alloc id already, so just return it.
                 // Its content is already filled in or will be filled in by functions
                 // further up the stack.
-                return Ok(alloc_id);
-                
+
+                // Skip the allocation
+                let pos = decoder.position();
+                decoder.set_position(pos + alloc_size);
+                return Ok(alloc_id)
             }
 
             let allocation = <&'tcx Allocation as Decodable>::decode(decoder)?;

src/librustc/ty/codec.rs

@@ -52,13 +52,18 @@ impl<'tcx> EncodableWithShorthand for ty::Predicate<'tcx> {
 
 pub trait TyEncoder: Encoder {
     fn position(&self) -> usize;
+    fn set_position(&mut self, usize);
 }
 
 impl<'buf> TyEncoder for opaque::Encoder<'buf> {
     #[inline]
     fn position(&self) -> usize {
         self.position()
     }
+    #[inline]
+    fn set_position(&mut self, p: usize) {
+        self.cursor.set_position(p as u64)
+    }
 }
 
 /// Encode the given value or a previously cached shorthand.
@@ -123,6 +128,8 @@ pub trait TyDecoder<'a, 'tcx: 'a>: Decoder {
 
     fn position(&self) -> usize;
 
+    fn set_position(&mut self, usize);
+
     fn cached_ty_for_shorthand<F>(&mut self,
                                   shorthand: usize,
                                   or_insert_with: F)

src/librustc/ty/maps/on_disk_cache.rs

@@ -525,6 +525,11 @@ impl<'a, 'tcx: 'a, 'x> ty_codec::TyDecoder<'a, 'tcx> for CacheDecoder<'a, 'tcx,
         self.opaque.position()
     }
 
+    #[inline]
+    fn set_position(&mut self, p: usize) {
+        self.opaque.set_position(p)
+    }
+
     #[inline]
     fn peek_byte(&self) -> u8 {
         self.opaque.data[self.opaque.position()]
@@ -860,6 +865,10 @@ impl<'enc, 'a, 'tcx, E> ty_codec::TyEncoder for CacheEncoder<'enc, 'a, 'tcx, E>
     fn position(&self) -> usize {
         self.encoder.position()
     }
+    #[inline]
+    fn set_position(&mut self, p: usize) {
+        self.encoder.set_position(p)
+    }
 }
 
 impl<'enc, 'a, 'tcx, E> SpecializedEncoder<CrateNum> for CacheEncoder<'enc, 'a, 'tcx, E>

src/librustc_metadata/decoder.rs

@@ -195,6 +195,11 @@ impl<'a, 'tcx: 'a> TyDecoder<'a, 'tcx> for DecodeContext<'a, 'tcx> {
         self.opaque.position()
     }
 
+    #[inline]
+    fn set_position(&mut self, p: usize) {
+        self.opaque.set_position(p)
+    }
+
     fn cached_ty_for_shorthand<F>(&mut self,
                                   shorthand: usize,
                                   or_insert_with: F)

src/librustc_metadata/encoder.rs

@@ -234,6 +234,9 @@ impl<'a, 'tcx> TyEncoder for EncodeContext<'a, 'tcx> {
     fn position(&self) -> usize {
         self.opaque.position()
     }
+    fn set_position(&mut self, p: usize) {
+        self.opaque.set_position(p)
+    }
 }
 
 impl<'a, 'tcx> EncodeContext<'a, 'tcx> {