Auto merge of #86255 - Smittyvb:mir-alloc-oom, r=RalfJung,oli-obk

Support allocation failures when interpreting MIR

This closes #79601 by handling the case where memory allocation fails during MIR interpretation, and translates that failure into an `InterpError`. The error message is "tried to allocate more memory than available to compiler" to make it clear that the memory shortage is happening at compile-time by the compiler itself, and that it is not a runtime issue.

Now that memory allocation can fail, it would be neat if Miri could simulate low-memory devices to make it easy to see how much memory a Rust program needs.

Note that this breaks Miri because it assumes that allocation can never fail.

compiler/rustc_mir/src/const_eval/eval_queries.rs

@@ -48,7 +48,7 @@ fn eval_body_using_ecx<'mir, 'tcx>(
     );
     let layout = ecx.layout_of(body.return_ty().subst(tcx, cid.instance.substs))?;
     assert!(!layout.is_unsized());
-    let ret = ecx.allocate(layout, MemoryKind::Stack);
+    let ret = ecx.allocate(layout, MemoryKind::Stack)?;
 
     let name =
         with_no_trimmed_paths(|| ty::tls::with(|tcx| tcx.def_path_str(cid.instance.def_id())));

compiler/rustc_mir/src/const_eval/machine.rs

@@ -201,6 +201,8 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
 
     type MemoryExtra = MemoryExtra;
 
+    const PANIC_ON_ALLOC_FAIL: bool = false; // will be raised as a proper error
+
     fn load_mir(
         ecx: &InterpCx<'mir, 'tcx, Self>,
         instance: ty::InstanceDef<'tcx>,
@@ -306,7 +308,7 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
                     Size::from_bytes(size as u64),
                     align,
                     interpret::MemoryKind::Machine(MemoryKind::Heap),
-                );
+                )?;
                 ecx.write_scalar(Scalar::Ptr(ptr), dest)?;
             }
             _ => {

compiler/rustc_mir/src/interpret/intern.rs

@@ -428,7 +428,7 @@ impl<'mir, 'tcx: 'mir, M: super::intern::CompileTimeMachine<'mir, 'tcx, !>>
             &MPlaceTy<'tcx, M::PointerTag>,
         ) -> InterpResult<'tcx, ()>,
     ) -> InterpResult<'tcx, &'tcx Allocation> {
-        let dest = self.allocate(layout, MemoryKind::Stack);
+        let dest = self.allocate(layout, MemoryKind::Stack)?;
         f(self, &dest)?;
         let ptr = dest.ptr.assert_ptr();
         assert_eq!(ptr.offset, Size::ZERO);

compiler/rustc_mir/src/interpret/intrinsics/caller_location.rs

@@ -91,7 +91,9 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             .type_of(self.tcx.require_lang_item(LangItem::PanicLocation, None))
             .subst(*self.tcx, self.tcx.mk_substs([self.tcx.lifetimes.re_erased.into()].iter()));
         let loc_layout = self.layout_of(loc_ty).unwrap();
-        let location = self.allocate(loc_layout, MemoryKind::CallerLocation);
+        // This can fail if rustc runs out of memory right here. Trying to emit an error would be
+        // pointless, since that would require allocating more memory than a Location.
+        let location = self.allocate(loc_layout, MemoryKind::CallerLocation).unwrap();
 
         // Initialize fields.
         self.write_immediate(file.to_ref(), &self.mplace_field(&location, 0).unwrap().into())

compiler/rustc_mir/src/interpret/machine.rs

@@ -122,6 +122,9 @@ pub trait Machine<'mir, 'tcx>: Sized {
     /// that is added to the memory so that the work is not done twice.
     const GLOBAL_KIND: Option<Self::MemoryKind>;
 
+    /// Should the machine panic on allocation failures?
+    const PANIC_ON_ALLOC_FAIL: bool;
+
     /// Whether memory accesses should be alignment-checked.
     fn enforce_alignment(memory_extra: &Self::MemoryExtra) -> bool;
 

compiler/rustc_mir/src/interpret/memory.rs

@@ -207,9 +207,9 @@ impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> Memory<'mir, 'tcx, M> {
         size: Size,
         align: Align,
         kind: MemoryKind<M::MemoryKind>,
-    ) -> Pointer<M::PointerTag> {
-        let alloc = Allocation::uninit(size, align);
-        self.allocate_with(alloc, kind)
+    ) -> InterpResult<'static, Pointer<M::PointerTag>> {
+        let alloc = Allocation::uninit(size, align, M::PANIC_ON_ALLOC_FAIL)?;
+        Ok(self.allocate_with(alloc, kind))
     }
 
     pub fn allocate_bytes(
@@ -257,7 +257,7 @@ impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> Memory<'mir, 'tcx, M> {
 
         // For simplicities' sake, we implement reallocate as "alloc, copy, dealloc".
         // This happens so rarely, the perf advantage is outweighed by the maintenance cost.
-        let new_ptr = self.allocate(new_size, new_align, kind);
+        let new_ptr = self.allocate(new_size, new_align, kind)?;
         let old_size = match old_size_and_align {
             Some((size, _align)) => size,
             None => self.get_raw(ptr.alloc_id)?.size(),

compiler/rustc_mir/src/interpret/place.rs

@@ -982,7 +982,7 @@ where
                         let (size, align) = self
                             .size_and_align_of(&meta, &local_layout)?
                             .expect("Cannot allocate for non-dyn-sized type");
-                        let ptr = self.memory.allocate(size, align, MemoryKind::Stack);
+                        let ptr = self.memory.allocate(size, align, MemoryKind::Stack)?;
                         let mplace = MemPlace { ptr: ptr.into(), align, meta };
                         if let LocalValue::Live(Operand::Immediate(value)) = local_val {
                             // Preserve old value.
@@ -1018,9 +1018,9 @@ where
         &mut self,
         layout: TyAndLayout<'tcx>,
         kind: MemoryKind<M::MemoryKind>,
-    ) -> MPlaceTy<'tcx, M::PointerTag> {
-        let ptr = self.memory.allocate(layout.size, layout.align.abi, kind);
-        MPlaceTy::from_aligned_ptr(ptr, layout)
+    ) -> InterpResult<'static, MPlaceTy<'tcx, M::PointerTag>> {
+        let ptr = self.memory.allocate(layout.size, layout.align.abi, kind)?;
+        Ok(MPlaceTy::from_aligned_ptr(ptr, layout))
     }
 
     /// Returns a wide MPlace of type `&'static [mut] str` to a new 1-aligned allocation.

compiler/rustc_mir/src/lib.rs

@@ -29,6 +29,7 @@ Rust MIR: a lowered representation of Rust.
 #![feature(option_get_or_insert_default)]
 #![feature(once_cell)]
 #![feature(control_flow_enum)]
+#![feature(try_reserve)]
 #![recursion_limit = "256"]
 
 #[macro_use]

compiler/rustc_mir/src/transform/const_prop.rs

@@ -181,6 +181,7 @@ impl<'mir, 'tcx> ConstPropMachine<'mir, 'tcx> {
 
 impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for ConstPropMachine<'mir, 'tcx> {
     compile_time_machine!(<'mir, 'tcx>);
+    const PANIC_ON_ALLOC_FAIL: bool = true; // all allocations are small (see `MAX_ALLOC_LIMIT`)
 
     type MemoryKind = !;
 
@@ -393,7 +394,11 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
             .filter(|ret_layout| {
                 !ret_layout.is_zst() && ret_layout.size < Size::from_bytes(MAX_ALLOC_LIMIT)
             })
-            .map(|ret_layout| ecx.allocate(ret_layout, MemoryKind::Stack).into());
+            .map(|ret_layout| {
+                ecx.allocate(ret_layout, MemoryKind::Stack)
+                    .expect("couldn't perform small allocation")
+                    .into()
+            });
 
         ecx.push_stack_frame(
             Instance::new(def_id, substs),