Auto merge of #79594 - vn-ki:const-eval-intrinsic, r=oli-obk

add const_allocate intrinsic

r? `@oli-obk`

fixes #75390

compiler/rustc_mir/src/const_eval/machine.rs

@@ -1,5 +1,6 @@
 use rustc_middle::mir;
 use rustc_middle::ty::layout::HasTyCtxt;
+use rustc_middle::ty::InstanceDef;
 use rustc_middle::ty::{self, Ty};
 use std::borrow::Borrow;
 use std::collections::hash_map::Entry;
@@ -12,6 +13,7 @@ use rustc_hir::def_id::DefId;
 use rustc_middle::mir::AssertMessage;
 use rustc_session::Limit;
 use rustc_span::symbol::{sym, Symbol};
+use rustc_target::abi::{Align, Size};
 
 use crate::interpret::{
     self, compile_time_machine, AllocId, Allocation, Frame, GlobalId, ImmTy, InterpCx,
@@ -37,6 +39,14 @@ impl<'mir, 'tcx> InterpCx<'mir, 'tcx, CompileTimeInterpreter<'mir, 'tcx>> {
         if instance.def.requires_caller_location(self.tcx()) {
             return Ok(false);
         }
+        // Only memoize instrinsics. This was added in #79594 while adding the `const_allocate` intrinsic.
+        // We only memoize intrinsics because it would be unsound to memoize functions
+        // which might interact with the heap.
+        // Additionally, const_allocate intrinsic is impure and thus should not be memoized;
+        // it will not be memoized because it has non-ZST args
+        if !matches!(instance.def, InstanceDef::Intrinsic(_)) {
+            return Ok(false);
+        }
         // For the moment we only do this for functions which take no arguments
         // (or all arguments are ZSTs) so that we don't memoize too much.
         if args.iter().any(|a| !a.layout.is_zst()) {
@@ -295,6 +305,22 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
                 };
                 ecx.write_scalar(Scalar::from_bool(cmp), dest)?;
             }
+            sym::const_allocate => {
+                let size = ecx.read_scalar(args[0])?.to_machine_usize(ecx)?;
+                let align = ecx.read_scalar(args[1])?.to_machine_usize(ecx)?;
+
+                let align = match Align::from_bytes(align) {
+                    Ok(a) => a,
+                    Err(err) => throw_ub_format!("align has to be a power of 2, {}", err),
+                };
+
+                let ptr = ecx.memory.allocate(
+                    Size::from_bytes(size as u64),
+                    align,
+                    interpret::MemoryKind::ConstHeap,
+                );
+                ecx.write_scalar(Scalar::Ptr(ptr), dest)?;
+            }
             _ => {
                 return Err(ConstEvalErrKind::NeedsRfc(format!(
                     "calling intrinsic `{}`",

compiler/rustc_mir/src/interpret/intern.rs

@@ -104,7 +104,10 @@ fn intern_shallow<'rt, 'mir, 'tcx, M: CompileTimeMachine<'mir, 'tcx>>(
     // This match is just a canary for future changes to `MemoryKind`, which most likely need
     // changes in this function.
     match kind {
-        MemoryKind::Stack | MemoryKind::Vtable | MemoryKind::CallerLocation => {}
+        MemoryKind::Stack
+        | MemoryKind::ConstHeap
+        | MemoryKind::Vtable
+        | MemoryKind::CallerLocation => {}
     }
     // 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

compiler/rustc_mir/src/interpret/memory.rs

@@ -27,6 +27,9 @@ use crate::util::pretty;
 pub enum MemoryKind<T> {
     /// Stack memory. Error if deallocated except during a stack pop.
     Stack,
+    /// Heap memory.
+    /// FIXME: this variant should be in const_eval
+    ConstHeap,
     /// Memory backing vtables. Error if ever deallocated.
     Vtable,
     /// Memory allocated by `caller_location` intrinsic. Error if ever deallocated.
@@ -40,6 +43,7 @@ impl<T: MayLeak> MayLeak for MemoryKind<T> {
     fn may_leak(self) -> bool {
         match self {
             MemoryKind::Stack => false,
+            MemoryKind::ConstHeap => false,
             MemoryKind::Vtable => true,
             MemoryKind::CallerLocation => true,
             MemoryKind::Machine(k) => k.may_leak(),
@@ -51,6 +55,7 @@ impl<T: fmt::Display> fmt::Display for MemoryKind<T> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             MemoryKind::Stack => write!(f, "stack variable"),
+            MemoryKind::ConstHeap => write!(f, "heap allocation"),
             MemoryKind::Vtable => write!(f, "vtable"),
             MemoryKind::CallerLocation => write!(f, "caller location"),
             MemoryKind::Machine(m) => write!(f, "{}", m),

compiler/rustc_span/src/symbol.rs

@@ -356,6 +356,7 @@ symbols! {
         concat_idents,
         conservative_impl_trait,
         console,
+        const_allocate,
         const_compare_raw_pointers,
         const_constructor,
         const_eval_limit,

compiler/rustc_typeck/src/check/intrinsic.rs

@@ -286,6 +286,10 @@ pub fn check_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>) {
                 (1, vec![tcx.mk_imm_ptr(param(0)), tcx.mk_imm_ptr(param(0))], tcx.types.bool)
             }
 
+            sym::const_allocate => {
+                (0, vec![tcx.types.usize, tcx.types.usize], tcx.mk_mut_ptr(tcx.types.u8))
+            }
+
             sym::ptr_offset_from => {
                 (1, vec![tcx.mk_imm_ptr(param(0)), tcx.mk_imm_ptr(param(0))], tcx.types.isize)
             }