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.

Author: bors

Diff for: compiler/rustc_middle/src/lib.rs

@@ -48,6 +48,7 @@
 #![feature(associated_type_defaults)]
 #![feature(iter_zip)]
 #![feature(thread_local_const_init)]
+#![feature(try_reserve)]
 #![recursion_limit = "512"]
 
 #[macro_use]

Diff for: compiler/rustc_middle/src/mir/interpret/allocation.rs

@@ -8,12 +8,15 @@ use std::ptr;
 
 use rustc_ast::Mutability;
 use rustc_data_structures::sorted_map::SortedMap;
+use rustc_span::DUMMY_SP;
 use rustc_target::abi::{Align, HasDataLayout, Size};
 
 use super::{
-    read_target_uint, write_target_uint, AllocId, InterpError, Pointer, Scalar, ScalarMaybeUninit,
-    UndefinedBehaviorInfo, UninitBytesAccess, UnsupportedOpInfo,
+    read_target_uint, write_target_uint, AllocId, InterpError, InterpResult, Pointer,
+    ResourceExhaustionInfo, Scalar, ScalarMaybeUninit, UndefinedBehaviorInfo, UninitBytesAccess,
+    UnsupportedOpInfo,
 };
+use crate::ty;
 
 /// This type represents an Allocation in the Miri/CTFE core engine.
 ///
@@ -121,15 +124,33 @@ impl<Tag> Allocation<Tag> {
         Allocation::from_bytes(slice, Align::ONE, Mutability::Not)
     }
 
-    pub fn uninit(size: Size, align: Align) -> Self {
-        Allocation {
-            bytes: vec![0; size.bytes_usize()],
+    /// Try to create an Allocation of `size` bytes, failing if there is not enough memory
+    /// available to the compiler to do so.
+    pub fn uninit(size: Size, align: Align, panic_on_fail: bool) -> InterpResult<'static, Self> {
+        let mut bytes = Vec::new();
+        bytes.try_reserve(size.bytes_usize()).map_err(|_| {
+            // This results in an error that can happen non-deterministically, since the memory
+            // available to the compiler can change between runs. Normally queries are always
+            // deterministic. However, we can be non-determinstic here because all uses of const
+            // evaluation (including ConstProp!) will make compilation fail (via hard error
+            // or ICE) upon encountering a `MemoryExhausted` error.
+            if panic_on_fail {
+                panic!("Allocation::uninit called with panic_on_fail had allocation failure")
+            }
+            ty::tls::with(|tcx| {
+                tcx.sess.delay_span_bug(DUMMY_SP, "exhausted memory during interpreation")
+            });
+            InterpError::ResourceExhaustion(ResourceExhaustionInfo::MemoryExhausted)
+        })?;
+        bytes.resize(size.bytes_usize(), 0);
+        Ok(Allocation {
+            bytes,
             relocations: Relocations::new(),
             init_mask: InitMask::new(size, false),
             align,
             mutability: Mutability::Mut,
             extra: (),
-        }
+        })
     }
 }
 

Diff for: compiler/rustc_middle/src/mir/interpret/error.rs

@@ -423,6 +423,8 @@ pub enum ResourceExhaustionInfo {
     ///
     /// The exact limit is set by the `const_eval_limit` attribute.
     StepLimitReached,
+    /// There is not enough memory to perform an allocation.
+    MemoryExhausted,
 }
 
 impl fmt::Display for ResourceExhaustionInfo {
@@ -435,6 +437,9 @@ impl fmt::Display for ResourceExhaustionInfo {
             StepLimitReached => {
                 write!(f, "exceeded interpreter step limit (see `#[const_eval_limit]`)")
             }
+            MemoryExhausted => {
+                write!(f, "tried to allocate more memory than available to compiler")
+            }
         }
     }
 }
@@ -525,7 +530,8 @@ impl InterpError<'_> {
         use InterpError::*;
         match *self {
             MachineStop(ref err) => err.is_hard_err(),
-            InterpError::UndefinedBehavior(_) => true,
+            UndefinedBehavior(_) => true,
+            ResourceExhaustion(ResourceExhaustionInfo::MemoryExhausted) => true,
             _ => false,
         }
     }

Diff for: compiler/rustc_middle/src/ty/vtable.rs

@@ -60,7 +60,8 @@ impl<'tcx> TyCtxt<'tcx> {
         let ptr_align = tcx.data_layout.pointer_align.abi;
 
         let vtable_size = ptr_size * u64::try_from(vtable_entries.len()).unwrap();
-        let mut vtable = Allocation::uninit(vtable_size, ptr_align);
+        let mut vtable =
+            Allocation::uninit(vtable_size, ptr_align, /* panic_on_fail */ true).unwrap();
 
         // No need to do any alignment checks on the memory accesses below, because we know the
         // allocation is correctly aligned as we created it above. Also we're only offsetting by

Diff for: 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())));

Diff for: 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)?;
             }
             _ => {

Diff for: 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);

Diff for: 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())

Diff for: 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;
 

Diff for: 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(),

Diff for: 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.

Diff for: 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]

Diff for: 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),

Diff for: src/test/ui/consts/large_const_alloc.rs

@@ -0,0 +1,18 @@
+// only-64bit
+// on 32bit and 16bit platforms it is plausible that the maximum allocation size will succeed
+
+const FOO: () = {
+    // 128 TiB, unlikely anyone has that much RAM
+    let x = [0_u8; (1 << 47) - 1];
+    //~^ ERROR evaluation of constant value failed
+};
+
+static FOO2: () = {
+    let x = [0_u8; (1 << 47) - 1];
+    //~^ ERROR could not evaluate static initializer
+};
+
+fn main() {
+    let _ = FOO;
+    let _ = FOO2;
+}

Diff for: src/test/ui/consts/large_const_alloc.stderr

@@ -0,0 +1,15 @@
+error[E0080]: evaluation of constant value failed
+  --> $DIR/large_const_alloc.rs:6:13
+   |
+LL |     let x = [0_u8; (1 << 47) - 1];
+   |             ^^^^^^^^^^^^^^^^^^^^^ tried to allocate more memory than available to compiler
+
+error[E0080]: could not evaluate static initializer
+  --> $DIR/large_const_alloc.rs:11:13
+   |
+LL |     let x = [0_u8; (1 << 47) - 1];
+   |             ^^^^^^^^^^^^^^^^^^^^^ tried to allocate more memory than available to compiler
+
+error: aborting due to 2 previous errors
+
+For more information about this error, try `rustc --explain E0080`.