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,