Auto merge of #114520 - RalfJung:unsized-valtrees, r=oli-obk

simplify handling of valtrees for unsized types

4 files changed, 66 insertions, 121 deletions

diff --git a/compiler/rustc_const_eval/src/const_eval/valtrees.rs b/compiler/rustc_const_eval/src/const_eval/valtrees.rs
index f785bcfed6c..b15a65d67a3 100644
--- a/compiler/rustc_const_eval/src/const_eval/valtrees.rs
+++ b/compiler/rustc_const_eval/src/const_eval/valtrees.rs
@@ -7,9 +7,10 @@ use crate::interpret::{
     intern_const_alloc_recursive, ConstValue, ImmTy, Immediate, InternKind, MemPlaceMeta,
     MemoryKind, Place, Projectable, Scalar,
 };
+use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
 use rustc_middle::ty::{self, ScalarInt, Ty, TyCtxt};
 use rustc_span::source_map::DUMMY_SP;
-use rustc_target::abi::{Align, FieldIdx, VariantIdx, FIRST_VARIANT};
+use rustc_target::abi::VariantIdx;
 
 #[instrument(skip(ecx), level = "debug")]
 fn branches<'tcx>(
@@ -154,52 +155,37 @@ pub(crate) fn const_to_valtree_inner<'tcx>(
     }
 }
 
-#[instrument(skip(ecx), level = "debug")]
-fn create_mplace_from_layout<'tcx>(
-    ecx: &mut CompileTimeEvalContext<'tcx, 'tcx>,
-    ty: Ty<'tcx>,
-) -> MPlaceTy<'tcx> {
-    let tcx = ecx.tcx;
-    let param_env = ecx.param_env;
-    let layout = tcx.layout_of(param_env.and(ty)).unwrap();
-    debug!(?layout);
-
-    ecx.allocate(layout, MemoryKind::Stack).unwrap()
-}
-
-// Walks custom DSTs and gets the type of the unsized field and the number of elements
-// in the unsized field.
-fn get_info_on_unsized_field<'tcx>(
-    ty: Ty<'tcx>,
+/// Valtrees don't store the `MemPlaceMeta` that all dynamically sized values have in the interpreter.
+/// This function reconstructs it.
+fn reconstruct_place_meta<'tcx>(
+    layout: TyAndLayout<'tcx>,
     valtree: ty::ValTree<'tcx>,
     tcx: TyCtxt<'tcx>,
-) -> (Ty<'tcx>, usize) {
+) -> MemPlaceMeta {
+    if layout.is_sized() {
+        return MemPlaceMeta::None;
+    }
+
     let mut last_valtree = valtree;
+    // Traverse the type, and update `last_valtree` as we go.
     let tail = tcx.struct_tail_with_normalize(
-        ty,
+        layout.ty,
         |ty| ty,
         || {
             let branches = last_valtree.unwrap_branch();
-            last_valtree = branches[branches.len() - 1];
+            last_valtree = *branches.last().unwrap();
             debug!(?branches, ?last_valtree);
         },
     );
-    let unsized_inner_ty = match tail.kind() {
-        ty::Slice(t) => *t,
-        ty::Str => tail,
-        _ => bug!("expected Slice or Str"),
-    };
-
-    // Have to adjust type for ty::Str
-    let unsized_inner_ty = match unsized_inner_ty.kind() {
-        ty::Str => tcx.types.u8,
-        _ => unsized_inner_ty,
+    // Sanity-check that we got a tail we support.
+    match tail.kind() {
+        ty::Slice(..) | ty::Str => {}
+        _ => bug!("unsized tail of a valtree must be Slice or Str"),
     };
 
-    // Get the number of elements in the unsized field
+    // Get the number of elements in the unsized field.
     let num_elems = last_valtree.unwrap_branch().len();
-
-    (unsized_inner_ty, num_elems)
+    MemPlaceMeta::Meta(Scalar::from_target_usize(num_elems as u64, &tcx))
 }
 
 #[instrument(skip(ecx), level = "debug", ret)]
@@ -208,41 +194,9 @@ fn create_pointee_place<'tcx>(
     ty: Ty<'tcx>,
     valtree: ty::ValTree<'tcx>,
 ) -> MPlaceTy<'tcx> {
-    let tcx = ecx.tcx.tcx;
-
-    if !ty.is_sized(*ecx.tcx, ty::ParamEnv::empty()) {
-        // We need to create `Allocation`s for custom DSTs
-
-        let (unsized_inner_ty, num_elems) = get_info_on_unsized_field(ty, valtree, tcx);
-        let unsized_inner_ty = match unsized_inner_ty.kind() {
-            ty::Str => tcx.types.u8,
-            _ => unsized_inner_ty,
-        };
-        let unsized_inner_ty_size =
-            tcx.layout_of(ty::ParamEnv::empty().and(unsized_inner_ty)).unwrap().layout.size();
-        debug!(?unsized_inner_ty, ?unsized_inner_ty_size, ?num_elems);
-
-        // for custom DSTs only the last field/element is unsized, but we need to also allocate
-        // space for the other fields/elements
-        let layout = tcx.layout_of(ty::ParamEnv::empty().and(ty)).unwrap();
-        let size_of_sized_part = layout.layout.size();
-
-        // Get the size of the memory behind the DST
-        let dst_size = unsized_inner_ty_size.checked_mul(num_elems as u64, &tcx).unwrap();
-
-        let size = size_of_sized_part.checked_add(dst_size, &tcx).unwrap();
-        let align = Align::from_bytes(size.bytes().next_power_of_two()).unwrap();
-        let ptr = ecx.allocate_ptr(size, align, MemoryKind::Stack).unwrap();
-        debug!(?ptr);
-
-        MPlaceTy::from_aligned_ptr_with_meta(
-            ptr.into(),
-            layout,
-            MemPlaceMeta::Meta(Scalar::from_target_usize(num_elems as u64, &tcx)),
-        )
-    } else {
-        create_mplace_from_layout(ecx, ty)
-    }
+    let layout = ecx.layout_of(ty).unwrap();
+    let meta = reconstruct_place_meta(layout, valtree, ecx.tcx.tcx);
+    ecx.allocate_dyn(layout, MemoryKind::Stack, meta).unwrap()
 }
 
 /// Converts a `ValTree` to a `ConstValue`, which is needed after mir
@@ -282,10 +236,13 @@ pub fn valtree_to_const_value<'tcx>(
         ty::Ref(_, _, _) | ty::Tuple(_) | ty::Array(_, _) | ty::Adt(..) => {
             let place = match ty.kind() {
                 ty::Ref(_, inner_ty, _) => {
-                    // Need to create a place for the pointee to fill for Refs
+                    // Need to create a place for the pointee (the reference itself will be an immediate)
                     create_pointee_place(&mut ecx, *inner_ty, valtree)
                 }
-                _ => create_mplace_from_layout(&mut ecx, ty),
+                _ => {
+                    // Need to create a place for this valtree.
+                    create_pointee_place(&mut ecx, ty, valtree)
+                }
             };
             debug!(?place);
 
@@ -399,45 +356,8 @@ fn valtree_into_mplace<'tcx>(
                 debug!(?i, ?inner_valtree);
 
                 let place_inner = match ty.kind() {
-                    ty::Str | ty::Slice(_) => ecx.project_index(place, i as u64).unwrap(),
-                    _ if !ty.is_sized(*ecx.tcx, ty::ParamEnv::empty())
-                        && i == branches.len() - 1 =>
-                    {
-                        // Note: For custom DSTs we need to manually process the last unsized field.
-                        // We created a `Pointer` for the `Allocation` of the complete sized version of
-                        // the Adt in `create_pointee_place` and now we fill that `Allocation` with the
-                        // values in the ValTree. For the unsized field we have to additionally add the meta
-                        // data.
-
-                        let (unsized_inner_ty, num_elems) =
-                            get_info_on_unsized_field(ty, valtree, tcx);
-                        debug!(?unsized_inner_ty);
-
-                        let inner_ty = match ty.kind() {
-                            ty::Adt(def, args) => {
-                                let i = FieldIdx::from_usize(i);
-                                def.variant(FIRST_VARIANT).fields[i].ty(tcx, args)
-                            }
-                            ty::Tuple(inner_tys) => inner_tys[i],
-                            _ => bug!("unexpected unsized type {:?}", ty),
-                        };
-
-                        let inner_layout =
-                            tcx.layout_of(ty::ParamEnv::empty().and(inner_ty)).unwrap();
-                        debug!(?inner_layout);
-
-                        let offset = place_adjusted.layout.fields.offset(i);
-                        place
-                            .offset_with_meta(
-                                offset,
-                                MemPlaceMeta::Meta(Scalar::from_target_usize(
-                                    num_elems as u64,
-                                    &tcx,
-                                )),
-                                inner_layout,
-                                &tcx,
-                            )
-                            .unwrap()
+                    ty::Str | ty::Slice(_) | ty::Array(..) => {
+                        ecx.project_index(place, i as u64).unwrap()
                     }
                     _ => ecx.project_field(&place_adjusted, i).unwrap(),
                 };
diff --git a/compiler/rustc_const_eval/src/interpret/place.rs b/compiler/rustc_const_eval/src/interpret/place.rs
index ec226808f1b..daadb758964 100644
--- a/compiler/rustc_const_eval/src/interpret/place.rs
+++ b/compiler/rustc_const_eval/src/interpret/place.rs
@@ -934,14 +934,26 @@ where
         Ok(MPlaceTy { mplace, layout: place.layout, align: place.align })
     }
 
+    pub fn allocate_dyn(
+        &mut self,
+        layout: TyAndLayout<'tcx>,
+        kind: MemoryKind<M::MemoryKind>,
+        meta: MemPlaceMeta<M::Provenance>,
+    ) -> InterpResult<'tcx, MPlaceTy<'tcx, M::Provenance>> {
+        let Some((size, align)) = self.size_and_align_of(&meta, &layout)? else {
+            span_bug!(self.cur_span(), "cannot allocate space for `extern` type, size is not known")
+        };
+        let ptr = self.allocate_ptr(size, align, kind)?;
+        Ok(MPlaceTy::from_aligned_ptr_with_meta(ptr.into(), layout, meta))
+    }
+
     pub fn allocate(
         &mut self,
         layout: TyAndLayout<'tcx>,
         kind: MemoryKind<M::MemoryKind>,
     ) -> InterpResult<'tcx, MPlaceTy<'tcx, M::Provenance>> {
         assert!(layout.is_sized());
-        let ptr = self.allocate_ptr(layout.size, layout.align.abi, kind)?;
-        Ok(MPlaceTy::from_aligned_ptr(ptr.into(), layout))
+        self.allocate_dyn(layout, kind, MemPlaceMeta::None)
     }
 
     /// Returns a wide MPlace of type `&'static [mut] str` to a new 1-aligned allocation.
diff --git a/compiler/rustc_const_eval/src/interpret/terminator.rs b/compiler/rustc_const_eval/src/interpret/terminator.rs
index 1bd21473182..3c03172bbef 100644
--- a/compiler/rustc_const_eval/src/interpret/terminator.rs
+++ b/compiler/rustc_const_eval/src/interpret/terminator.rs
@@ -393,15 +393,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             ));
             assert_eq!(dest_offset, None);
             // Allocate enough memory to hold `src`.
-            let Some((size, align)) = self.size_and_align_of_mplace(&src)? else {
-                span_bug!(
-                    self.cur_span(),
-                    "unsized fn arg with `extern` type tail should not be allowed"
-                )
-            };
-            let ptr = self.allocate_ptr(size, align, MemoryKind::Stack)?;
-            let dest_place =
-                MPlaceTy::from_aligned_ptr_with_meta(ptr.into(), callee_arg.layout, src.meta);
+            let dest_place = self.allocate_dyn(src.layout, MemoryKind::Stack, src.meta)?;
             // Update the local to be that new place.
             *M::access_local_mut(self, dest_frame, dest_local)? = Operand::Indirect(*dest_place);
         }
diff --git a/tests/ui/const-generics/slice-const-param.rs b/tests/ui/const-generics/slice-const-param.rs
index 05d21e08d74..90c573ab365 100644
--- a/tests/ui/const-generics/slice-const-param.rs
+++ b/tests/ui/const-generics/slice-const-param.rs
@@ -11,9 +11,30 @@ pub fn function_with_bytes<const BYTES: &'static [u8]>() -> &'static [u8] {
     BYTES
 }
 
+// Also check the codepaths for custom DST
+#[derive(PartialEq, Eq)]
+struct MyStr(str);
+impl std::marker::ConstParamTy for MyStr {}
+
+fn function_with_my_str<const S: &'static MyStr>() -> &'static MyStr {
+    S
+}
+
+impl MyStr {
+    const fn new(s: &'static str) -> &'static MyStr {
+        unsafe { std::mem::transmute(s) }
+    }
+
+    fn as_str(&self) -> &str {
+        &self.0
+    }
+}
+
 pub fn main() {
     assert_eq!(function_with_str::<"Rust">(), "Rust");
     assert_eq!(function_with_str::<"ℇ㇈↦">(), "ℇ㇈↦");
     assert_eq!(function_with_bytes::<b"AAAA">(), &[0x41, 0x41, 0x41, 0x41]);
     assert_eq!(function_with_bytes::<{&[0x41, 0x41, 0x41, 0x41]}>(), b"AAAA");
+
+    assert_eq!(function_with_my_str::<{ MyStr::new("hello") }>().as_str(), "hello");
 }