[clang] [AllocToken] Enable alloc token instrumentation for size-returning functions (PR #168840)

[via cfe-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-clang-codegen

Author: Aleksandr Nogikh (a-nogikh)

<details>
<summary>Changes</summary>

Consider a newly added "malloc_span" attribute in the allocation token instrumentation to ensure that __size_returning_new variants are correctly identified as memory allocation functions.

Adjust the allocation token tests to verify this new behavior.

---
Full diff: https://github.com/llvm/llvm-project/pull/168840.diff


2 Files Affected:

- (modified) clang/lib/CodeGen/CGExpr.cpp (+1) 
- (modified) clang/test/CodeGenCXX/alloc-token.cpp (+8-9) 


``````````diff
diff --git a/clang/lib/CodeGen/CGExpr.cpp b/clang/lib/CodeGen/CGExpr.cpp
index f2451b16e78be..712bec62f0a68 100644
--- a/clang/lib/CodeGen/CGExpr.cpp
+++ b/clang/lib/CodeGen/CGExpr.cpp
@@ -6644,6 +6644,7 @@ RValue CodeGenFunction::EmitCall(QualType CalleeType,
 
   if (auto *CalleeDecl = dyn_cast_or_null<FunctionDecl>(TargetDecl)) {
     if (CalleeDecl->hasAttr<RestrictAttr>() ||
+        CalleeDecl->hasAttr<MallocSpanAttr>() ||
         CalleeDecl->hasAttr<AllocSizeAttr>()) {
       // Function has 'malloc' (aka. 'restrict') or 'alloc_size' attribute.
       if (SanOpts.has(SanitizerKind::AllocToken)) {
diff --git a/clang/test/CodeGenCXX/alloc-token.cpp b/clang/test/CodeGenCXX/alloc-token.cpp
index feed808a3b89b..98842206dfc00 100644
--- a/clang/test/CodeGenCXX/alloc-token.cpp
+++ b/clang/test/CodeGenCXX/alloc-token.cpp
@@ -17,10 +17,10 @@ struct __sized_ptr_t {
   size_t n;
 };
 enum class __hot_cold_t : uint8_t;
-__sized_ptr_t __size_returning_new(size_t size);
-__sized_ptr_t __size_returning_new_hot_cold(size_t, __hot_cold_t);
-__sized_ptr_t __size_returning_new_aligned(size_t, std::align_val_t);
-__sized_ptr_t __size_returning_new_aligned_hot_cold(size_t, std::align_val_t,  __hot_cold_t);
+__sized_ptr_t __size_returning_new(size_t size) __attribute__((malloc_span));
+__sized_ptr_t __size_returning_new_hot_cold(size_t, __hot_cold_t) __attribute__((malloc_span));
+__sized_ptr_t __size_returning_new_aligned(size_t, std::align_val_t) __attribute__((malloc_span));
+__sized_ptr_t __size_returning_new_aligned_hot_cold(size_t, std::align_val_t,  __hot_cold_t) __attribute__((malloc_span));
 }
 
 void *sink; // prevent optimizations from removing the calls
@@ -101,12 +101,11 @@ int *test_new_array_nothrow() {
 }
 
 // CHECK-LABEL: define dso_local void @_Z23test_size_returning_newv(
-// CHECK: call { ptr, i64 } @__size_returning_new(i64 noundef 8)
-// CHECK: call { ptr, i64 } @__size_returning_new_hot_cold(i64 noundef 8, i8 noundef zeroext 1)
-// CHECK: call { ptr, i64 } @__size_returning_new_aligned(i64 noundef 8, i64 noundef 32)
-// CHECK: call { ptr, i64 } @__size_returning_new_aligned_hot_cold(i64 noundef 8, i64 noundef 32, i8 noundef zeroext 1)
+// CHECK: call { ptr, i64 } @__size_returning_new(i64 noundef 8){{.*}} !alloc_token [[META_LONG]]
+// CHECK: call { ptr, i64 } @__size_returning_new_hot_cold(i64 noundef 8, i8 noundef zeroext 1){{.*}} !alloc_token [[META_LONG]]
+// CHECK: call { ptr, i64 } @__size_returning_new_aligned(i64 noundef 8, i64 noundef 32){{.*}} !alloc_token [[META_LONG]]
+// CHECK: call { ptr, i64 } @__size_returning_new_aligned_hot_cold(i64 noundef 8, i64 noundef 32, i8 noundef zeroext 1){{.*}}_token [[META_LONG]]
 void test_size_returning_new() {
-  // FIXME: Support __size_returning_new variants.
   sink = __size_returning_new(sizeof(long)).p;
   sink = __size_returning_new_hot_cold(sizeof(long), __hot_cold_t{1}).p;
   sink = __size_returning_new_aligned(sizeof(long), std::align_val_t{32}).p;

``````````

</details>


https://github.com/llvm/llvm-project/pull/168840