add const_leak and reject interning non-leaked const_allocate ptrs

Author: fee1-dead

compiler/rustc_const_eval/messages.ftl

@@ -56,6 +56,9 @@ const_eval_const_context = {$kind ->
     *[other] {""}
 }
 
+const_eval_const_heap_ptr_in_final = encountered `const_allocate` pointer in final value that was not leaked
+    .note = use `const_leak` to leak allocated pointers before returning
+
 const_eval_copy_nonoverlapping_overlapping =
     `copy_nonoverlapping` called on overlapping ranges
 
@@ -338,6 +341,7 @@ const_eval_realloc_or_alloc_with_offset =
     {$kind ->
         [dealloc] deallocating
         [realloc] reallocating
+        [leak] leaking
         *[other] {""}
     } {$ptr} which does not point to the beginning of an object
 

compiler/rustc_const_eval/src/const_eval/eval_queries.rs

@@ -93,7 +93,7 @@ fn eval_body_using_ecx<'tcx, R: InterpretationResult<'tcx>>(
     // Since evaluation had no errors, validate the resulting constant.
     const_validate_mplace(ecx, &ret, cid)?;
 
-    // Only report this after validation, as validaiton produces much better diagnostics.
+    // Only report this after validation, as validation produces much better diagnostics.
     // FIXME: ensure validation always reports this and stop making interning care about it.
 
     match intern_result {

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -170,12 +170,14 @@ pub type CompileTimeInterpCx<'tcx> = InterpCx<'tcx, CompileTimeMachine<'tcx>>;
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
 pub enum MemoryKind {
     Heap,
+    HeapLeaked,
 }
 
 impl fmt::Display for MemoryKind {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             MemoryKind::Heap => write!(f, "heap allocation"),
+            MemoryKind::HeapLeaked => write!(f, "leaked heap allocation"),
         }
     }
 }
@@ -185,6 +187,17 @@ impl interpret::MayLeak for MemoryKind {
     fn may_leak(self) -> bool {
         match self {
             MemoryKind::Heap => false,
+            MemoryKind::HeapLeaked => true,
+        }
+    }
+}
+
+impl interpret::IsConstHeap for MemoryKind {
+    #[inline(always)]
+    fn is_const_heap(&self) -> bool {
+        match self {
+            MemoryKind::Heap => true,
+            MemoryKind::HeapLeaked => false,
         }
     }
 }
@@ -197,6 +210,13 @@ impl interpret::MayLeak for ! {
     }
 }
 
+impl interpret::IsConstHeap for ! {
+    #[inline(always)]
+    fn is_const_heap(&self) -> bool {
+        *self
+    }
+}
+
 impl<'tcx> CompileTimeInterpCx<'tcx> {
     fn location_triple_for_span(&self, span: Span) -> (Symbol, u32, u32) {
         let topmost = span.ctxt().outer_expn().expansion_cause().unwrap_or(span);
@@ -457,6 +477,46 @@ impl<'tcx> interpret::Machine<'tcx> for CompileTimeMachine<'tcx> {
                     )?;
                 }
             }
+
+            sym::const_leak => {
+                let ptr = ecx.read_pointer(&args[0])?;
+                let size = ecx.read_scalar(&args[1])?.to_target_usize(ecx)?;
+                let align = ecx.read_scalar(&args[2])?.to_target_usize(ecx)?;
+
+                let size = Size::from_bytes(size);
+                let align = match Align::from_bytes(align) {
+                    Ok(a) => a,
+                    Err(err) => throw_ub_custom!(
+                        fluent::const_eval_invalid_align_details,
+                        name = "const_leak",
+                        err_kind = err.diag_ident(),
+                        align = err.align()
+                    ),
+                };
+
+                // If an allocation is created in an another const,
+                // we don't reallocate it.
+                let (alloc_id, _, _) = ecx.ptr_get_alloc_id(ptr, 0)?;
+                let is_allocated_in_another_const = matches!(
+                    ecx.tcx.try_get_global_alloc(alloc_id),
+                    Some(interpret::GlobalAlloc::Memory(_))
+                );
+
+                if is_allocated_in_another_const {
+                    // just return the pointer that was passed in
+                    ecx.write_pointer(ptr, dest)?;
+                } else {
+                    let ptr = ecx.leak_const_heap_ptr(
+                        ptr,
+                        size,
+                        align,
+                        interpret::MemoryKind::Machine(MemoryKind::Heap),
+                        interpret::MemoryKind::Machine(MemoryKind::HeapLeaked),
+                    )?;
+                    ecx.write_pointer(ptr, dest)?;
+                }
+            }
+
             // The intrinsic represents whether the value is known to the optimizer (LLVM).
             // We're not doing any optimizations here, so there is no optimizer that could know the value.
             // (We know the value here in the machine of course, but this is the runtime of that code,

compiler/rustc_const_eval/src/errors.rs

@@ -43,6 +43,14 @@ pub(crate) struct MutablePtrInFinal {
     pub kind: InternKind,
 }
 
+#[derive(Diagnostic)]
+#[diag(const_eval_const_heap_ptr_in_final)]
+#[note]
+pub(crate) struct ConstHeapPtrInFinal {
+    #[primary_span]
+    pub span: Span,
+}
+
 #[derive(Diagnostic)]
 #[diag(const_eval_unstable_in_stable_exposed)]
 pub(crate) struct UnstableInStableExposed {

compiler/rustc_const_eval/src/interpret/intern.rs

@@ -27,11 +27,12 @@ use rustc_span::def_id::LocalDefId;
 use tracing::{instrument, trace};
 
 use super::{
-    AllocId, Allocation, InterpCx, MPlaceTy, Machine, MemoryKind, PlaceTy, err_ub, interp_ok,
+    AllocId, Allocation, InterpCx, IsConstHeap, MPlaceTy, Machine, MemoryKind, PlaceTy, err_ub,
+    interp_ok,
 };
 use crate::const_eval;
 use crate::const_eval::DummyMachine;
-use crate::errors::NestedStaticInThreadLocal;
+use crate::errors::{ConstHeapPtrInFinal, NestedStaticInThreadLocal};
 
 pub trait CompileTimeMachine<'tcx, T> = Machine<
         'tcx,
@@ -62,7 +63,7 @@ impl HasStaticRootDefId for const_eval::CompileTimeMachine<'_> {
 /// already mutable (as a sanity check).
 ///
 /// Returns an iterator over all relocations referred to by this allocation.
-fn intern_shallow<'tcx, T, M: CompileTimeMachine<'tcx, T>>(
+fn intern_shallow<'tcx, T: IsConstHeap, M: CompileTimeMachine<'tcx, T>>(
     ecx: &mut InterpCx<'tcx, M>,
     alloc_id: AllocId,
     mutability: Mutability,
@@ -71,9 +72,16 @@ fn intern_shallow<'tcx, T, M: CompileTimeMachine<'tcx, T>>(
     trace!("intern_shallow {:?}", alloc_id);
     // remove allocation
     // FIXME(#120456) - is `swap_remove` correct?
-    let Some((_kind, mut alloc)) = ecx.memory.alloc_map.swap_remove(&alloc_id) else {
+    let Some((kind, mut alloc)) = ecx.memory.alloc_map.swap_remove(&alloc_id) else {
         return Err(());
     };
+
+    if matches!(kind, MemoryKind::Machine(x) if x.is_const_heap()) {
+        // emit an error but don't return an `Err` as if we did the caller assumes we found
+        // a dangling pointer.
+        ecx.tcx.dcx().emit_err(ConstHeapPtrInFinal { span: ecx.tcx.span });
+    }
+
     // Set allocation mutability as appropriate. This is used by LLVM to put things into
     // read-only memory, and also by Miri when evaluating other globals that
     // access this one.
@@ -321,7 +329,7 @@ pub fn intern_const_alloc_recursive<'tcx, M: CompileTimeMachine<'tcx, const_eval
 
 /// Intern `ret`. This function assumes that `ret` references no other allocation.
 #[instrument(level = "debug", skip(ecx))]
-pub fn intern_const_alloc_for_constprop<'tcx, T, M: CompileTimeMachine<'tcx, T>>(
+pub fn intern_const_alloc_for_constprop<'tcx, T: IsConstHeap, M: CompileTimeMachine<'tcx, T>>(
     ecx: &mut InterpCx<'tcx, M>,
     alloc_id: AllocId,
 ) -> InterpResult<'tcx, ()> {

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -46,6 +46,11 @@ pub trait MayLeak: Copy {
     fn may_leak(self) -> bool;
 }
 
+/// Whether this kind of memory is const heap (allocation that did not call `const_leak`)
+pub trait IsConstHeap {
+    fn is_const_heap(&self) -> bool;
+}
+
 /// The functionality needed by memory to manage its allocations
 pub trait AllocMap<K: Hash + Eq, V> {
     /// Tests if the map contains the given key.

compiler/rustc_const_eval/src/interpret/memory.rs

@@ -274,15 +274,14 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         M::adjust_alloc_root_pointer(self, Pointer::from(id), Some(kind))
     }
 
-    /// If this grows the allocation, `init_growth` determines
-    /// whether the additional space will be initialized.
-    pub fn reallocate_ptr(
+    fn reallocate_ptr_inner(
         &mut self,
         ptr: Pointer<Option<M::Provenance>>,
         old_size_and_align: Option<(Size, Align)>,
         new_size: Size,
         new_align: Align,
         kind: MemoryKind<M::MemoryKind>,
+        new_kind: MemoryKind<M::MemoryKind>,
         init_growth: AllocInit,
     ) -> InterpResult<'tcx, Pointer<M::Provenance>> {
         let (alloc_id, offset, _prov) = self.ptr_get_alloc_id(ptr, 0)?;
@@ -299,7 +298,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         // If requested, we zero-init the entire allocation, to ensure that a growing
         // allocation has its new bytes properly set. For the part that is copied,
         // `mem_copy` below will de-initialize things as necessary.
-        let new_ptr = self.allocate_ptr(new_size, new_align, kind, init_growth)?;
+        let new_ptr = self.allocate_ptr(new_size, new_align, new_kind, init_growth)?;
         let old_size = match old_size_and_align {
             Some((size, _align)) => size,
             None => self.get_alloc_raw(alloc_id)?.size(),
@@ -311,6 +310,47 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         interp_ok(new_ptr)
     }
 
+    /// If this grows the allocation, `init_growth` determines
+    /// whether the additional space will be initialized.
+    pub fn reallocate_ptr(
+        &mut self,
+        ptr: Pointer<Option<M::Provenance>>,
+        old_size_and_align: Option<(Size, Align)>,
+        new_size: Size,
+        new_align: Align,
+        kind: MemoryKind<M::MemoryKind>,
+        init_growth: AllocInit,
+    ) -> InterpResult<'tcx, Pointer<M::Provenance>> {
+        self.reallocate_ptr_inner(
+            ptr,
+            old_size_and_align,
+            new_size,
+            new_align,
+            kind,
+            kind,
+            init_growth,
+        )
+    }
+
+    pub fn leak_const_heap_ptr(
+        &mut self,
+        ptr: Pointer<Option<M::Provenance>>,
+        size: Size,
+        align: Align,
+        kind: MemoryKind<M::MemoryKind>,
+        new_kind: MemoryKind<M::MemoryKind>,
+    ) -> InterpResult<'tcx, Pointer<M::Provenance>> {
+        self.reallocate_ptr_inner(
+            ptr,
+            Some((size, align)),
+            size,
+            align,
+            kind,
+            new_kind,
+            AllocInit::Uninit,
+        )
+    }
+
     #[instrument(skip(self), level = "debug")]
     pub fn deallocate_ptr(
         &mut self,

compiler/rustc_const_eval/src/interpret/mod.rs

@@ -29,7 +29,9 @@ pub use self::intern::{
     HasStaticRootDefId, InternKind, InternResult, intern_const_alloc_for_constprop,
     intern_const_alloc_recursive,
 };
-pub use self::machine::{AllocMap, Machine, MayLeak, ReturnAction, compile_time_machine};
+pub use self::machine::{
+    AllocMap, IsConstHeap, Machine, MayLeak, ReturnAction, compile_time_machine,
+};
 pub use self::memory::{AllocInfo, AllocKind, AllocRef, AllocRefMut, FnVal, Memory, MemoryKind};
 use self::operand::Operand;
 pub use self::operand::{ImmTy, Immediate, OpTy};

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -415,6 +415,12 @@ pub(crate) fn check_intrinsic_type(
             vec![Ty::new_mut_ptr(tcx, tcx.types.u8), tcx.types.usize, tcx.types.usize],
             tcx.types.unit,
         ),
+        sym::const_leak => (
+            0,
+            0,
+            vec![Ty::new_mut_ptr(tcx, tcx.types.u8), tcx.types.usize, tcx.types.usize],
+            Ty::new_imm_ptr(tcx, tcx.types.u8),
+        ),
 
         sym::ptr_offset_from => (
             1,

compiler/rustc_span/src/symbol.rs

@@ -712,6 +712,7 @@ symbols! {
         const_impl_trait,
         const_in_array_repeat_expressions,
         const_indexing,
+        const_leak,
         const_let,
         const_loop,
         const_mut_refs,

library/alloc/src/boxed/thin.rs

@@ -5,7 +5,7 @@
 use core::error::Error;
 use core::fmt::{self, Debug, Display, Formatter};
 #[cfg(not(no_global_oom_handling))]
-use core::intrinsics::const_allocate;
+use core::intrinsics::{const_allocate, const_leak};
 use core::marker::PhantomData;
 #[cfg(not(no_global_oom_handling))]
 use core::marker::Unsize;
@@ -342,7 +342,11 @@ impl<H> WithHeader<H> {
                 metadata_ptr.write(ptr::metadata::<Dyn>(ptr::dangling::<T>() as *const Dyn));
 
                 // SAFETY: we have just written the metadata.
-                &*(metadata_ptr)
+                const_leak(alloc, alloc_size, alloc_align)
+                    .add(metadata_offset)
+                    .cast::<<Dyn as Pointee>::Metadata>()
+                    .as_ref()
+                    .unwrap()
             }
         };
 

library/core/src/intrinsics/mod.rs

@@ -2533,6 +2533,15 @@ pub const unsafe fn const_deallocate(_ptr: *mut u8, _size: usize, _align: usize)
     // Runtime NOP
 }
 
+#[rustc_const_unstable(feature = "const_heap", issue = "79597")]
+#[rustc_nounwind]
+#[rustc_intrinsic]
+#[miri::intrinsic_fallback_is_spec]
+pub const unsafe fn const_leak(ptr: *mut u8, _size: usize, _align: usize) -> *const u8 {
+    // const eval overrides this function.
+    ptr
+}
+
 /// Returns whether we should perform contract-checking at runtime.
 ///
 /// This is meant to be similar to the ub_checks intrinsic, in terms

tests/ui/consts/const-eval/heap/alloc_intrinsic_nontransient.rs

@@ -8,11 +8,11 @@ const FOO_RAW: *const i32 = foo();
 
 const fn foo() -> &'static i32 {
     let t = unsafe {
-        let i = intrinsics::const_allocate(4, 4) as * mut i32;
+        let i = intrinsics::const_allocate(4, 4) as *mut i32;
         *i = 20;
         i
     };
-    unsafe { &*t }
+    unsafe { &*(intrinsics::const_leak(t as *mut u8, 4, 4) as *mut i32) }
 }
 fn main() {
     assert_eq!(*FOO, 20);

tests/ui/consts/const-eval/heap/alloc_intrinsic_uninit.64bit.stderr

@@ -1,7 +1,7 @@
 error[E0080]: constructing invalid value at .<deref>: encountered uninitialized memory, but expected an integer
   --> $DIR/alloc_intrinsic_uninit.rs:7:1
    |
-LL | const BAR: &i32 = unsafe { &*(intrinsics::const_allocate(4, 4) as *mut i32) };
+LL | const BAR: &i32 = unsafe {
    | ^^^^^^^^^^^^^^^ it is undefined behavior to use this value
    |
    = note: The rules on what exactly is undefined behavior aren't clear, so this check might be overzealous. Please open an issue on the rustc repository if you believe it should not be considered undefined behavior.

tests/ui/consts/const-eval/heap/alloc_intrinsic_uninit.rs

@@ -4,6 +4,7 @@
 #![feature(const_heap)]
 use std::intrinsics;
 
-const BAR: &i32 = unsafe { &*(intrinsics::const_allocate(4, 4) as *mut i32) };
-//~^ ERROR: uninitialized memory
+const BAR: &i32 = unsafe { //~ ERROR: uninitialized memory
+    &*(intrinsics::const_leak(intrinsics::const_allocate(4, 4), 4, 4) as *mut i32)
+};
 fn main() {}

tests/ui/consts/const-eval/heap/alloc_intrinsic_untyped.rs

@@ -7,5 +7,6 @@ use std::intrinsics;
 
 const BAR: *mut i32 = unsafe { intrinsics::const_allocate(4, 4) as *mut i32 };
 //~^ error: mutable pointer in final value of constant
+//~| error: encountered `const_allocate` pointer in final value that was not leaked
 
 fn main() {}