[Clang] Allow malloc and alloc_size attributes for functions returning structs

These attributes restrict 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.

Relax the mentioned restriction and allow "malloc" and "alloc_size"
attributes to be applied not just to the functions returning pointer
types, but also 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.

As the previous restriction was dictated by the need to set the noalias
attribute, adjust the code generation logic to only assign the attribute
when the returned value type is a pointer type.

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

This change also 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/include/clang/AST/TypeBase.h

@@ -2601,6 +2601,7 @@ class alignas(TypeAlignment) Type : public ExtQualsTypeCommonBase {
   bool isFunctionProtoType() const { return getAs<FunctionProtoType>(); }
   bool isPointerType() const;
   bool isPointerOrReferenceType() const;
+  bool isSpanLikeType() const;
   bool isSignableType(const ASTContext &Ctx) const;
   bool isSignablePointerType() const;
   bool isSignableIntegerType(const ASTContext &Ctx) const;

clang/lib/AST/Type.cpp

@@ -5753,3 +5753,23 @@ StringRef PredefinedSugarType::getName(Kind KD) {
   }
   llvm_unreachable("unexpected kind");
 }
+
+bool Type::isSpanLikeType() const {
+  // 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.
+  const RecordDecl *RD = 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();
+}

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

@@ -571,7 +571,7 @@ static void handleAllocSizeAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   assert(isFuncOrMethodForAttrSubject(D) && hasFunctionProto(D));
 
   QualType RetTy = getFunctionOrMethodResultType(D);
-  if (!RetTy->isPointerType()) {
+  if (!RetTy->isPointerType() && !RetTy->isSpanLikeType()) {
     S.Diag(AL.getLoc(), diag::warn_attribute_return_pointers_only) << AL;
     return;
   }
@@ -1750,7 +1750,8 @@ static void handleTLSModelAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
 
 static void handleRestrictAttr(Sema &S, Decl *D, const ParsedAttr &AL) {
   QualType ResultType = getFunctionOrMethodResultType(D);
-  if (!ResultType->isAnyPointerType() && !ResultType->isBlockPointerType()) {
+  if (!ResultType->isAnyPointerType() && !ResultType->isBlockPointerType() &&
+      !ResultType->isSpanLikeType()) {
     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/alloc-size.c

@@ -30,6 +30,14 @@ void *KR() __attribute__((alloc_size(1))); //expected-warning{{'alloc_size' attr
 void *(__attribute__((alloc_size(1))) * func_ptr1)(int);
 void *(__attribute__((alloc_size(1, 2))) func_ptr2)(int, int);
 
+// Applying alloc_size to functions returning a struct with a pointer as a first field should work.
+typedef struct {
+  void* p;
+  int n;
+} sized_ptr;
+
+sized_ptr sized_ptr_alloc(int len) __attribute__((alloc_size(1)));
+
 // TODO: according to GCC documentation the following should actually be the type
 // “pointer to pointer to alloc_size attributed function returning void*” and should
 // therefore be supported

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}}