Auto merge of #115764 - RalfJung:const-by-ref-alloc-id, r=<try>

use AllocId instead of Allocation in ConstValue::ByRef This helps avoid redundant AllocIds when a `ByRef` constant gets put back into the interpreter. r? `@oli-obk`
rust-lang · Sep 11, 2023 · 993d146 · 993d146
2 parents 3ebb562 + 50ad68a
commit 993d146
Show file tree

Hide file tree

Showing 15 changed files with 75 additions and 57 deletions.
diff --git a/compiler/rustc_codegen_cranelift/src/constant.rs b/compiler/rustc_codegen_cranelift/src/constant.rs
@@ -222,11 +222,15 @@ pub(crate) fn codegen_const_value<'tcx>(
                 CValue::by_val(val, layout)
             }
         },
-        ConstValue::ByRef { alloc, offset } => CValue::by_ref(
-            pointer_for_allocation(fx, alloc)
-                .offset_i64(fx, i64::try_from(offset.bytes()).unwrap()),
-            layout,
-        ),
+        ConstValue::ByRef { alloc_id, offset } => {
+            let alloc = fx.tcx.global_alloc(alloc_id).unwrap_memory();
+            // FIXME: avoid creating multiple allocations for the same AllocId?
+            CValue::by_ref(
+                pointer_for_allocation(fx, alloc)
+                    .offset_i64(fx, i64::try_from(offset.bytes()).unwrap()),
+                layout,
+            )
+        }
         ConstValue::Slice { data, start, end } => {
             let ptr = pointer_for_allocation(fx, data)
                 .offset_i64(fx, i64::try_from(start).unwrap())

diff --git a/compiler/rustc_codegen_cranelift/src/intrinsics/simd.rs b/compiler/rustc_codegen_cranelift/src/intrinsics/simd.rs
@@ -172,7 +172,8 @@ pub(super) fn codegen_simd_intrinsic_call<'tcx>(
                     .expect("simd_shuffle idx not const");
 
                 let idx_bytes = match idx_const {
-                    ConstValue::ByRef { alloc, offset } => {
+                    ConstValue::ByRef { alloc_id, offset } => {
+                        let alloc = fx.tcx.global_alloc(alloc_id).unwrap_memory();
                         let size = Size::from_bytes(
                             4 * ret_lane_count, /* size_of([u32; ret_lane_count]) */
                         );

diff --git a/compiler/rustc_codegen_ssa/src/mir/operand.rs b/compiler/rustc_codegen_ssa/src/mir/operand.rs
@@ -116,7 +116,9 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
                 let b_llval = bx.const_usize((end - start) as u64);
                 OperandValue::Pair(a_llval, b_llval)
             }
-            ConstValue::ByRef { alloc, offset } => {
+            ConstValue::ByRef { alloc_id, offset } => {
+                let alloc = bx.tcx().global_alloc(alloc_id).unwrap_memory();
+                // FIXME: should we attempt to avoid building the same AllocId multiple times?
                 return Self::from_const_alloc(bx, layout, alloc, offset);
             }
         };

diff --git a/compiler/rustc_const_eval/src/const_eval/eval_queries.rs b/compiler/rustc_const_eval/src/const_eval/eval_queries.rs
@@ -148,10 +148,7 @@ pub(super) fn op_to_const<'tcx>(
     let to_const_value = |mplace: &MPlaceTy<'_>| {
         debug!("to_const_value(mplace: {:?})", mplace);
         match mplace.ptr().into_parts() {
-            (Some(alloc_id), offset) => {
-                let alloc = ecx.tcx.global_alloc(alloc_id).unwrap_memory();
-                ConstValue::ByRef { alloc, offset }
-            }
+            (Some(alloc_id), offset) => ConstValue::ByRef { alloc_id, offset },
             (None, offset) => {
                 assert!(mplace.layout.is_zst());
                 assert_eq!(

diff --git a/compiler/rustc_const_eval/src/interpret/intern.rs b/compiler/rustc_const_eval/src/interpret/intern.rs
@@ -24,8 +24,8 @@ use rustc_middle::ty::{self, layout::TyAndLayout, Ty};
 use rustc_ast::Mutability;
 
 use super::{
-    AllocId, Allocation, ConstAllocation, InterpCx, MPlaceTy, Machine, MemoryKind, PlaceTy,
-    Projectable, ValueVisitor,
+    AllocId, Allocation, InterpCx, MPlaceTy, Machine, MemoryKind, PlaceTy, Projectable,
+    ValueVisitor,
 };
 use crate::const_eval;
 use crate::errors::{DanglingPtrInFinal, UnsupportedUntypedPointer};
@@ -455,19 +455,23 @@ impl<'mir, 'tcx: 'mir, M: super::intern::CompileTimeMachine<'mir, 'tcx, !>>
 {
     /// A helper function that allocates memory for the layout given and gives you access to mutate
     /// it. Once your own mutation code is done, the backing `Allocation` is removed from the
-    /// current `Memory` and returned.
+    /// current `Memory` and interned as read-only into the global memory.
     pub fn intern_with_temp_alloc(
         &mut self,
         layout: TyAndLayout<'tcx>,
         f: impl FnOnce(
             &mut InterpCx<'mir, 'tcx, M>,
             &PlaceTy<'tcx, M::Provenance>,
         ) -> InterpResult<'tcx, ()>,
-    ) -> InterpResult<'tcx, ConstAllocation<'tcx>> {
+    ) -> InterpResult<'tcx, AllocId> {
+        // `allocate` picks a fresh AllocId that we will associate with its data below.
         let dest = self.allocate(layout, MemoryKind::Stack)?;
         f(self, &dest.clone().into())?;
         let mut alloc = self.memory.alloc_map.remove(&dest.ptr().provenance.unwrap()).unwrap().1;
         alloc.mutability = Mutability::Not;
-        Ok(self.tcx.mk_const_alloc(alloc))
+        let alloc = self.tcx.mk_const_alloc(alloc);
+        let alloc_id = dest.ptr().provenance.unwrap(); // this was just allocated, it must have provenance
+        self.tcx.set_alloc_id_memory(alloc_id, alloc);
+        Ok(alloc_id)
     }
 }
diff --git a/compiler/rustc_const_eval/src/interpret/operand.rs b/compiler/rustc_const_eval/src/interpret/operand.rs
@@ -756,11 +756,10 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         };
         let layout = from_known_layout(self.tcx, self.param_env, layout, || self.layout_of(ty))?;
         let op = match val_val {
-            ConstValue::ByRef { alloc, offset } => {
-                let id = self.tcx.create_memory_alloc(alloc);
+            ConstValue::ByRef { alloc_id, offset } => {
                 // We rely on mutability being set correctly in that allocation to prevent writes
                 // where none should happen.
-                let ptr = self.global_base_pointer(Pointer::new(id, offset))?;
+                let ptr = self.global_base_pointer(Pointer::new(alloc_id, offset))?;
                 Operand::Indirect(MemPlace::from_ptr(ptr.into()))
             }
             ConstValue::Scalar(x) => Operand::Immediate(adjust_scalar(x)?.into()),

diff --git a/compiler/rustc_middle/src/mir/interpret/value.rs b/compiler/rustc_middle/src/mir/interpret/value.rs
@@ -43,9 +43,13 @@ pub enum ConstValue<'tcx> {
 
     /// A value not represented/representable by `Scalar` or `Slice`
     ByRef {
-        /// The backing memory of the value, may contain more memory than needed for just the value
-        /// in order to share `ConstAllocation`s between values
-        alloc: ConstAllocation<'tcx>,
+        /// The backing memory of the value. May contain more memory than needed for just the value
+        /// if this points into some other larger ConstValue.
+        ///
+        /// We use an `AllocId` here instead of a `ConstAllocation<'tcx>` to make sure that when a
+        /// raw constant (which is basically just an `AllocId`) is turned into a `ConstValue` and
+        /// back, we can preserve the original `AllocId`.
+        alloc_id: AllocId,
         /// Offset into `alloc`
         offset: Size,
     },

diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
@@ -2883,8 +2883,9 @@ fn pretty_print_const_value<'tcx>(
                     _ => {}
                 }
             }
-            (ConstValue::ByRef { alloc, offset }, ty::Array(t, n)) if *t == u8_type => {
+            (ConstValue::ByRef { alloc_id, offset }, ty::Array(t, n)) if *t == u8_type => {
                 let n = n.try_to_bits(tcx.data_layout.pointer_size).unwrap();
+                let alloc = tcx.global_alloc(alloc_id).unwrap_memory();
                 // cast is ok because we already checked for pointer size (32 or 64 bit) above
                 let range = AllocRange { start: offset, size: Size::from_bytes(n) };
                 let byte_str = alloc.inner().get_bytes_strip_provenance(&tcx, range).unwrap();

diff --git a/compiler/rustc_middle/src/mir/pretty.rs b/compiler/rustc_middle/src/mir/pretty.rs
@@ -707,8 +707,9 @@ pub fn write_allocations<'tcx>(
                 Either::Left(Either::Right(std::iter::empty()))
             }
             ConstValue::ZeroSized => Either::Left(Either::Right(std::iter::empty())),
-            ConstValue::ByRef { alloc, .. } | ConstValue::Slice { data: alloc, .. } => {
-                Either::Right(alloc_ids_from_alloc(alloc))
+            ConstValue::Slice { data, .. } => Either::Right(alloc_ids_from_alloc(data)),
+            ConstValue::ByRef { alloc_id, .. } => {
+                Either::Left(Either::Left(std::iter::once(alloc_id)))
             }
         }
     }

diff --git a/compiler/rustc_middle/src/ty/structural_impls.rs b/compiler/rustc_middle/src/ty/structural_impls.rs
@@ -502,6 +502,7 @@ TrivialTypeTraversalAndLiftImpls! {
     ::rustc_span::symbol::Ident,
     ::rustc_errors::ErrorGuaranteed,
     interpret::Scalar,
+    interpret::AllocId,
     rustc_target::abi::Size,
     ty::BoundVar,
 }

diff --git a/compiler/rustc_mir_transform/src/const_prop.rs b/compiler/rustc_mir_transform/src/const_prop.rs
@@ -570,15 +570,15 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
                 Some(self.operand_from_scalar(scalar, value.layout.ty))
             }
             Immediate::ScalarPair(l, r) if l.try_to_int().is_ok() && r.try_to_int().is_ok() => {
-                let alloc = self
+                let alloc_id = self
                     .ecx
                     .intern_with_temp_alloc(value.layout, |ecx, dest| {
                         ecx.write_immediate(*imm, dest)
                     })
                     .ok()?;
 
                 let literal = ConstantKind::Val(
-                    ConstValue::ByRef { alloc, offset: Size::ZERO },
+                    ConstValue::ByRef { alloc_id, offset: Size::ZERO },
                     value.layout.ty,
                 );
                 Some(Operand::Constant(Box::new(Constant {

diff --git a/compiler/rustc_mir_transform/src/large_enums.rs b/compiler/rustc_mir_transform/src/large_enums.rs
@@ -139,7 +139,6 @@ impl EnumSizeOpt {
 
                     let (adt_def, num_variants, alloc_id) =
                         self.candidate(tcx, param_env, ty, &mut alloc_cache)?;
-                    let alloc = tcx.global_alloc(alloc_id).unwrap_memory();
 
                     let tmp_ty = Ty::new_array(tcx, tcx.types.usize, num_variants as u64);
 
@@ -154,7 +153,7 @@ impl EnumSizeOpt {
                         span,
                         user_ty: None,
                         literal: ConstantKind::Val(
-                            interpret::ConstValue::ByRef { alloc, offset: Size::ZERO },
+                            interpret::ConstValue::ByRef { alloc_id, offset: Size::ZERO },
                             tmp_ty,
                         ),
                     };

diff --git a/compiler/rustc_monomorphize/src/collector.rs b/compiler/rustc_monomorphize/src/collector.rs
@@ -1470,8 +1470,9 @@ fn collect_const_value<'tcx>(
 ) {
     match value {
         ConstValue::Scalar(Scalar::Ptr(ptr, _size)) => collect_alloc(tcx, ptr.provenance, output),
-        ConstValue::Slice { data: alloc, start: _, end: _ } | ConstValue::ByRef { alloc, .. } => {
-            for &id in alloc.inner().provenance().ptrs().values() {
+        ConstValue::ByRef { alloc_id, .. } => collect_alloc(tcx, alloc_id, output),
+        ConstValue::Slice { data, start: _, end: _ } => {
+            for &id in data.inner().provenance().ptrs().values() {
                 collect_alloc(tcx, id, output);
             }
         }

diff --git a/compiler/rustc_smir/src/rustc_smir/alloc.rs b/compiler/rustc_smir/src/rustc_smir/alloc.rs
@@ -72,7 +72,8 @@ pub fn new_allocation<'tcx>(
                 .unwrap();
             allocation.stable(tables)
         }
-        ConstValue::ByRef { alloc, offset } => {
+        ConstValue::ByRef { alloc_id, offset } => {
+            let alloc = tables.tcx.global_alloc(alloc_id).unwrap_memory();
             let ty_size = tables
                 .tcx
                 .layout_of(rustc_middle::ty::ParamEnv::reveal_all().and(ty))

diff --git a/src/tools/clippy/clippy_utils/src/consts.rs b/src/tools/clippy/clippy_utils/src/consts.rs
@@ -684,34 +684,37 @@ pub fn miri_to_const<'tcx>(lcx: &LateContext<'tcx>, result: mir::ConstantKind<'t
             },
             _ => None,
         },
-        mir::ConstantKind::Val(ConstValue::ByRef { alloc, offset: _ }, _) => match result.ty().kind() {
-            ty::Adt(adt_def, _) if adt_def.is_struct() => Some(Constant::Adt(result)),
-            ty::Array(sub_type, len) => match sub_type.kind() {
-                ty::Float(FloatTy::F32) => match len.try_to_target_usize(lcx.tcx) {
-                    Some(len) => alloc
-                        .inner()
-                        .inspect_with_uninit_and_ptr_outside_interpreter(0..(4 * usize::try_from(len).unwrap()))
-                        .to_owned()
-                        .array_chunks::<4>()
-                        .map(|&chunk| Some(Constant::F32(f32::from_le_bytes(chunk))))
-                        .collect::<Option<Vec<Constant<'tcx>>>>()
-                        .map(Constant::Vec),
-                    _ => None,
-                },
-                ty::Float(FloatTy::F64) => match len.try_to_target_usize(lcx.tcx) {
-                    Some(len) => alloc
-                        .inner()
-                        .inspect_with_uninit_and_ptr_outside_interpreter(0..(8 * usize::try_from(len).unwrap()))
-                        .to_owned()
-                        .array_chunks::<8>()
-                        .map(|&chunk| Some(Constant::F64(f64::from_le_bytes(chunk))))
-                        .collect::<Option<Vec<Constant<'tcx>>>>()
-                        .map(Constant::Vec),
+        mir::ConstantKind::Val(ConstValue::ByRef { alloc_id, offset: _ }, _) => {
+            let alloc = lcx.tcx.global_alloc(alloc_id).unwrap_memory();
+            match result.ty().kind() {
+                ty::Adt(adt_def, _) if adt_def.is_struct() => Some(Constant::Adt(result)),
+                ty::Array(sub_type, len) => match sub_type.kind() {
+                    ty::Float(FloatTy::F32) => match len.try_to_target_usize(lcx.tcx) {
+                        Some(len) => alloc
+                            .inner()
+                            .inspect_with_uninit_and_ptr_outside_interpreter(0..(4 * usize::try_from(len).unwrap()))
+                            .to_owned()
+                            .array_chunks::<4>()
+                            .map(|&chunk| Some(Constant::F32(f32::from_le_bytes(chunk))))
+                            .collect::<Option<Vec<Constant<'tcx>>>>()
+                            .map(Constant::Vec),
+                        _ => None,
+                    },
+                    ty::Float(FloatTy::F64) => match len.try_to_target_usize(lcx.tcx) {
+                        Some(len) => alloc
+                            .inner()
+                            .inspect_with_uninit_and_ptr_outside_interpreter(0..(8 * usize::try_from(len).unwrap()))
+                            .to_owned()
+                            .array_chunks::<8>()
+                            .map(|&chunk| Some(Constant::F64(f64::from_le_bytes(chunk))))
+                            .collect::<Option<Vec<Constant<'tcx>>>>()
+                            .map(Constant::Vec),
+                        _ => None,
+                    },
                     _ => None,
                 },
                 _ => None,
-            },
-            _ => None,
+            }
         },
         _ => None,
     }