Auto merge of #99574 - durin42:allocator-patch-redux, r=nikic

codegen: use new {re,de,}allocator annotations in llvm

This obviates the patch that teaches LLVM internals about
_rust_{re,de}alloc functions by putting annotations directly in the IR
for the optimizer.

The sole test change is required to anchor FileCheck to the body of the
`box_uninitialized` method, so it doesn't see the `allocalign` on
`__rust_alloc` and get mad about the string `alloca` showing up. Since I
was there anyway, I added some checks on the attributes to prove the
right attributes got set.

r? `@nikic`

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -13,7 +13,7 @@ use smallvec::SmallVec;
 
 use crate::attributes;
 use crate::llvm::AttributePlace::Function;
-use crate::llvm::{self, Attribute, AttributeKind, AttributePlace};
+use crate::llvm::{self, AllocKindFlags, Attribute, AttributeKind, AttributePlace};
 use crate::llvm_util;
 pub use rustc_attr::{InlineAttr, InstructionSetAttr, OptimizeAttr};
 
@@ -227,6 +227,10 @@ pub(crate) fn default_optimisation_attrs<'ll>(
     attrs
 }
 
+fn create_alloc_family_attr(llcx: &llvm::Context) -> &llvm::Attribute {
+    llvm::CreateAttrStringValue(llcx, "alloc-family", "__rust_alloc")
+}
+
 /// Composite function which sets LLVM attributes for function depending on its AST (`#[attribute]`)
 /// attributes.
 pub fn from_fn_attrs<'ll, 'tcx>(
@@ -309,11 +313,54 @@ pub fn from_fn_attrs<'ll, 'tcx>(
         // Need this for AArch64.
         to_add.push(llvm::CreateAttrStringValue(cx.llcx, "branch-target-enforcement", "false"));
     }
-    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::ALLOCATOR) {
+    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::ALLOCATOR)
+        || codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::ALLOCATOR_ZEROED)
+    {
+        if llvm_util::get_version() >= (15, 0, 0) {
+            to_add.push(create_alloc_family_attr(cx.llcx));
+            // apply to argument place instead of function
+            let alloc_align = AttributeKind::AllocAlign.create_attr(cx.llcx);
+            attributes::apply_to_llfn(llfn, AttributePlace::Argument(1), &[alloc_align]);
+            to_add.push(llvm::CreateAllocSizeAttr(cx.llcx, 0));
+            let mut flags = AllocKindFlags::Alloc | AllocKindFlags::Aligned;
+            if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::ALLOCATOR) {
+                flags |= AllocKindFlags::Uninitialized;
+            } else {
+                flags |= AllocKindFlags::Zeroed;
+            }
+            to_add.push(llvm::CreateAllocKindAttr(cx.llcx, flags));
+        }
         // apply to return place instead of function (unlike all other attributes applied in this function)
         let no_alias = AttributeKind::NoAlias.create_attr(cx.llcx);
         attributes::apply_to_llfn(llfn, AttributePlace::ReturnValue, &[no_alias]);
     }
+    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::REALLOCATOR) {
+        if llvm_util::get_version() >= (15, 0, 0) {
+            to_add.push(create_alloc_family_attr(cx.llcx));
+            to_add.push(llvm::CreateAllocKindAttr(
+                cx.llcx,
+                AllocKindFlags::Realloc | AllocKindFlags::Aligned,
+            ));
+            // applies to argument place instead of function place
+            let allocated_pointer = AttributeKind::AllocatedPointer.create_attr(cx.llcx);
+            attributes::apply_to_llfn(llfn, AttributePlace::Argument(0), &[allocated_pointer]);
+            // apply to argument place instead of function
+            let alloc_align = AttributeKind::AllocAlign.create_attr(cx.llcx);
+            attributes::apply_to_llfn(llfn, AttributePlace::Argument(2), &[alloc_align]);
+            to_add.push(llvm::CreateAllocSizeAttr(cx.llcx, 3));
+        }
+        let no_alias = AttributeKind::NoAlias.create_attr(cx.llcx);
+        attributes::apply_to_llfn(llfn, AttributePlace::ReturnValue, &[no_alias]);
+    }
+    if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::DEALLOCATOR) {
+        if llvm_util::get_version() >= (15, 0, 0) {
+            to_add.push(create_alloc_family_attr(cx.llcx));
+            to_add.push(llvm::CreateAllocKindAttr(cx.llcx, AllocKindFlags::Free));
+            // applies to argument place instead of function place
+            let allocated_pointer = AttributeKind::AllocatedPointer.create_attr(cx.llcx);
+            attributes::apply_to_llfn(llfn, AttributePlace::Argument(0), &[allocated_pointer]);
+        }
+    }
     if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::CMSE_NONSECURE_ENTRY) {
         to_add.push(llvm::CreateAttrString(cx.llcx, "cmse_nonsecure_entry"));
     }

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -193,6 +193,9 @@ pub enum AttributeKind {
     SanitizeMemTag = 34,
     NoCfCheck = 35,
     ShadowCallStack = 36,
+    AllocSize = 37,
+    AllocatedPointer = 38,
+    AllocAlign = 39,
 }
 
 /// LLVMIntPredicate
@@ -986,6 +989,22 @@ pub mod debuginfo {
     }
 }
 
+use bitflags::bitflags;
+// These values **must** match with LLVMRustAllocKindFlags
+bitflags! {
+    #[repr(transparent)]
+    #[derive(Default)]
+    pub struct AllocKindFlags : u64 {
+        const Unknown = 0;
+        const Alloc = 1;
+        const Realloc = 1 << 1;
+        const Free = 1 << 2;
+        const Uninitialized = 1 << 3;
+        const Zeroed = 1 << 4;
+        const Aligned = 1 << 5;
+    }
+}
+
 extern "C" {
     pub type ModuleBuffer;
 }
@@ -1193,6 +1212,8 @@ extern "C" {
     pub fn LLVMRustCreateByValAttr<'a>(C: &'a Context, ty: &'a Type) -> &'a Attribute;
     pub fn LLVMRustCreateStructRetAttr<'a>(C: &'a Context, ty: &'a Type) -> &'a Attribute;
     pub fn LLVMRustCreateUWTableAttr(C: &Context, async_: bool) -> &Attribute;
+    pub fn LLVMRustCreateAllocSizeAttr(C: &Context, size_arg: u32) -> &Attribute;
+    pub fn LLVMRustCreateAllocKindAttr(C: &Context, size_arg: u64) -> &Attribute;
 
     // Operations on functions
     pub fn LLVMRustGetOrInsertFunction<'a>(

compiler/rustc_codegen_llvm/src/llvm/mod.rs

@@ -95,6 +95,14 @@ pub fn CreateUWTableAttr(llcx: &Context, async_: bool) -> &Attribute {
     unsafe { LLVMRustCreateUWTableAttr(llcx, async_) }
 }
 
+pub fn CreateAllocSizeAttr(llcx: &Context, size_arg: u32) -> &Attribute {
+    unsafe { LLVMRustCreateAllocSizeAttr(llcx, size_arg) }
+}
+
+pub fn CreateAllocKindAttr(llcx: &Context, kind_arg: AllocKindFlags) -> &Attribute {
+    unsafe { LLVMRustCreateAllocKindAttr(llcx, kind_arg.bits()) }
+}
+
 #[derive(Copy, Clone)]
 pub enum AttributePlace {
     ReturnValue,