interpret: factor out common code for place mutation

3 files changed, 89 insertions, 81 deletions

diff --git a/compiler/rustc_const_eval/src/interpret/operand.rs b/compiler/rustc_const_eval/src/interpret/operand.rs
index d9da0b3e00a..72842fd6f1d 100644
--- a/compiler/rustc_const_eval/src/interpret/operand.rs
+++ b/compiler/rustc_const_eval/src/interpret/operand.rs
@@ -111,6 +111,32 @@ impl<Prov: Provenance> Immediate<Prov> {
             Immediate::Uninit => bug!("Got uninit where a scalar or scalar pair was expected"),
         }
     }
+
+    /// Assert that this immediate is a valid value for the given ABI.
+    pub fn assert_matches_abi(self, abi: Abi, cx: &impl HasDataLayout) {
+        match (self, abi) {
+            (Immediate::Scalar(scalar), Abi::Scalar(s)) => {
+                assert_eq!(scalar.size(), s.size(cx));
+                if !matches!(s.primitive(), abi::Pointer(..)) {
+                    assert!(matches!(scalar, Scalar::Int(..)));
+                }
+            }
+            (Immediate::ScalarPair(a_val, b_val), Abi::ScalarPair(a, b)) => {
+                assert_eq!(a_val.size(), a.size(cx));
+                if !matches!(a.primitive(), abi::Pointer(..)) {
+                    assert!(matches!(a_val, Scalar::Int(..)));
+                }
+                assert_eq!(b_val.size(), b.size(cx));
+                if !matches!(b.primitive(), abi::Pointer(..)) {
+                    assert!(matches!(b_val, Scalar::Int(..)));
+                }
+            }
+            (Immediate::Uninit, _) => {}
+            _ => {
+                bug!("value {self:?} does not match ABI {abi:?})",)
+            }
+        }
+    }
 }
 
 // ScalarPair needs a type to interpret, so we often have an immediate and a type together
diff --git a/compiler/rustc_const_eval/src/interpret/place.rs b/compiler/rustc_const_eval/src/interpret/place.rs
index 45695186ff0..551a4012ddd 100644
--- a/compiler/rustc_const_eval/src/interpret/place.rs
+++ b/compiler/rustc_const_eval/src/interpret/place.rs
@@ -180,7 +180,8 @@ pub(super) enum Place<Prov: Provenance = CtfeProvenance> {
     Ptr(MemPlace<Prov>),
 
     /// To support alloc-free locals, we are able to write directly to a local. The offset indicates
-    /// where in the local this place is located; if it is `None`, no projection has been applied.
+    /// where in the local this place is located; if it is `None`, no projection has been applied
+    /// and the type of the place is exactly the type of the local.
     /// Such projections are meaningful even if the offset is 0, since they can change layouts.
     /// (Without that optimization, we'd just always be a `MemPlace`.)
     /// `Local` places always refer to the current stack frame, so they are unstable under
@@ -557,6 +558,40 @@ where
         Ok(place)
     }
 
+    /// Given a place, returns either the underlying mplace or a reference to where the value of
+    /// this place is stored.
+    fn as_mplace_or_mutable_local(
+        &mut self,
+        place: &PlaceTy<'tcx, M::Provenance>,
+    ) -> InterpResult<
+        'tcx,
+        Either<MPlaceTy<'tcx, M::Provenance>, (&mut Immediate<M::Provenance>, TyAndLayout<'tcx>)>,
+    > {
+        Ok(match place.to_place().as_mplace_or_local() {
+            Left(mplace) => Left(mplace),
+            Right((local, offset, locals_addr, layout)) => {
+                if offset.is_some() {
+                    // This has been projected to a part of this local, or had the type changed.
+                    // FIMXE: there are cases where we could still avoid allocating an mplace.
+                    Left(place.force_mplace(self)?)
+                } else {
+                    debug_assert_eq!(locals_addr, self.frame().locals_addr());
+                    debug_assert_eq!(self.layout_of_local(self.frame(), local, None)?, layout);
+                    match self.frame_mut().locals[local].access_mut()? {
+                        Operand::Indirect(mplace) => {
+                            // The local is in memory.
+                            Left(MPlaceTy { mplace: *mplace, layout })
+                        }
+                        Operand::Immediate(local_val) => {
+                            // The local still has the optimized representation.
+                            Right((local_val, layout))
+                        }
+                    }
+                }
+            }
+        })
+    }
+
     /// Write an immediate to a place
     #[inline(always)]
     #[instrument(skip(self), level = "trace")]
@@ -608,60 +643,20 @@ where
     ) -> InterpResult<'tcx> {
         assert!(dest.layout().is_sized(), "Cannot write unsized immediate data");
 
-        // See if we can avoid an allocation. This is the counterpart to `read_immediate_raw`,
-        // but not factored as a separate function.
-        let mplace = match dest.to_place().as_mplace_or_local() {
-            Right((local, offset, locals_addr, layout)) => {
-                if offset.is_some() {
-                    // This has been projected to a part of this local. We could have complicated
-                    // logic to still keep this local as an `Operand`... but it's much easier to
-                    // just fall back to the indirect path.
-                    dest.force_mplace(self)?
-                } else {
-                    debug_assert_eq!(locals_addr, self.frame().locals_addr());
-                    match self.frame_mut().locals[local].access_mut()? {
-                        Operand::Immediate(local_val) => {
-                            // Local can be updated in-place.
-                            *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.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) => {
-                            // The local is in memory, go on below.
-                            MPlaceTy { mplace: *mplace, layout }
-                        }
-                    }
+        match self.as_mplace_or_mutable_local(&dest.to_place())? {
+            Right((local_val, local_layout)) => {
+                // Local can be updated in-place.
+                *local_val = src;
+                // Double-check that the value we are storing and the local fit to each other.
+                if cfg!(debug_assertions) {
+                    src.assert_matches_abi(local_layout.abi, self);
                 }
             }
-            Left(mplace) => mplace, // already referring to memory
-        };
-
-        // This is already in memory, write there.
-        self.write_immediate_to_mplace_no_validate(src, mplace.layout, mplace.mplace)
+            Left(mplace) => {
+                self.write_immediate_to_mplace_no_validate(src, mplace.layout, mplace.mplace)?;
+            }
+        }
+        Ok(())
     }
 
     /// Write an immediate to memory.
@@ -673,6 +668,9 @@ where
         layout: TyAndLayout<'tcx>,
         dest: MemPlace<M::Provenance>,
     ) -> InterpResult<'tcx> {
+        if cfg!(debug_assertions) {
+            value.assert_matches_abi(layout.abi, self);
+        }
         // Note that it is really important that the type here is the right one, and matches the
         // type things are read at. In case `value` is a `ScalarPair`, we don't do any magic here
         // to handle padding properly, which is only correct if we never look at this data with the
@@ -723,35 +721,18 @@ where
         &mut self,
         dest: &impl Writeable<'tcx, M::Provenance>,
     ) -> InterpResult<'tcx> {
-        let mplace = match dest.to_place().as_mplace_or_local() {
-            Left(mplace) => mplace,
-            Right((local, offset, locals_addr, layout)) => {
-                if offset.is_some() {
-                    // This has been projected to a part of this local. We could have complicated
-                    // logic to still keep this local as an `Operand`... but it's much easier to
-                    // just fall back to the indirect path.
-                    // FIXME: share the logic with `write_immediate_no_validate`.
-                    dest.force_mplace(self)?
-                } else {
-                    debug_assert_eq!(locals_addr, self.frame().locals_addr());
-                    match self.frame_mut().locals[local].access_mut()? {
-                        Operand::Immediate(local) => {
-                            *local = Immediate::Uninit;
-                            return Ok(());
-                        }
-                        Operand::Indirect(mplace) => {
-                            // The local is in memory, go on below.
-                            MPlaceTy { mplace: *mplace, layout }
-                        }
-                    }
-                }
+        match self.as_mplace_or_mutable_local(&dest.to_place())? {
+            Right((local_val, _local_layout)) => {
+                *local_val = Immediate::Uninit;
             }
-        };
-        let Some(mut alloc) = self.get_place_alloc_mut(&mplace)? else {
-            // Zero-sized access
-            return Ok(());
-        };
-        alloc.write_uninit()?;
+            Left(mplace) => {
+                let Some(mut alloc) = self.get_place_alloc_mut(&mplace)? else {
+                    // Zero-sized access
+                    return Ok(());
+                };
+                alloc.write_uninit()?;
+            }
+        }
         Ok(())
     }
 
diff --git a/compiler/rustc_const_eval/src/interpret/visitor.rs b/compiler/rustc_const_eval/src/interpret/visitor.rs
index b02f12e3c7f..08b39fe8751 100644
--- a/compiler/rustc_const_eval/src/interpret/visitor.rs
+++ b/compiler/rustc_const_eval/src/interpret/visitor.rs
@@ -82,6 +82,7 @@ pub trait ValueVisitor<'tcx, M: Machine<'tcx>>: Sized {
         self.visit_value(new_val)
     }
 
+    /// Traversal logic; should not be overloaded.
     fn walk_value(&mut self, v: &Self::V) -> InterpResult<'tcx> {
         let ty = v.layout().ty;
         trace!("walk_value: type: {ty}");