Auto merge of #131068 - RalfJung:immediate-offset-sanity-check, r=nnethercote

Don't use Immediate::offset to transmute pointers to integers

This applies the relatively new `assert_matches_abi` check in the `offset` operation on immediates, which makes sure that if offsets are used to alter the layout (which is possible because the field layout is arbitrarily picked by the caller), this is not done in a way that breaks the invariant of the `Immediate` type.

This leads to ICEs in a GVN mir-opt test, so the second commit fixes GVN.

Fixes https://github.com/rust-lang/rust/issues/131064.

compiler/rustc_const_eval/src/interpret/operand.rs

@@ -113,28 +113,47 @@ impl<Prov: Provenance> Immediate<Prov> {
     }
 
     /// Assert that this immediate is a valid value for the given ABI.
-    pub fn assert_matches_abi(self, abi: Abi, cx: &impl HasDataLayout) {
+    pub fn assert_matches_abi(self, abi: Abi, msg: &str, cx: &impl HasDataLayout) {
         match (self, abi) {
             (Immediate::Scalar(scalar), Abi::Scalar(s)) => {
-                assert_eq!(scalar.size(), s.size(cx));
+                assert_eq!(scalar.size(), s.size(cx), "{msg}: scalar value has wrong size");
                 if !matches!(s.primitive(), abi::Pointer(..)) {
                     // This is not a pointer, it should not carry provenance.
-                    assert!(matches!(scalar, Scalar::Int(..)));
+                    assert!(
+                        matches!(scalar, Scalar::Int(..)),
+                        "{msg}: scalar value should be an integer, but has provenance"
+                    );
                 }
             }
             (Immediate::ScalarPair(a_val, b_val), Abi::ScalarPair(a, b)) => {
-                assert_eq!(a_val.size(), a.size(cx));
+                assert_eq!(
+                    a_val.size(),
+                    a.size(cx),
+                    "{msg}: first component of scalar pair has wrong size"
+                );
                 if !matches!(a.primitive(), abi::Pointer(..)) {
-                    assert!(matches!(a_val, Scalar::Int(..)));
+                    assert!(
+                        matches!(a_val, Scalar::Int(..)),
+                        "{msg}: first component of scalar pair should be an integer, but has provenance"
+                    );
                 }
-                assert_eq!(b_val.size(), b.size(cx));
+                assert_eq!(
+                    b_val.size(),
+                    b.size(cx),
+                    "{msg}: second component of scalar pair has wrong size"
+                );
                 if !matches!(b.primitive(), abi::Pointer(..)) {
-                    assert!(matches!(b_val, Scalar::Int(..)));
+                    assert!(
+                        matches!(b_val, Scalar::Int(..)),
+                        "{msg}: second component of scalar pair should be an integer, but has provenance"
+                    );
                 }
             }
-            (Immediate::Uninit, _) => {}
+            (Immediate::Uninit, _) => {
+                assert!(abi.is_sized(), "{msg}: unsized immediates are not a thing");
+            }
             _ => {
-                bug!("value {self:?} does not match ABI {abi:?})",)
+                bug!("{msg}: value {self:?} does not match ABI {abi:?})",)
             }
         }
     }
@@ -241,6 +260,7 @@ impl<'tcx, Prov: Provenance> ImmTy<'tcx, Prov> {
 
     #[inline(always)]
     pub fn from_immediate(imm: Immediate<Prov>, layout: TyAndLayout<'tcx>) -> Self {
+        // Without a `cx` we cannot call `assert_matches_abi`.
         debug_assert!(
             match (imm, layout.abi) {
                 (Immediate::Scalar(..), Abi::Scalar(..)) => true,
@@ -261,7 +281,6 @@ impl<'tcx, Prov: Provenance> ImmTy<'tcx, Prov> {
 
     #[inline]
     pub fn from_scalar_int(s: ScalarInt, layout: TyAndLayout<'tcx>) -> Self {
-        assert_eq!(s.size(), layout.size);
         Self::from_scalar(Scalar::from(s), layout)
     }
 
@@ -334,7 +353,10 @@ impl<'tcx, Prov: Provenance> ImmTy<'tcx, Prov> {
     /// given layout.
     // Not called `offset` to avoid confusion with the trait method.
     fn offset_(&self, offset: Size, layout: TyAndLayout<'tcx>, cx: &impl HasDataLayout) -> Self {
-        debug_assert!(layout.is_sized(), "unsized immediates are not a thing");
+        // Verify that the input matches its type.
+        if cfg!(debug_assertions) {
+            self.assert_matches_abi(self.layout.abi, "invalid input to Immediate::offset", cx);
+        }
         // `ImmTy` have already been checked to be in-bounds, so we can just check directly if this
         // remains in-bounds. This cannot actually be violated since projections are type-checked
         // and bounds-checked.
@@ -368,32 +390,14 @@ impl<'tcx, Prov: Provenance> ImmTy<'tcx, Prov> {
             // the field covers the entire type
             _ if layout.size == self.layout.size => {
                 assert_eq!(offset.bytes(), 0);
-                assert!(
-                    match (self.layout.abi, layout.abi) {
-                        (Abi::Scalar(l), Abi::Scalar(r)) => l.size(cx) == r.size(cx),
-                        (Abi::ScalarPair(l1, l2), Abi::ScalarPair(r1, r2)) =>
-                            l1.size(cx) == r1.size(cx) && l2.size(cx) == r2.size(cx),
-                        _ => false,
-                    },
-                    "cannot project into {} immediate with equally-sized field {}\nouter ABI: {:#?}\nfield ABI: {:#?}",
-                    self.layout.ty,
-                    layout.ty,
-                    self.layout.abi,
-                    layout.abi,
-                );
                 **self
             }
             // extract fields from types with `ScalarPair` ABI
             (Immediate::ScalarPair(a_val, b_val), Abi::ScalarPair(a, b)) => {
-                assert_matches!(layout.abi, Abi::Scalar(..));
                 Immediate::from(if offset.bytes() == 0 {
-                    // It is "okay" to transmute from `usize` to a pointer (GVN relies on that).
-                    // So only compare the size.
-                    assert_eq!(layout.size, a.size(cx));
                     a_val
                 } else {
                     assert_eq!(offset, a.size(cx).align_to(b.align(cx).abi));
-                    assert_eq!(layout.size, b.size(cx));
                     b_val
                 })
             }
@@ -405,6 +409,8 @@ impl<'tcx, Prov: Provenance> ImmTy<'tcx, Prov> {
                 self.layout
             ),
         };
+        // Ensure the new layout matches the new value.
+        inner_val.assert_matches_abi(layout.abi, "invalid field type in Immediate::offset", cx);
 
         ImmTy::from_immediate(inner_val, layout)
     }

compiler/rustc_const_eval/src/interpret/place.rs

@@ -658,7 +658,11 @@ where
                 // Things can ge wrong in quite weird ways when this is violated.
                 // Unfortunately this is too expensive to do in release builds.
                 if cfg!(debug_assertions) {
-                    src.assert_matches_abi(local_layout.abi, self);
+                    src.assert_matches_abi(
+                        local_layout.abi,
+                        "invalid immediate for given destination place",
+                        self,
+                    );
                 }
             }
             Left(mplace) => {
@@ -679,7 +683,7 @@ where
     ) -> InterpResult<'tcx> {
         // We use the sizes from `value` below.
         // Ensure that matches the type of the place it is written to.
-        value.assert_matches_abi(layout.abi, self);
+        value.assert_matches_abi(layout.abi, "invalid immediate for given destination place", 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

compiler/rustc_const_eval/src/interpret/projection.rs

@@ -82,6 +82,10 @@ pub trait Projectable<'tcx, Prov: Provenance>: Sized + std::fmt::Debug {
         self.offset_with_meta(offset, OffsetMode::Inbounds, MemPlaceMeta::None, layout, ecx)
     }
 
+    /// This does an offset-by-zero, which is effectively a transmute. Note however that
+    /// not all transmutes are supported by all projectables -- specifically, if this is an
+    /// `OpTy` or `ImmTy`, the new layout must have almost the same ABI as the old one
+    /// (only changing the `valid_range` is allowed and turning integers into pointers).
     fn transmute<M: Machine<'tcx, Provenance = Prov>>(
         &self,
         layout: TyAndLayout<'tcx>,

compiler/rustc_mir_transform/src/gvn.rs

@@ -103,7 +103,7 @@ use rustc_middle::ty::layout::{HasParamEnv, LayoutOf};
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_span::DUMMY_SP;
 use rustc_span::def_id::DefId;
-use rustc_target::abi::{self, Abi, FIRST_VARIANT, FieldIdx, Size, VariantIdx};
+use rustc_target::abi::{self, Abi, FIRST_VARIANT, FieldIdx, Primitive, Size, VariantIdx};
 use smallvec::SmallVec;
 use tracing::{debug, instrument, trace};
 
@@ -568,13 +568,29 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
                 CastKind::Transmute => {
                     let value = self.evaluated[value].as_ref()?;
                     let to = self.ecx.layout_of(to).ok()?;
-                    // `offset` for immediates only supports scalar/scalar-pair ABIs,
-                    // so bail out if the target is not one.
+                    // `offset` for immediates generally only supports projections that match the
+                    // type of the immediate. However, as a HACK, we exploit that it can also do
+                    // limited transmutes: it only works between types with the same layout, and
+                    // cannot transmute pointers to integers.
                     if value.as_mplace_or_imm().is_right() {
-                        match (value.layout.abi, to.abi) {
-                            (Abi::Scalar(..), Abi::Scalar(..)) => {}
-                            (Abi::ScalarPair(..), Abi::ScalarPair(..)) => {}
-                            _ => return None,
+                        let can_transmute = match (value.layout.abi, to.abi) {
+                            (Abi::Scalar(s1), Abi::Scalar(s2)) => {
+                                s1.size(&self.ecx) == s2.size(&self.ecx)
+                                    && !matches!(s1.primitive(), Primitive::Pointer(..))
+                            }
+                            (Abi::ScalarPair(a1, b1), Abi::ScalarPair(a2, b2)) => {
+                                a1.size(&self.ecx) == a2.size(&self.ecx) &&
+                                b1.size(&self.ecx) == b2.size(&self.ecx) &&
+                                // The alignment of the second component determines its offset, so that also needs to match.
+                                b1.align(&self.ecx) == b2.align(&self.ecx) &&
+                                // None of the inputs may be a pointer.
+                                !matches!(a1.primitive(), Primitive::Pointer(..))
+                                    && !matches!(b1.primitive(), Primitive::Pointer(..))
+                            }
+                            _ => false,
+                        };
+                        if !can_transmute {
+                            return None;
                         }
                     }
                     value.offset(Size::ZERO, to, &self.ecx).discard_err()?

tests/coverage/closure_macro.cov-map

@@ -23,19 +23,19 @@ Number of file 0 mappings: 6
 - Code(Counter(0)) at (prev + 3, 1) to (start + 0, 2)
 
 Function name: closure_macro::main::{closure#0}
-Raw bytes (35): 0x[01, 01, 03, 01, 05, 05, 0b, 09, 00, 05, 01, 10, 1c, 03, 21, 05, 04, 11, 01, 27, 02, 03, 11, 00, 16, 00, 00, 17, 00, 1e, 07, 02, 09, 00, 0a]
+Raw bytes (35): 0x[01, 01, 03, 01, 05, 05, 0b, 09, 0d, 05, 01, 10, 1c, 03, 21, 05, 04, 11, 01, 27, 02, 03, 11, 00, 16, 0d, 00, 17, 00, 1e, 07, 02, 09, 00, 0a]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 3
 - expression 0 operands: lhs = Counter(0), rhs = Counter(1)
 - expression 1 operands: lhs = Counter(1), rhs = Expression(2, Add)
-- expression 2 operands: lhs = Counter(2), rhs = Zero
+- expression 2 operands: lhs = Counter(2), rhs = Counter(3)
 Number of file 0 mappings: 5
 - Code(Counter(0)) at (prev + 16, 28) to (start + 3, 33)
 - Code(Counter(1)) at (prev + 4, 17) to (start + 1, 39)
 - Code(Expression(0, Sub)) at (prev + 3, 17) to (start + 0, 22)
     = (c0 - c1)
-- Code(Zero) at (prev + 0, 23) to (start + 0, 30)
+- Code(Counter(3)) at (prev + 0, 23) to (start + 0, 30)
 - Code(Expression(1, Add)) at (prev + 2, 9) to (start + 0, 10)
-    = (c1 + (c2 + Zero))
+    = (c1 + (c2 + c3))
 

tests/coverage/closure_macro_async.cov-map

@@ -31,19 +31,19 @@ Number of file 0 mappings: 6
 - Code(Counter(0)) at (prev + 3, 1) to (start + 0, 2)
 
 Function name: closure_macro_async::test::{closure#0}::{closure#0}
-Raw bytes (35): 0x[01, 01, 03, 01, 05, 05, 0b, 09, 00, 05, 01, 15, 1c, 03, 21, 05, 04, 11, 01, 27, 02, 03, 11, 00, 16, 00, 00, 17, 00, 1e, 07, 02, 09, 00, 0a]
+Raw bytes (35): 0x[01, 01, 03, 01, 05, 05, 0b, 09, 0d, 05, 01, 15, 1c, 03, 21, 05, 04, 11, 01, 27, 02, 03, 11, 00, 16, 0d, 00, 17, 00, 1e, 07, 02, 09, 00, 0a]
 Number of files: 1
 - file 0 => global file 1
 Number of expressions: 3
 - expression 0 operands: lhs = Counter(0), rhs = Counter(1)
 - expression 1 operands: lhs = Counter(1), rhs = Expression(2, Add)
-- expression 2 operands: lhs = Counter(2), rhs = Zero
+- expression 2 operands: lhs = Counter(2), rhs = Counter(3)
 Number of file 0 mappings: 5
 - Code(Counter(0)) at (prev + 21, 28) to (start + 3, 33)
 - Code(Counter(1)) at (prev + 4, 17) to (start + 1, 39)
 - Code(Expression(0, Sub)) at (prev + 3, 17) to (start + 0, 22)
     = (c0 - c1)
-- Code(Zero) at (prev + 0, 23) to (start + 0, 30)
+- Code(Counter(3)) at (prev + 0, 23) to (start + 0, 30)
 - Code(Expression(1, Add)) at (prev + 2, 9) to (start + 0, 10)
-    = (c1 + (c2 + Zero))
+    = (c1 + (c2 + c3))