[Clang] Introduce malloc_span attribute

The "malloc" attribute restricts the possible function signatures to
the ones returning a pointer, which is not the case for some non-standard
allocation functions variants. For example, P0901R11 proposed ::operator new
overloads that return a return_size_t result - a struct that contains
a pointer to the allocated memory as well as the actual size of the
allocated memory. Another example is __size_returning_new.

Introduce a new "malloc_span" attribute that exhibits similar semantics,
but can be applied to the functions returning records whose first member is
a pointer (which would be assumed to point to the allocated memory).
This is the case for return_size_t as well as std::span, should it be
returned from such an annotated function.

An alternative approach would be to relax the restrictions for the
existing "malloc" attribute to be applied to both functions returning
pointers and functions returning span-like structs. However, it would
complicate the user-space code that would now need to check for the
specific Clang version while the presence of a new attribute is
straightforward to check via __has_attribute. Another concern would have
been the compatibility between "malloc" semantics between Clang and GCC,
which is not the case for a new separate "malloc_span" attribute.

In future commits, codegen can be improved to recognize the
noalias-ness of the pointer returned inside a span-like struct.

This change helps unlock the alloc token instrumentation for such
non-standard allocation functions:
https://clang.llvm.org/docs/AllocToken.html#instrumenting-non-standard-allocation-functions

Author: a-nogikh

clang/docs/ReleaseNotes.rst

@@ -327,6 +327,10 @@ Attribute Changes in Clang
 - New format attributes ``gnu_printf``, ``gnu_scanf``, ``gnu_strftime`` and ``gnu_strfmon`` are added
   as aliases for ``printf``, ``scanf``, ``strftime`` and ``strfmon``. (#GH16219)
 
+- New function attribute `malloc_span` is added. It has the `malloc` semantics, but must be applied
+  not to functions returning pointers, but to functions returning span-like structures (i.e. those
+  that contain a pointer field and a size integer field).
+
 Improvements to Clang's diagnostics
 -----------------------------------
 - Diagnostics messages now refer to ``structured binding`` instead of ``decomposition``,

clang/include/clang/Basic/Attr.td

@@ -2068,6 +2068,12 @@ def Restrict : InheritableAttr {
   let Documentation = [RestrictDocs];
 }
 
+def MallocSpan : InheritableAttr {
+  let Spellings = [Clang<"malloc_span">];
+  let Subjects = SubjectList<[Function]>;
+  let Documentation = [MallocSpanDocs];
+}
+
 def LayoutVersion : InheritableAttr, TargetSpecificAttr<TargetMicrosoftRecordLayout> {
   let Spellings = [Declspec<"layout_version">];
   let Args = [UnsignedArgument<"Version">];

clang/include/clang/Basic/AttrDocs.td

@@ -5247,6 +5247,20 @@ yet implemented in clang.
   }];
 }
 
+def MallocSpanDocs : Documentation {
+  let Category = DocCatFunction;
+  let Heading = "malloc_span";
+  let Content = [{
+The ``malloc_span`` marks that a function acts like a system memory allocation
+function returning a span-like structure does not alias storage from any other
+object accessible to the caller.
+
+In this context, a span-like structure is assumed to have a pointer as its first
+field and an integer with the size of the actually allocated memory as the
+second field.
+  }];
+}
+
 def ReturnsNonNullDocs : Documentation {
   let Category = NullabilityDocs;
   let Content = [{

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -3449,6 +3449,10 @@ def err_attribute_integers_only : Error<
 def warn_attribute_return_pointers_only : Warning<
   "%0 attribute only applies to return values that are pointers">,
   InGroup<IgnoredAttributes>;
+def warn_attribute_return_span_only
+    : Warning<"%0 attribute only applies to return values that are span-like "
+              "structures">,
+      InGroup<IgnoredAttributes>;
 def warn_attribute_return_pointers_refs_only : Warning<
   "%0 attribute only applies to return values that are pointers or references">,
   InGroup<IgnoredAttributes>;

clang/lib/CodeGen/CGExpr.cpp

@@ -6642,7 +6642,8 @@ RValue CodeGenFunction::EmitCall(QualType CalleeType,
                                   CalleeDecl);
     }
     if (CalleeDecl->hasAttr<RestrictAttr>() ||
-        CalleeDecl->hasAttr<AllocSizeAttr>()) {
+        CalleeDecl->hasAttr<AllocSizeAttr>() ||
+        CalleeDecl->hasAttr<MallocSpanAttr>()) {
       // Function has 'malloc' (aka. 'restrict') or 'alloc_size' attribute.
       if (SanOpts.has(SanitizerKind::AllocToken)) {
         // Set !alloc_token metadata.

clang/lib/Sema/SemaDeclAttr.cpp

@@ -1839,6 +1839,41 @@ static void handleRestrictAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
                  RestrictAttr(S.Context, AL, DeallocE, DeallocPtrIdx));
 }
 
+static bool isSpanLikeType(const QualType &Ty) {
+  // Check that the type is a plain record with the first field being a pointer
+  // type and the second field being an integer.
+  // This matches the common implementation of std::span or sized_allocation_t
+  // in P0901R11.
+  // Note that there may also be numerous cases of pointer+integer structures
+  // not actually exhibiting a std::span-like semantics, so sometimes
+  // this heuristic expectedly leads to false positive results.
+  const RecordDecl *RD = Ty->getAsRecordDecl();
+  if (!RD || RD->isUnion())
+    return false;
+  const RecordDecl *Def = RD->getDefinition();
+  if (!Def)
+    return false; // This is an incomplete type.
+  auto FieldsBegin = Def->field_begin();
+  if (std::distance(FieldsBegin, Def->field_end()) != 2)
+    return false;
+  const FieldDecl *FirstField = *FieldsBegin;
+  const FieldDecl *SecondField = *std::next(FieldsBegin);
+  return FirstField->getType()->isAnyPointerType() &&
+         SecondField->getType()->isIntegerType();
+}
+
+static void handleMallocSpanAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
+  QualType ResultType = getFunctionOrMethodResultType(D);
+  if (!isSpanLikeType(ResultType)) {
+    S.Diag(AL.getLoc(), diag::warn_attribute_return_span_only)
+        << AL << getFunctionOrMethodResultSourceRange(D);
+    return;
+  }
+  // TODO: do we need a special check for the number of args to be 0?
+  // Or is it auto-generated?
+  D->addAttr(::new (S.Context) MallocSpanAttr(S.Context, AL));
+}
+
 static void handleCPUSpecificAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   // Ensure we don't combine these with themselves, since that causes some
   // confusing behavior.
@@ -7278,6 +7313,9 @@ ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, const ParsedAttr &AL,
   case ParsedAttr::AT_Restrict:
     handleRestrictAttr(S, D, AL);
     break;
+  case ParsedAttr::AT_MallocSpan:
+    handleMallocSpanAttr(S, D, AL);
+    break;
   case ParsedAttr::AT_Mode:
     handleModeAttr(S, D, AL);
     break;

clang/test/CodeGen/attr-malloc_span.c

@@ -0,0 +1,12 @@
+// RUN: %clang_cc1 -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s 
+
+int *Mem;
+
+typedef struct {
+  void *p;
+  int size;
+} sized_ptr;
+
+// It should not set the no alias attribute (for now).
+__attribute__((malloc_span)) sized_ptr MallocSpan(){ return (sized_ptr){ .p = Mem };}
+// CHECK: define dso_local { ptr, i32 } @MallocSpan

clang/test/Misc/pragma-attribute-supported-attributes-list.test

@@ -102,6 +102,7 @@
 // CHECK-NEXT: MIGServerRoutine (SubjectMatchRule_function, SubjectMatchRule_objc_method, SubjectMatchRule_block)
 // CHECK-NEXT: MSConstexpr (SubjectMatchRule_function)
 // CHECK-NEXT: MSStruct (SubjectMatchRule_record)
+// CHECK-NEXT: MallocSpan (SubjectMatchRule_function)
 // CHECK-NEXT: MaybeUndef (SubjectMatchRule_variable_is_parameter)
 // CHECK-NEXT: MicroMips (SubjectMatchRule_function)
 // CHECK-NEXT: MinSize (SubjectMatchRule_function, SubjectMatchRule_objc_method)

clang/test/Sema/attr-malloc_span.c

@@ -0,0 +1,31 @@
+// RUN: %clang_cc1 -verify -fsyntax-only %s
+// RUN: %clang_cc1 -emit-llvm -o %t %s
+
+#include <stddef.h>
+
+typedef struct {
+  void *ptr;
+  size_t n;
+} sized_ptr;
+sized_ptr  returns_sized_ptr  (void) __attribute((malloc_span)); // no-warning
+
+// The first struct field must be pointer and the second must be an integer.
+// Check the possible ways to violate it.
+typedef struct {
+  size_t n;
+  void *ptr;
+} invalid_span1;
+invalid_span1  returns_non_std_span1  (void) __attribute((malloc_span)); // expected-warning {{attribute only applies to return values that are span-like structures}}
+
+typedef struct {
+  void *ptr;
+  void *ptr2;
+} invalid_span2;
+invalid_span2  returns_non_std_span2  (void) __attribute((malloc_span)); // expected-warning {{attribute only applies to return values that are span-like structures}}
+
+typedef struct {
+  void *ptr;
+  size_t n;
+  size_t n2;
+} invalid_span3;
+invalid_span3  returns_non_std_span3  (void) __attribute((malloc_span)); // expected-warning {{attribute only applies to return values that are span-like structures}}