From e9d69d9f8eb888a6e124c567f804c2e464c7b00a Mon Sep 17 00:00:00 2001
From: Smitty <me@smitop.com>
Date: Fri, 2 Jul 2021 16:06:12 -0400
Subject: [PATCH] Allocation failure in constprop panics right away

---
 compiler/rustc_middle/src/mir/interpret/allocation.rs | 5 ++++-
 compiler/rustc_middle/src/ty/vtable.rs                | 2 +-
 compiler/rustc_mir/src/const_eval/machine.rs          | 2 ++
 compiler/rustc_mir/src/interpret/machine.rs           | 3 +++
 compiler/rustc_mir/src/interpret/memory.rs            | 2 +-
 compiler/rustc_mir/src/transform/const_prop.rs        | 1 +
 6 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/compiler/rustc_middle/src/mir/interpret/allocation.rs b/compiler/rustc_middle/src/mir/interpret/allocation.rs
index 49e0af9a3a48..f1b1bb9ab88a 100644
--- a/compiler/rustc_middle/src/mir/interpret/allocation.rs
+++ b/compiler/rustc_middle/src/mir/interpret/allocation.rs
@@ -126,7 +126,7 @@ impl<Tag> Allocation<Tag> {
 
     /// 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) -> InterpResult<'static, Self> {
+    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
@@ -134,6 +134,9 @@ impl<Tag> Allocation<Tag> {
             // 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")
             });
diff --git a/compiler/rustc_middle/src/ty/vtable.rs b/compiler/rustc_middle/src/ty/vtable.rs
index 0230e05c12e8..ef1c941bee4e 100644
--- a/compiler/rustc_middle/src/ty/vtable.rs
+++ b/compiler/rustc_middle/src/ty/vtable.rs
@@ -60,7 +60,7 @@ 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, true)?;
 
         // 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 --git a/compiler/rustc_mir/src/const_eval/machine.rs b/compiler/rustc_mir/src/const_eval/machine.rs
index ddc87084e9fd..992e32e298f8 100644
--- a/compiler/rustc_mir/src/const_eval/machine.rs
+++ b/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>,
diff --git a/compiler/rustc_mir/src/interpret/machine.rs b/compiler/rustc_mir/src/interpret/machine.rs
index 0d01dc3c219b..5b8c0788cbc8 100644
--- a/compiler/rustc_mir/src/interpret/machine.rs
+++ b/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 --git a/compiler/rustc_mir/src/interpret/memory.rs b/compiler/rustc_mir/src/interpret/memory.rs
index 671e3d278f36..cb929c21850c 100644
--- a/compiler/rustc_mir/src/interpret/memory.rs
+++ b/compiler/rustc_mir/src/interpret/memory.rs
@@ -208,7 +208,7 @@ impl<'mir, 'tcx, M: Machine<'mir, 'tcx>> Memory<'mir, 'tcx, M> {
         align: Align,
         kind: MemoryKind<M::MemoryKind>,
     ) -> InterpResult<'static, Pointer<M::PointerTag>> {
-        let alloc = Allocation::uninit(size, align)?;
+        let alloc = Allocation::uninit(size, align, M::PANIC_ON_ALLOC_FAIL)?;
         Ok(self.allocate_with(alloc, kind))
     }
 
diff --git a/compiler/rustc_mir/src/transform/const_prop.rs b/compiler/rustc_mir/src/transform/const_prop.rs
index 2a6924d354a8..743ba95a9afa 100644
--- a/compiler/rustc_mir/src/transform/const_prop.rs
+++ b/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
 
     type MemoryKind = !;