From d1276005119341f381ce1ec04894b24a9275d5f0 Mon Sep 17 00:00:00 2001
From: Ralf Jung <post@ralfj.de>
Date: Tue, 25 Jul 2023 09:47:19 +0200
Subject: [PATCH] add some sanity checks in write_immediate_no_validate

---
 .../rustc_const_eval/src/interpret/place.rs   | 31 +++++++++++++++++--
 .../src/mir/interpret/allocation.rs           |  2 +-
 .../rustc_middle/src/mir/interpret/value.rs   |  8 +++++
 3 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/compiler/rustc_const_eval/src/interpret/place.rs b/compiler/rustc_const_eval/src/interpret/place.rs
index 9c492fb8b93..db1239c7136 100644
--- a/compiler/rustc_const_eval/src/interpret/place.rs
+++ b/compiler/rustc_const_eval/src/interpret/place.rs
@@ -577,7 +577,7 @@ where
         src: Immediate<M::Provenance>,
         dest: &PlaceTy<'tcx, M::Provenance>,
     ) -> InterpResult<'tcx> {
-        assert!(dest.layout.is_sized(), "Cannot write unsized data");
+        assert!(dest.layout.is_sized(), "Cannot write unsized immediate data");
         trace!("write_immediate: {:?} <- {:?}: {}", *dest, src, dest.layout.ty);
 
         // See if we can avoid an allocation. This is the counterpart to `read_immediate_raw`,
@@ -591,9 +591,34 @@ where
                     *self.force_allocation(dest)?
                 } else {
                     match M::access_local_mut(self, frame, local)? {
-                        Operand::Immediate(local) => {
+                        Operand::Immediate(local_val) => {
                             // Local can be updated in-place.
-                            *local = src;
+                            *local_val = src;
+                            // Double-check that the value we are storing and the local fit to each other.
+                            // (*After* doing the update for borrow checker reasons.)
+                            if cfg!(debug_assertions) {
+                                let local_layout =
+                                    self.layout_of_local(&self.stack()[frame], local, None)?;
+                                match (src, local_layout.abi) {
+                                    (Immediate::Scalar(scalar), Abi::Scalar(s)) => {
+                                        assert_eq!(scalar.size(), s.size(self))
+                                    }
+                                    (
+                                        Immediate::ScalarPair(a_val, b_val),
+                                        Abi::ScalarPair(a, b),
+                                    ) => {
+                                        assert_eq!(a_val.size(), a.size(self));
+                                        assert_eq!(b_val.size(), b.size(self));
+                                    }
+                                    (Immediate::Uninit, _) => {}
+                                    (src, abi) => {
+                                        bug!(
+                                            "value {src:?} cannot be written into local with type {} (ABI {abi:?})",
+                                            local_layout.ty
+                                        )
+                                    }
+                                };
+                            }
                             return Ok(());
                         }
                         Operand::Indirect(mplace) => {
diff --git a/compiler/rustc_middle/src/mir/interpret/allocation.rs b/compiler/rustc_middle/src/mir/interpret/allocation.rs
index b8030d9db13..c1cb2f2e497 100644
--- a/compiler/rustc_middle/src/mir/interpret/allocation.rs
+++ b/compiler/rustc_middle/src/mir/interpret/allocation.rs
@@ -571,7 +571,7 @@ impl<Prov: Provenance, Extra, Bytes: AllocBytes> Allocation<Prov, Extra, Bytes>
         assert!(self.mutability == Mutability::Mut);
 
         // `to_bits_or_ptr_internal` is the right method because we just want to store this data
-        // as-is into memory.
+        // as-is into memory. This also double-checks that `val.size()` matches `range.size`.
         let (bytes, provenance) = match val.to_bits_or_ptr_internal(range.size)? {
             Right(ptr) => {
                 let (provenance, offset) = ptr.into_parts();
diff --git a/compiler/rustc_middle/src/mir/interpret/value.rs b/compiler/rustc_middle/src/mir/interpret/value.rs
index 0416411dfe1..47421d0f037 100644
--- a/compiler/rustc_middle/src/mir/interpret/value.rs
+++ b/compiler/rustc_middle/src/mir/interpret/value.rs
@@ -320,6 +320,14 @@ impl<Prov> Scalar<Prov> {
             }
         })
     }
+
+    #[inline]
+    pub fn size(self) -> Size {
+        match self {
+            Scalar::Int(int) => int.size(),
+            Scalar::Ptr(_ptr, sz) => Size::from_bytes(sz),
+        }
+    }
 }
 
 impl<'tcx, Prov: Provenance> Scalar<Prov> {