miri engine: lazily allocate memory for locals on first write

Author: RalfJung

src/librustc_mir/interpret/eval_context.rs

@@ -116,26 +116,41 @@ pub struct LocalState<'tcx, Tag=(), Id=AllocId> {
 /// State of a local variable
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 pub enum LocalValue<Tag=(), Id=AllocId> {
+    /// This local is not currently alive, and cannot be used at all.
     Dead,
-    // Mostly for convenience, we re-use the `Operand` type here.
-    // This is an optimization over just always having a pointer here;
-    // we can thus avoid doing an allocation when the local just stores
-    // immediate values *and* never has its address taken.
+    /// This local is alive but not yet initialized. It can be written to
+    /// but not read from or its address taken. Locals get initialized on
+    /// first write because for unsized locals, we do not know their size
+    /// before that.
+    Uninitialized,
+    /// A normal, live local.
+    /// Mostly for convenience, we re-use the `Operand` type here.
+    /// This is an optimization over just always having a pointer here;
+    /// we can thus avoid doing an allocation when the local just stores
+    /// immediate values *and* never has its address taken.
     Live(Operand<Tag, Id>),
 }
 
-impl<'tcx, Tag> LocalState<'tcx, Tag> {
+impl<'tcx, Tag: Copy> LocalState<'tcx, Tag> {
     pub fn access(&self) -> EvalResult<'tcx, &Operand<Tag>> {
         match self.state {
-            LocalValue::Dead => err!(DeadLocal),
+            LocalValue::Dead | LocalValue::Uninitialized => err!(DeadLocal),
             LocalValue::Live(ref val) => Ok(val),
         }
     }
 
-    pub fn access_mut(&mut self) -> EvalResult<'tcx, &mut Operand<Tag>> {
+    /// Overwrite the local.  If the local can be overwritten in place, return a reference
+    /// to do so; otherwise return the `MemPlace` to consult instead.
+    pub fn access_mut(
+        &mut self,
+    ) -> EvalResult<'tcx, Result<&mut LocalValue<Tag>, MemPlace<Tag>>> {
         match self.state {
             LocalValue::Dead => err!(DeadLocal),
-            LocalValue::Live(ref mut val) => Ok(val),
+            LocalValue::Live(Operand::Indirect(mplace)) => Ok(Err(mplace)),
+            ref mut local @ LocalValue::Live(Operand::Immediate(_)) |
+            ref mut local @ LocalValue::Uninitialized => {
+                Ok(Ok(local))
+            }
         }
     }
 }
@@ -327,6 +342,7 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'a, 'mir, 'tcx>> InterpretCx<'a, 'mir, 'tc
                     let local_ty = self.monomorphize_with_substs(local_ty, frame.instance.substs);
                     self.layout_of(local_ty)
                 })?;
+                // Layouts of locals are requested a lot, so we cache them.
                 frame.locals[local].layout.set(Some(layout));
                 Ok(layout)
             }
@@ -473,13 +489,9 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'a, 'mir, 'tcx>> InterpretCx<'a, 'mir, 'tc
 
         // don't allocate at all for trivial constants
         if mir.local_decls.len() > 1 {
-            // We put some marker immediate into the locals that we later want to initialize.
-            // This can be anything except for LocalValue::Dead -- because *that* is the
-            // value we use for things that we know are initially dead.
+            // Locals are initially uninitialized.
             let dummy = LocalState {
-                state: LocalValue::Live(Operand::Immediate(Immediate::Scalar(
-                    ScalarMaybeUndef::Undef,
-                ))),
+                state: LocalValue::Uninitialized,
                 layout: Cell::new(None),
             };
             let mut locals = IndexVec::from_elem(dummy, &mir.local_decls);
@@ -506,19 +518,25 @@ impl<'a, 'mir, 'tcx: 'mir, M: Machine<'a, 'mir, 'tcx>> InterpretCx<'a, 'mir, 'tc
                     }
                 },
             }
-            // Finally, properly initialize all those that still have the dummy value
+            // FIXME: We initialize live ZST here.  This should not be needed if MIR was
+            // consistently generated for ZST, but that seems to not be the case -- there
+            // is MIR (around promoteds in particular) that reads local ZSTs that never
+            // were written to.
             for (idx, local) in locals.iter_enumerated_mut() {
                 match local.state {
-                    LocalValue::Live(_) => {
+                    LocalValue::Uninitialized => {
                         // This needs to be properly initialized.
                         let ty = self.monomorphize(mir.local_decls[idx].ty)?;
                         let layout = self.layout_of(ty)?;
-                        local.state = LocalValue::Live(self.uninit_operand(layout)?);
+                        if layout.is_zst() {
+                            local.state = LocalValue::Live(self.uninit_operand(layout)?);
+                        }
                         local.layout = Cell::new(Some(layout));
                     }
                     LocalValue::Dead => {
                         // Nothing to do
                     }
+                    LocalValue::Live(_) => bug!("Locals cannot be live yet"),
                 }
             }
             // done

src/librustc_mir/interpret/place.rs

@@ -15,7 +15,7 @@ use rustc::ty::TypeFoldable;
 use super::{
     GlobalId, AllocId, Allocation, Scalar, EvalResult, Pointer, PointerArithmetic,
     InterpretCx, Machine, AllocMap, AllocationExtra,
-    RawConst, Immediate, ImmTy, ScalarMaybeUndef, Operand, OpTy, MemoryKind
+    RawConst, Immediate, ImmTy, ScalarMaybeUndef, Operand, OpTy, MemoryKind, LocalValue
 };
 
 #[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
@@ -639,6 +639,7 @@ where
                 None => return err!(InvalidNullPointerUsage),
             },
             Base(PlaceBase::Local(local)) => PlaceTy {
+                // This works even for dead/uninitialized locals; we check further when writing
                 place: Place::Local {
                     frame: self.cur_frame(),
                     local,
@@ -714,16 +715,19 @@ where
         // but not factored as a separate function.
         let mplace = match dest.place {
             Place::Local { frame, local } => {
-                match *self.stack[frame].locals[local].access_mut()? {
-                    Operand::Immediate(ref mut dest_val) => {
-                        // Yay, we can just change the local directly.
-                        *dest_val = src;
+                match self.stack[frame].locals[local].access_mut()? {
+                    Ok(local) => {
+                        // Local can be updated in-place.
+                        *local = LocalValue::Live(Operand::Immediate(src));
                         return Ok(());
-                    },
-                    Operand::Indirect(mplace) => mplace, // already in memory
+                    }
+                    Err(mplace) => {
+                        // The local is in memory, go on below.
+                        mplace
+                    }
                 }
             },
-            Place::Ptr(mplace) => mplace, // already in memory
+            Place::Ptr(mplace) => mplace, // already referring to memory
         };
         let dest = MPlaceTy { mplace, layout: dest.layout };
 
@@ -904,27 +908,40 @@ where
     ) -> EvalResult<'tcx, MPlaceTy<'tcx, M::PointerTag>> {
         let mplace = match place.place {
             Place::Local { frame, local } => {
-                match *self.stack[frame].locals[local].access()? {
-                    Operand::Indirect(mplace) => mplace,
-                    Operand::Immediate(value) => {
+                match self.stack[frame].locals[local].access_mut()? {
+                    Ok(local_val) => {
                         // We need to make an allocation.
                         // FIXME: Consider not doing anything for a ZST, and just returning
                         // a fake pointer?  Are we even called for ZST?
 
+                        // We cannot hold on to the reference `local_val` while allocating,
+                        // but we can hold on to the value in there.
+                        let old_val =
+                            if let LocalValue::Live(Operand::Immediate(value)) = *local_val {
+                                Some(value)
+                            } else {
+                                None
+                            };
+
                         // We need the layout of the local.  We can NOT use the layout we got,
                         // that might e.g., be an inner field of a struct with `Scalar` layout,
                         // that has different alignment than the outer field.
                         let local_layout = self.layout_of_local(&self.stack[frame], local, None)?;
                         let ptr = self.allocate(local_layout, MemoryKind::Stack);
-                        // We don't have to validate as we can assume the local
-                        // was already valid for its type.
-                        self.write_immediate_to_mplace_no_validate(value, ptr)?;
+                        if let Some(value) = old_val {
+                            // Preserve old value.
+                            // We don't have to validate as we can assume the local
+                            // was already valid for its type.
+                            self.write_immediate_to_mplace_no_validate(value, ptr)?;
+                        }
                         let mplace = ptr.mplace;
-                        // Update the local
-                        *self.stack[frame].locals[local].access_mut()? =
-                            Operand::Indirect(mplace);
+                        // Now we can call `access_mut` again, asserting it goes well,
+                        // and actually overwrite things.
+                        *self.stack[frame].locals[local].access_mut().unwrap().unwrap() =
+                            LocalValue::Live(Operand::Indirect(mplace));
                         mplace
                     }
+                    Err(mplace) => mplace, // this already was an indirect local
                 }
             }
             Place::Ptr(mplace) => mplace

src/librustc_mir/interpret/snapshot.rs

@@ -114,10 +114,11 @@ macro_rules! impl_snapshot_for {
             fn snapshot(&self, __ctx: &'a Ctx) -> Self::Item {
                 match *self {
                     $(
-                        $enum_name::$variant $( ( $(ref $field),* ) )? =>
+                        $enum_name::$variant $( ( $(ref $field),* ) )? => {
                             $enum_name::$variant $(
-                                ( $( __impl_snapshot_field!($field, __ctx $(, $delegate)?) ),* ),
+                                ( $( __impl_snapshot_field!($field, __ctx $(, $delegate)?) ),* )
                             )?
+                        }
                     )*
                 }
             }
@@ -250,11 +251,13 @@ impl_snapshot_for!(enum Operand {
 
 impl_stable_hash_for!(enum crate::interpret::LocalValue {
     Dead,
+    Uninitialized,
     Live(x),
 });
 impl_snapshot_for!(enum LocalValue {
-    Live(v),
     Dead,
+    Uninitialized,
+    Live(v),
 });
 
 impl<'a, Ctx> Snapshot<'a, Ctx> for Relocations

src/librustc_mir/interpret/terminator.rs

@@ -315,12 +315,13 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> InterpretCx<'a, 'mir, 'tcx, M>
                     );
 
                     // Figure out how to pass which arguments.
-                    // We have two iterators: Where the arguments come from,
-                    // and where they go to.
+                    // The Rust ABI is special: ZST get skipped.
                     let rust_abi = match caller_abi {
                         Abi::Rust | Abi::RustCall => true,
                         _ => false
                     };
+                    // We have two iterators: Where the arguments come from,
+                    // and where they go to.
 
                     // For where they come from: If the ABI is RustCall, we untuple the
                     // last incoming argument.  These two iterators do not have the same type,
@@ -368,7 +369,7 @@ impl<'a, 'mir, 'tcx, M: Machine<'a, 'mir, 'tcx>> InterpretCx<'a, 'mir, 'tcx, M>
                     }
                     // Now we should have no more caller args
                     if caller_iter.next().is_some() {
-                        trace!("Caller has too many args over");
+                        trace!("Caller has passed too many args");
                         return err!(FunctionArgCountMismatch);
                     }
                     // Don't forget to check the return type!