safe transmute: forbid reference lifetime extension

Modifies `BikeshedIntrinsicFrom` to forbid lifetime extensions on
references. This static check can be opted out of with the
`Assume::lifetimes` flag.

Fixes #129097

1 files changed, 88 insertions, 53 deletions

diff --git a/compiler/rustc_transmute/src/layout/tree.rs b/compiler/rustc_transmute/src/layout/tree.rs
index 5c25f913ffe..7c73f74e629 100644
--- a/compiler/rustc_transmute/src/layout/tree.rs
+++ b/compiler/rustc_transmute/src/layout/tree.rs
@@ -171,10 +171,12 @@ where
 
 #[cfg(feature = "rustc")]
 pub(crate) mod rustc {
-    use rustc_middle::ty::layout::{HasTyCtxt, LayoutCx, LayoutError, LayoutOf};
+    use rustc_middle::ty::layout::{HasTyCtxt, LayoutCx, LayoutError};
     use rustc_middle::ty::{self, AdtDef, AdtKind, List, ScalarInt, Ty, TyCtxt, TypeVisitableExt};
     use rustc_span::ErrorGuaranteed;
-    use rustc_target::abi::{FieldsShape, Size, TyAndLayout, Variants};
+    use rustc_target::abi::{
+        FieldIdx, FieldsShape, Layout, Size, TyAndLayout, VariantIdx, Variants,
+    };
 
     use super::Tree;
     use crate::layout::rustc::{Def, Ref};
@@ -202,20 +204,18 @@ pub(crate) mod rustc {
     }
 
     impl<'tcx> Tree<Def<'tcx>, Ref<'tcx>> {
-        pub fn from_ty(
-            ty_and_layout: TyAndLayout<'tcx, Ty<'tcx>>,
-            cx: LayoutCx<'tcx, TyCtxt<'tcx>>,
-        ) -> Result<Self, Err> {
+        pub fn from_ty(ty: Ty<'tcx>, cx: LayoutCx<'tcx, TyCtxt<'tcx>>) -> Result<Self, Err> {
             use rustc_target::abi::HasDataLayout;
+            let layout = ty_layout(cx, ty);
 
-            if let Err(e) = ty_and_layout.ty.error_reported() {
+            if let Err(e) = ty.error_reported() {
                 return Err(Err::TypeError(e));
             }
 
             let target = cx.tcx.data_layout();
             let pointer_size = target.pointer_size;
 
-            match ty_and_layout.ty.kind() {
+            match ty.kind() {
                 ty::Bool => Ok(Self::bool()),
 
                 ty::Float(nty) => {
@@ -233,32 +233,30 @@ pub(crate) mod rustc {
                     Ok(Self::number(width as _))
                 }
 
-                ty::Tuple(members) => Self::from_tuple(ty_and_layout, members, cx),
+                ty::Tuple(members) => Self::from_tuple((ty, layout), members, cx),
 
                 ty::Array(inner_ty, len) => {
-                    let FieldsShape::Array { stride, count } = &ty_and_layout.fields else {
+                    let FieldsShape::Array { stride, count } = &layout.fields else {
                         return Err(Err::NotYetSupported);
                     };
-                    let inner_ty_and_layout = cx.layout_of(*inner_ty)?;
-                    assert_eq!(*stride, inner_ty_and_layout.size);
-                    let elt = Tree::from_ty(inner_ty_and_layout, cx)?;
+                    let inner_layout = ty_layout(cx, *inner_ty);
+                    assert_eq!(*stride, inner_layout.size);
+                    let elt = Tree::from_ty(*inner_ty, cx)?;
                     Ok(std::iter::repeat(elt)
                         .take(*count as usize)
                         .fold(Tree::unit(), |tree, elt| tree.then(elt)))
                 }
 
-                ty::Adt(adt_def, _args_ref) if !ty_and_layout.ty.is_box() => {
-                    match adt_def.adt_kind() {
-                        AdtKind::Struct => Self::from_struct(ty_and_layout, *adt_def, cx),
-                        AdtKind::Enum => Self::from_enum(ty_and_layout, *adt_def, cx),
-                        AdtKind::Union => Self::from_union(ty_and_layout, *adt_def, cx),
-                    }
-                }
+                ty::Adt(adt_def, _args_ref) if !ty.is_box() => match adt_def.adt_kind() {
+                    AdtKind::Struct => Self::from_struct((ty, layout), *adt_def, cx),
+                    AdtKind::Enum => Self::from_enum((ty, layout), *adt_def, cx),
+                    AdtKind::Union => Self::from_union((ty, layout), *adt_def, cx),
+                },
 
                 ty::Ref(lifetime, ty, mutability) => {
-                    let ty_and_layout = cx.layout_of(*ty)?;
-                    let align = ty_and_layout.align.abi.bytes_usize();
-                    let size = ty_and_layout.size.bytes_usize();
+                    let layout = ty_layout(cx, *ty);
+                    let align = layout.align.abi.bytes_usize();
+                    let size = layout.size.bytes_usize();
                     Ok(Tree::Ref(Ref {
                         lifetime: *lifetime,
                         ty: *ty,
@@ -274,21 +272,20 @@ pub(crate) mod rustc {
 
         /// Constructs a `Tree` from a tuple.
         fn from_tuple(
-            ty_and_layout: TyAndLayout<'tcx, Ty<'tcx>>,
+            (ty, layout): (Ty<'tcx>, Layout<'tcx>),
             members: &'tcx List<Ty<'tcx>>,
             cx: LayoutCx<'tcx, TyCtxt<'tcx>>,
         ) -> Result<Self, Err> {
-            match &ty_and_layout.fields {
+            match &layout.fields {
                 FieldsShape::Primitive => {
                     assert_eq!(members.len(), 1);
                     let inner_ty = members[0];
-                    let inner_ty_and_layout = cx.layout_of(inner_ty)?;
-                    assert_eq!(ty_and_layout.layout, inner_ty_and_layout.layout);
-                    Self::from_ty(inner_ty_and_layout, cx)
+                    let inner_layout = ty_layout(cx, inner_ty);
+                    Self::from_ty(inner_ty, cx)
                 }
                 FieldsShape::Arbitrary { offsets, .. } => {
                     assert_eq!(offsets.len(), members.len());
-                    Self::from_variant(Def::Primitive, None, ty_and_layout, ty_and_layout.size, cx)
+                    Self::from_variant(Def::Primitive, None, (ty, layout), layout.size, cx)
                 }
                 FieldsShape::Array { .. } | FieldsShape::Union(_) => Err(Err::NotYetSupported),
             }
@@ -300,13 +297,13 @@ pub(crate) mod rustc {
         ///
         /// Panics if `def` is not a struct definition.
         fn from_struct(
-            ty_and_layout: TyAndLayout<'tcx, Ty<'tcx>>,
+            (ty, layout): (Ty<'tcx>, Layout<'tcx>),
             def: AdtDef<'tcx>,
             cx: LayoutCx<'tcx, TyCtxt<'tcx>>,
         ) -> Result<Self, Err> {
             assert!(def.is_struct());
             let def = Def::Adt(def);
-            Self::from_variant(def, None, ty_and_layout, ty_and_layout.size, cx)
+            Self::from_variant(def, None, (ty, layout), layout.size, cx)
         }
 
         /// Constructs a `Tree` from an enum.
@@ -315,19 +312,18 @@ pub(crate) mod rustc {
         ///
         /// Panics if `def` is not an enum definition.
         fn from_enum(
-            ty_and_layout: TyAndLayout<'tcx, Ty<'tcx>>,
+            (ty, layout): (Ty<'tcx>, Layout<'tcx>),
             def: AdtDef<'tcx>,
             cx: LayoutCx<'tcx, TyCtxt<'tcx>>,
         ) -> Result<Self, Err> {
             assert!(def.is_enum());
-            let layout = ty_and_layout.layout;
 
             // Computes the variant of a given index.
             let layout_of_variant = |index| {
-                let tag = cx.tcx.tag_for_variant((ty_and_layout.ty, index));
+                let tag = cx.tcx.tag_for_variant((cx.tcx.erase_regions(ty), index));
                 let variant_def = Def::Variant(def.variant(index));
-                let variant_ty_and_layout = ty_and_layout.for_variant(&cx, index);
-                Self::from_variant(variant_def, tag, variant_ty_and_layout, layout.size, cx)
+                let variant_layout = ty_variant(cx, (ty, layout), index);
+                Self::from_variant(variant_def, tag, (ty, variant_layout), layout.size, cx)
             };
 
             // We consider three kinds of enums, each demanding a different
@@ -385,21 +381,20 @@ pub(crate) mod rustc {
         fn from_variant(
             def: Def<'tcx>,
             tag: Option<ScalarInt>,
-            ty_and_layout: TyAndLayout<'tcx, Ty<'tcx>>,
+            (ty, layout): (Ty<'tcx>, Layout<'tcx>),
             total_size: Size,
             cx: LayoutCx<'tcx, TyCtxt<'tcx>>,
         ) -> Result<Self, Err> {
             // This constructor does not support non-`FieldsShape::Arbitrary`
             // layouts.
-            let FieldsShape::Arbitrary { offsets, memory_index } = ty_and_layout.layout.fields()
-            else {
+            let FieldsShape::Arbitrary { offsets, memory_index } = layout.fields() else {
                 return Err(Err::NotYetSupported);
             };
 
             // When this function is invoked with enum variants,
             // `ty_and_layout.size` does not encompass the entire size of the
             // enum. We rely on `total_size` for this.
-            assert!(ty_and_layout.size <= total_size);
+            assert!(layout.size <= total_size);
 
             let mut size = Size::ZERO;
             let mut struct_tree = Self::def(def);
@@ -412,17 +407,18 @@ pub(crate) mod rustc {
 
             // Append the fields, in memory order, to the layout.
             let inverse_memory_index = memory_index.invert_bijective_mapping();
-            for (memory_idx, field_idx) in inverse_memory_index.iter_enumerated() {
+            for (memory_idx, &field_idx) in inverse_memory_index.iter_enumerated() {
                 // Add interfield padding.
-                let padding_needed = offsets[*field_idx] - size;
+                let padding_needed = offsets[field_idx] - size;
                 let padding = Self::padding(padding_needed.bytes_usize());
 
-                let field_ty_and_layout = ty_and_layout.field(&cx, field_idx.as_usize());
-                let field_tree = Self::from_ty(field_ty_and_layout, cx)?;
+                let field_ty = ty_field(cx, (ty, layout), field_idx);
+                let field_layout = ty_layout(cx, field_ty);
+                let field_tree = Self::from_ty(field_ty, cx)?;
 
                 struct_tree = struct_tree.then(padding).then(field_tree);
 
-                size += padding_needed + field_ty_and_layout.size;
+                size += padding_needed + field_layout.size;
             }
 
             // Add trailing padding.
@@ -457,28 +453,27 @@ pub(crate) mod rustc {
         ///
         /// Panics if `def` is not a union definition.
         fn from_union(
-            ty_and_layout: TyAndLayout<'tcx, Ty<'tcx>>,
+            (ty, layout): (Ty<'tcx>, Layout<'tcx>),
             def: AdtDef<'tcx>,
             cx: LayoutCx<'tcx, TyCtxt<'tcx>>,
         ) -> Result<Self, Err> {
             assert!(def.is_union());
 
-            let union_layout = ty_and_layout.layout;
-
             // This constructor does not support non-`FieldsShape::Union`
             // layouts. Fields of this shape are all placed at offset 0.
-            let FieldsShape::Union(fields) = union_layout.fields() else {
+            let FieldsShape::Union(fields) = layout.fields() else {
                 return Err(Err::NotYetSupported);
             };
 
             let fields = &def.non_enum_variant().fields;
             let fields = fields.iter_enumerated().try_fold(
                 Self::uninhabited(),
-                |fields, (idx, ref field_def)| {
+                |fields, (idx, field_def)| {
                     let field_def = Def::Field(field_def);
-                    let field_ty_and_layout = ty_and_layout.field(&cx, idx.as_usize());
-                    let field = Self::from_ty(field_ty_and_layout, cx)?;
-                    let trailing_padding_needed = union_layout.size - field_ty_and_layout.size;
+                    let field_ty = ty_field(cx, (ty, layout), idx);
+                    let field_layout = ty_layout(cx, field_ty);
+                    let field = Self::from_ty(field_ty, cx)?;
+                    let trailing_padding_needed = layout.size - field_layout.size;
                     let trailing_padding = Self::padding(trailing_padding_needed.bytes_usize());
                     let field_and_padding = field.then(trailing_padding);
                     Result::<Self, Err>::Ok(fields.or(field_and_padding))
@@ -488,4 +483,44 @@ pub(crate) mod rustc {
             Ok(Self::def(Def::Adt(def)).then(fields))
         }
     }
+
+    pub(crate) fn ty_layout<'tcx>(cx: LayoutCx<'tcx, TyCtxt<'tcx>>, ty: Ty<'tcx>) -> Layout<'tcx> {
+        crate::layout::rustc::layout_of(cx, ty).unwrap()
+    }
+
+    fn ty_field<'tcx>(
+        cx: LayoutCx<'tcx, TyCtxt<'tcx>>,
+        (ty, layout): (Ty<'tcx>, Layout<'tcx>),
+        i: FieldIdx,
+    ) -> Ty<'tcx> {
+        match ty.kind() {
+            ty::Adt(def, args) => {
+                match layout.variants {
+                    Variants::Single { index } => {
+                        let field = &def.variant(index).fields[i];
+                        field.ty(cx.tcx, args)
+                    }
+                    // Discriminant field for enums (where applicable).
+                    Variants::Multiple { tag, .. } => {
+                        assert_eq!(i.as_usize(), 0);
+                        ty::layout::PrimitiveExt::to_ty(&tag.primitive(), cx.tcx)
+                    }
+                }
+            }
+            ty::Tuple(fields) => fields[i.as_usize()],
+            kind @ _ => unimplemented!(
+                "only a subset of `Ty::ty_and_layout_field`'s functionality is implemented. implementation needed for {:?}",
+                kind
+            ),
+        }
+    }
+
+    fn ty_variant<'tcx>(
+        cx: LayoutCx<'tcx, TyCtxt<'tcx>>,
+        (ty, layout): (Ty<'tcx>, Layout<'tcx>),
+        i: VariantIdx,
+    ) -> Layout<'tcx> {
+        let ty = cx.tcx.erase_regions(ty);
+        TyAndLayout { ty, layout }.for_variant(&cx, i).layout
+    }
 }