Auto merge of #115764 - RalfJung:const-by-ref-alloc-id, r=oli-obk

some ConstValue refactoring

In particular, 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`

Fixes #105536

Author: bors

compiler/rustc_codegen_cranelift/src/constant.rs

@@ -3,7 +3,7 @@
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir::interpret::{
-    read_target_uint, AllocId, ConstAllocation, ConstValue, ErrorHandled, GlobalAlloc, Scalar,
+    read_target_uint, AllocId, ConstValue, ErrorHandled, GlobalAlloc, Scalar,
 };
 
 use cranelift_module::*;
@@ -116,7 +116,7 @@ pub(crate) fn codegen_const_value<'tcx>(
     }
 
     match const_val {
-        ConstValue::ZeroSized => unreachable!(), // we already handles ZST above
+        ConstValue::ZeroSized => unreachable!(), // we already handled ZST above
         ConstValue::Scalar(x) => match x {
             Scalar::Int(int) => {
                 if fx.clif_type(layout.ty).is_some() {
@@ -200,13 +200,14 @@ pub(crate) fn codegen_const_value<'tcx>(
                 CValue::by_val(val, layout)
             }
         },
-        ConstValue::ByRef { alloc, offset } => CValue::by_ref(
-            pointer_for_allocation(fx, alloc)
+        ConstValue::Indirect { alloc_id, offset } => CValue::by_ref(
+            pointer_for_allocation(fx, alloc_id)
                 .offset_i64(fx, i64::try_from(offset.bytes()).unwrap()),
             layout,
         ),
         ConstValue::Slice { data, start, end } => {
-            let ptr = pointer_for_allocation(fx, data)
+            let alloc_id = fx.tcx.reserve_and_set_memory_alloc(data);
+            let ptr = pointer_for_allocation(fx, alloc_id)
                 .offset_i64(fx, i64::try_from(start).unwrap())
                 .get_addr(fx);
             let len = fx
@@ -220,9 +221,9 @@ pub(crate) fn codegen_const_value<'tcx>(
 
 fn pointer_for_allocation<'tcx>(
     fx: &mut FunctionCx<'_, '_, 'tcx>,
-    alloc: ConstAllocation<'tcx>,
+    alloc_id: AllocId,
 ) -> crate::pointer::Pointer {
-    let alloc_id = fx.tcx.create_memory_alloc(alloc);
+    let alloc = fx.tcx.global_alloc(alloc_id).unwrap_memory();
     let data_id = data_id_for_alloc_id(
         &mut fx.constants_cx,
         &mut *fx.module,
@@ -353,6 +354,7 @@ fn define_all_allocs(tcx: TyCtxt<'_>, module: &mut dyn Module, cx: &mut Constant
                         unreachable!()
                     }
                 };
+                // FIXME: should we have a cache so we don't do this multiple times for the same `ConstAllocation`?
                 let data_id = *cx.anon_allocs.entry(alloc_id).or_insert_with(|| {
                     module.declare_anonymous_data(alloc.inner().mutability.is_mut(), false).unwrap()
                 });

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::Indirect { 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]) */
                         );

compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -105,7 +105,10 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
                     bug!("from_const: invalid ScalarPair layout: {:#?}", layout);
                 };
                 let a = Scalar::from_pointer(
-                    Pointer::new(bx.tcx().create_memory_alloc(data), Size::from_bytes(start)),
+                    Pointer::new(
+                        bx.tcx().reserve_and_set_memory_alloc(data),
+                        Size::from_bytes(start),
+                    ),
                     &bx.tcx(),
                 );
                 let a_llval = bx.scalar_to_backend(
@@ -116,7 +119,8 @@ 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::Indirect { alloc_id, offset } => {
+                let alloc = bx.tcx().global_alloc(alloc_id).unwrap_memory();
                 return Self::from_const_alloc(bx, layout, alloc, offset);
             }
         };
@@ -182,6 +186,8 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
             _ if layout.is_zst() => OperandRef::zero_sized(layout),
             _ => {
                 // Neither a scalar nor scalar pair. Load from a place
+                // FIXME: should we cache `const_data_from_alloc` to avoid repeating this for the
+                // same `ConstAllocation`?
                 let init = bx.const_data_from_alloc(alloc);
                 let base_addr = bx.static_addr_of(init, alloc_align, None);
 

compiler/rustc_const_eval/src/const_eval/eval_queries.rs

@@ -18,9 +18,9 @@ use super::{CanAccessStatics, CompileTimeEvalContext, CompileTimeInterpreter};
 use crate::errors;
 use crate::interpret::eval_nullary_intrinsic;
 use crate::interpret::{
-    intern_const_alloc_recursive, Allocation, ConstAlloc, ConstValue, CtfeValidationMode, GlobalId,
-    Immediate, InternKind, InterpCx, InterpError, InterpResult, MPlaceTy, MemoryKind, OpTy,
-    RefTracking, StackPopCleanup,
+    intern_const_alloc_recursive, ConstAlloc, ConstValue, CtfeValidationMode, GlobalId, Immediate,
+    InternKind, InterpCx, InterpError, InterpResult, MPlaceTy, MemoryKind, OpTy, RefTracking,
+    StackPopCleanup,
 };
 
 // Returns a pointer to where the result lives
@@ -105,91 +105,68 @@ pub(super) fn mk_eval_cx<'mir, 'tcx>(
     )
 }
 
-/// This function converts an interpreter value into a constant that is meant for use in the
-/// type system.
+/// This function converts an interpreter value into a MIR constant.
 #[instrument(skip(ecx), level = "debug")]
 pub(super) fn op_to_const<'tcx>(
     ecx: &CompileTimeEvalContext<'_, 'tcx>,
     op: &OpTy<'tcx>,
 ) -> ConstValue<'tcx> {
-    // We do not have value optimizations for everything.
-    // Only scalars and slices, since they are very common.
-    // Note that further down we turn scalars of uninitialized bits back to `ByRef`. These can result
-    // from scalar unions that are initialized with one of their zero sized variants. We could
-    // instead allow `ConstValue::Scalar` to store `ScalarMaybeUninit`, but that would affect all
-    // the usual cases of extracting e.g. a `usize`, without there being a real use case for the
-    // `Undef` situation.
-    let try_as_immediate = match op.layout.abi {
+    // Handle ZST consistently and early.
+    if op.layout.is_zst() {
+        return ConstValue::ZeroSized;
+    }
+
+    // All scalar types should be stored as `ConstValue::Scalar`. This is needed to make
+    // `ConstValue::try_to_scalar` efficient; we want that to work for *all* constants of scalar
+    // type (it's used throughout the compiler and having it work just on literals is not enough)
+    // and we want it to be fast (i.e., don't go to an `Allocation` and reconstruct the `Scalar`
+    // from its byte-serialized form).
+    let force_as_immediate = match op.layout.abi {
         Abi::Scalar(abi::Scalar::Initialized { .. }) => true,
-        Abi::ScalarPair(..) => match op.layout.ty.kind() {
-            ty::Ref(_, inner, _) => match *inner.kind() {
-                ty::Slice(elem) => elem == ecx.tcx.types.u8,
-                ty::Str => true,
-                _ => false,
-            },
-            _ => false,
-        },
+        // We don't *force* `ConstValue::Slice` for `ScalarPair`. This has the advantage that if the
+        // input `op` is a place, then turning it into a `ConstValue` and back into a `OpTy` will
+        // not have to generate any duplicate allocations (we preserve the original `AllocId` in
+        // `ConstValue::Indirect`). It means accessing the contents of a slice can be slow (since
+        // they can be stored as `ConstValue::Indirect`), but that's not relevant since we barely
+        // ever have to do this. (`try_get_slice_bytes_for_diagnostics` exists to provide this
+        // functionality.)
         _ => false,
     };
-    let immediate = if try_as_immediate {
+    let immediate = if force_as_immediate {
         Right(ecx.read_immediate(op).expect("normalization works on validated constants"))
     } else {
-        // It is guaranteed that any non-slice scalar pair is actually ByRef here.
-        // When we come back from raw const eval, we are always by-ref. The only way our op here is
-        // by-val is if we are in destructure_mir_constant, i.e., if this is (a field of) something that we
-        // "tried to make immediate" before. We wouldn't do that for non-slice scalar pairs or
-        // structs containing such.
         op.as_mplace_or_imm()
     };
 
     debug!(?immediate);
 
-    // We know `offset` is relative to the allocation, so we can use `into_parts`.
-    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 }
-            }
-            (None, offset) => {
-                assert!(mplace.layout.is_zst());
-                assert_eq!(
-                    offset.bytes() % mplace.layout.align.abi.bytes(),
-                    0,
-                    "this MPlaceTy must come from a validated constant, thus we can assume the \
-                alignment is correct",
-                );
-                ConstValue::ZeroSized
-            }
-        }
-    };
     match immediate {
-        Left(ref mplace) => to_const_value(mplace),
-        // see comment on `let try_as_immediate` above
+        Left(ref mplace) => {
+            // We know `offset` is relative to the allocation, so we can use `into_parts`.
+            let (alloc_id, offset) = mplace.ptr().into_parts();
+            let alloc_id = alloc_id.expect("cannot have `fake` place fot non-ZST type");
+            ConstValue::Indirect { alloc_id, offset }
+        }
+        // see comment on `let force_as_immediate` above
         Right(imm) => match *imm {
-            _ if imm.layout.is_zst() => ConstValue::ZeroSized,
             Immediate::Scalar(x) => ConstValue::Scalar(x),
             Immediate::ScalarPair(a, b) => {
                 debug!("ScalarPair(a: {:?}, b: {:?})", a, b);
+                // FIXME: assert that this has an appropriate type.
+                // Currently we actually get here for non-[u8] slices during valtree construction!
+                let msg = "`op_to_const` on an immediate scalar pair must only be used on slice references to actually allocated memory";
                 // We know `offset` is relative to the allocation, so we can use `into_parts`.
-                let (data, start) = match a.to_pointer(ecx).unwrap().into_parts() {
-                    (Some(alloc_id), offset) => {
-                        (ecx.tcx.global_alloc(alloc_id).unwrap_memory(), offset.bytes())
-                    }
-                    (None, _offset) => (
-                        ecx.tcx.mk_const_alloc(Allocation::from_bytes_byte_aligned_immutable(
-                            b"" as &[u8],
-                        )),
-                        0,
-                    ),
-                };
-                let len = b.to_target_usize(ecx).unwrap();
-                let start = start.try_into().unwrap();
+                // We use `ConstValue::Slice` so that we don't have to generate an allocation for
+                // `ConstValue::Indirect` here.
+                let (alloc_id, offset) = a.to_pointer(ecx).expect(msg).into_parts();
+                let alloc_id = alloc_id.expect(msg);
+                let data = ecx.tcx.global_alloc(alloc_id).unwrap_memory();
+                let start = offset.bytes_usize();
+                let len = b.to_target_usize(ecx).expect(msg);
                 let len: usize = len.try_into().unwrap();
                 ConstValue::Slice { data, start, end: start + len }
             }
-            Immediate::Uninit => to_const_value(&op.assert_mem_place()),
+            Immediate::Uninit => bug!("`Uninit` is not a valid value for {}", op.layout.ty),
         },
     }
 }

compiler/rustc_const_eval/src/interpret/cast.rs

@@ -84,7 +84,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                         )
                         .ok_or_else(|| err_inval!(TooGeneric))?;
 
-                        let fn_ptr = self.create_fn_alloc_ptr(FnVal::Instance(instance));
+                        let fn_ptr = self.fn_ptr(FnVal::Instance(instance));
                         self.write_pointer(fn_ptr, dest)?;
                     }
                     _ => span_bug!(self.cur_span(), "reify fn pointer on {:?}", src.layout.ty),
@@ -116,7 +116,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                             ty::ClosureKind::FnOnce,
                         )
                         .ok_or_else(|| err_inval!(TooGeneric))?;
-                        let fn_ptr = self.create_fn_alloc_ptr(FnVal::Instance(instance));
+                        let fn_ptr = self.fn_ptr(FnVal::Instance(instance));
                         self.write_pointer(fn_ptr, dest)?;
                     }
                     _ => span_bug!(self.cur_span(), "closure fn pointer on {:?}", src.layout.ty),

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)
     }
 }

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -555,7 +555,7 @@ pub macro compile_time_machine(<$mir: lifetime, $tcx: lifetime>) {
         def_id: DefId,
     ) -> InterpResult<$tcx, Pointer> {
         // Use the `AllocId` associated with the `DefId`. Any actual *access* will fail.
-        Ok(Pointer::new(ecx.tcx.create_static_alloc(def_id), Size::ZERO))
+        Ok(Pointer::new(ecx.tcx.reserve_and_set_static_alloc(def_id), Size::ZERO))
     }
 
     #[inline(always)]

compiler/rustc_const_eval/src/interpret/memory.rs

@@ -176,12 +176,9 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         M::adjust_alloc_base_pointer(self, ptr)
     }
 
-    pub fn create_fn_alloc_ptr(
-        &mut self,
-        fn_val: FnVal<'tcx, M::ExtraFnVal>,
-    ) -> Pointer<M::Provenance> {
+    pub fn fn_ptr(&mut self, fn_val: FnVal<'tcx, M::ExtraFnVal>) -> Pointer<M::Provenance> {
         let id = match fn_val {
-            FnVal::Instance(instance) => self.tcx.create_fn_alloc(instance),
+            FnVal::Instance(instance) => self.tcx.reserve_and_set_fn_alloc(instance),
             FnVal::Other(extra) => {
                 // FIXME(RalfJung): Should we have a cache here?
                 let id = self.tcx.reserve_alloc_id();

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::Indirect { 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()),
@@ -769,7 +768,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 // We rely on mutability being set correctly in `data` to prevent writes
                 // where none should happen.
                 let ptr = Pointer::new(
-                    self.tcx.create_memory_alloc(data),
+                    self.tcx.reserve_and_set_memory_alloc(data),
                     Size::from_bytes(start), // offset: `start`
                 );
                 Operand::Immediate(Immediate::new_slice(

compiler/rustc_const_eval/src/interpret/traits.rs

@@ -27,7 +27,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         ensure_monomorphic_enough(*self.tcx, ty)?;
         ensure_monomorphic_enough(*self.tcx, poly_trait_ref)?;
 
-        let vtable_symbolic_allocation = self.tcx.create_vtable_alloc(ty, poly_trait_ref);
+        let vtable_symbolic_allocation = self.tcx.reserve_and_set_vtable_alloc(ty, poly_trait_ref);
         let vtable_ptr = self.global_base_pointer(Pointer::from(vtable_symbolic_allocation))?;
         Ok(vtable_ptr.into())
     }