[Clang] Allow malloc for functions returning structs

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)
 
+- The malloc attribute can now be applied to functions returning span-like structures (structs
+  containing a pointer field and a size integer field). Currently, this is primarily used for
+  Allocation Token instrumentation.
+
 Improvements to Clang's diagnostics
 -----------------------------------
 - Diagnostics messages now refer to ``structured binding`` instead of ``decomposition``,

clang/include/clang/Basic/AttrDocs.td

@@ -5244,6 +5244,11 @@ deallocation function. When this form is used, it enables the compiler to
 diagnose when the incorrect deallocation function is used with this variable.
 However the associated warning, spelled `-Wmismatched-dealloc` in GCC, is not
 yet implemented in clang.
+
+It's also possible to place the ``malloc`` attribute on the functions returning
+span-like structures, that is, the structs having a pointer as the first field
+and an integer with the size of the actually allocated memory as the second field.
+Note that extended semantics is clang-specific and not currently supported in GCC.
   }];
 }
 

clang/lib/CodeGen/CGCall.cpp

@@ -2507,7 +2507,8 @@ void CodeGenModule::ConstructAttributeList(StringRef Name,
       FuncAttrs.addAttribute(llvm::Attribute::NoUnwind);
     }
     if (const auto *RA = TargetDecl->getAttr<RestrictAttr>();
-        RA && RA->getDeallocator() == nullptr)
+        RA && RA->getDeallocator() == nullptr &&
+        FI.getReturnType()->getAs<PointerType>())
       RetAttrs.addAttribute(llvm::Attribute::NoAlias);
     if (TargetDecl->hasAttr<ReturnsNonNullAttr>() &&
         !CodeGenOpts.NullPointerIsValid)

clang/lib/Sema/SemaDeclAttr.cpp

@@ -1748,9 +1748,33 @@ static void handleTLSModelAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   D->addAttr(::new (S.Context) TLSModelAttr(S.Context, AL, Model));
 }
 
+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 handleRestrictAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   QualType ResultType = getFunctionOrMethodResultType(D);
-  if (!ResultType->isAnyPointerType() && !ResultType->isBlockPointerType()) {
+  if (!ResultType->isAnyPointerType() && !ResultType->isBlockPointerType() &&
+      !isSpanLikeType(ResultType)) {
     S.Diag(AL.getLoc(), diag::warn_attribute_return_pointers_only)
         << AL << getFunctionOrMethodResultSourceRange(D);
     return;

clang/test/CodeGen/attr-malloc.c

@@ -3,10 +3,17 @@
 int *Mem;
 void dealloc(int*);
 
+typedef struct {
+  void *p;
+  int size;
+} sized_ptr;
+
 __attribute__((malloc)) int *MallocFunc(){ return Mem;}
 // CHECK: define[[BEFORE:.*]] noalias[[AFTER:.*]]@MallocFunc
-// Ensure these two do not generate noalias here.
+// Ensure these three do not generate noalias here.
 __attribute__((malloc(dealloc))) int *MallocFunc2(){ return Mem;}
 // CHECK: define[[BEFORE]][[AFTER]]@MallocFunc2
 __attribute__((malloc(dealloc, 1))) int *MallocFunc3(){ return Mem;}
 // CHECK: define[[BEFORE]][[AFTER]]@MallocFunc3
+__attribute__((malloc)) sized_ptr MallocFunc4(){ return (sized_ptr){ .p = Mem };}
+// CHECK: define[[BEFORE]] { ptr, i32 } @MallocFunc4

clang/test/Sema/attr-malloc.c

@@ -13,6 +13,32 @@ int   returns_int   (void) __attribute((malloc)); // expected-warning {{attribut
 int * returns_intptr(void) __attribute((malloc)); // no-warning
 typedef int * iptr;
 iptr  returns_iptr  (void) __attribute((malloc)); // no-warning
+typedef struct {
+  void *ptr;
+  size_t n;
+} sized_ptr;
+sized_ptr  returns_sized_ptr  (void) __attribute((malloc)); // 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)); // expected-warning {{attribute only applies to return values that are pointers}}
+
+typedef struct {
+  void *ptr;
+  void *ptr2;
+} invalid_span2;
+invalid_span2  returns_non_std_span2  (void) __attribute((malloc)); // expected-warning {{attribute only applies to return values that are pointers}}
+
+typedef struct {
+  void *ptr;
+  size_t n;
+  size_t n2;
+} invalid_span3;
+invalid_span3  returns_non_std_span3  (void) __attribute((malloc)); // expected-warning {{attribute only applies to return values that are pointers}}
 
 __attribute((malloc)) void *(*f)(void); //  expected-warning{{attribute only applies to functions}}
 __attribute((malloc)) int (*g)(void); // expected-warning{{attribute only applies to functions}}