[BoundsSafety][Sema] Allow counted_by and counted_by_or_null on pointers where the pointee type is incomplete but potentially completable

[delcypher]

Previously using the counted_by or counted_by_or_null attribute on a pointer with an incomplete pointee type was forbidden. Unfortunately this prevented a situation like the following from being allowed.

Header file:

struct EltTy; // Incomplete type
struct Buffer {
  size_t count;
  struct EltTy* __counted_by(count) buffer; // error before this patch
};

Implementation file:

struct EltTy {
  // definition
};

void allocBuffer(struct Buffer* b) {
  b->buffer = malloc(sizeof(EltTy)* b->count);
}

To allow code like the above but still enforce that the pointee type size is known in locations where -fbounds-safety needs to emit bounds checks the following scheme is used.

• For incomplete pointee types that can never be completed (e.g. void) these are treated as error where the attribute is written (just like before this patch).
• For incomplete pointee types that might be completable later on (struct, union, and enum forward declarations) in the translation unit, writing the attribute on the incomplete pointee type is allowed on the FieldDecl declaration but "uses" of the declared pointer are forbidden if at the point of "use" the pointee type is still incomplete.

For this patch a "use" of a FieldDecl covers:

• Explicit and Implicit initialization (note see Tentative Definition Initialization for an exception to this)
• Assignment
• Conversion to an lvalue (e.g. for use in an expression)

In Apple's internal fork of Clang the counted_by and counted_by_or_null attribute are allowed in many more contexts. That isn't the case for upstream Clang so the "use" checks for the attribute on VarDecl, ParamVarDecl, and function return type have been omitted from this patch because they can't be tested. However, the BoundsSafetyCheckAssignmentToCountAttrPtrWithIncompletePointeeTy and BoundsSafetyCheckUseOfCountAttrPtrWithIncompletePointeeTy functions retain the ability to emit diagnostics for these other contexts to avoid unnecessary divergence between upstream Clang and Apple's internal fork. Support for checking "uses" will be upstreamed when upstream Clang allows the counted_by and counted_by_or_null attribute in additional contexts.

This patch has a few limitations:

** 1. Tentative Defition Initialization **

This patch currently allows something like:

struct IncompleteTy;
struct Buffer {
  int count;
  struct IncompleteTy* __counted_by(count) buf;
};

// Tentative definition
struct Buffer GlobalBuf;

The Tentative definition in this example becomes an actual definition whose initialization should be checked but it currently isn't. Addressing this problem will be done in a subseqent patch.

** 2. When the incomplete pointee type is a typedef diagnostics are slightly misleading **

For this situation:

struct IncompleteTy;
typedef struct IncompleteTy Incomplete_t;

struct Buffer {
  int count;
  struct IncompleteTy* __counted_by(count) buf;
};

void use(struct Buffer b) {
  b.buf = 0x0;
}

This code emits note: forward declaration of 'Incomplete_t' (aka 'struct IncompleteTy') but the location is on the struct IncompleteTy; forward declaration. This is misleading because Incomplete_t isn't actually forward declared there (instead the underlying type is). This could be resolved by additional diagnostics that walk the chain of typedefs and explain each step of the walk. However, that would be very verbose and didn't seem like a direction worth pursuing.

rdar://133600117

[llvmbot]

@llvm/pr-subscribers-clang

Author: Dan Liew (delcypher)

Changes

Previously using the counted_by or counted_by_or_null attribute on a pointer with an incomplete pointee type was forbidden. Unfortunately this prevented a situation like the following from being allowed.

Header file:

struct EltTy; // Incomplete type
struct Buffer {
  size_t count;
  struct EltTy* __counted_by(count) buffer; // error before this patch
};

Implementation file:

struct EltTy {
  // definition
};

void allocBuffer(struct Buffer* b) {
  b->buffer = malloc(sizeof(EltTy)* b->count);
}

To allow code like the above but still enforce that the pointee type size is known in locations where -fbounds-safety needs to emit bounds checks the following scheme is used.

• For incomplete pointee types that can never be completed (e.g. void) these are treated as error where the attribute is written (just like before this patch).
• For incomplete pointee types that might be completable later on (struct, union, and enum forward declarations) in the translation unit, writing the attribute on the incomplete pointee type is allowed on the FieldDecl declaration but "uses" of the declared pointer are forbidden if at the point of "use" the pointee type is still incomplete.

For this patch a "use" of a FieldDecl covers:

• Explicit and Implicit initialization (note see Tentative Definition Initialization for an exception to this)
• Assignment
• Conversion to an lvalue (e.g. for use in an expression)

In Apple's internal fork of Clang the counted_by and counted_by_or_null attribute are allowed in many more contexts. That isn't the case for upstream Clang so the "use" checks for the attribute on VarDecl, ParamVarDecl, and function return type have been omitted from this patch because they can't be tested. However, the BoundsSafetyCheckAssignmentToCountAttrPtrWithIncompletePointeeTy and BoundsSafetyCheckUseOfCountAttrPtrWithIncompletePointeeTy functions retain the ability to emit diagnostics for these other contexts to avoid unnecessary divergence between upstream Clang and Apple's internal fork. Support for checking "uses" will be upstreamed when upstream Clang allows the counted_by and counted_by_or_null attribute in additional contexts.

This patch also includes the new Sema::BoundsSafetySourceRangeFor method which allows a more accurate SourceRange for the counted_by and counted_by_or_null attributes to be shown in diagnostics. This is used by the new diag::note_named_attribute diagnostic introduced in this patch. This method has several shortcomings but it didn't seem worth covering them given this is essentially a workaround for the fact CountAttributedType doesn't store the SourceLocation for its attribute.

This patch has a few limitations:

** 1. Tentative Defition Initialization **

This patch currently allows something like:

struct IncompleteTy;
struct Buffer {
  int count;
  struct IncompleteTy* __counted_by(count) buf;
};

// Tentative definition
struct Buffer GlobalBuf;

The Tentative definition in this example becomes an actual definition whose initialization should be checked but it currently isn't. Addressing this problem will be done in a subseqent patch.

** 2. When the incomplete pointee type is a typedef diagnostics are slightly misleading **

For this situation:

struct IncompleteTy;
typedef struct IncompleteTy Incomplete_t;

struct Buffer {
  int count;
  struct IncompleteTy* __counted_by(count) buf;
};

void use(struct Buffer b) {
  b.buf = 0x0;
}

This code emits note: forward declaration of 'Incomplete_t' (aka 'struct IncompleteTy') but the location is on the struct IncompleteTy; forward declaration. This is misleading because Incomplete_t isn't actually forward declared there (instead the underlying type is). This could be resolved by additional diagnostics that walk the chain of typedefs and explain each step of the walk. However, that would be very verbose and didn't seem like a direction worth pursuing.

rdar://133600117

---

Patch is 109.87 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/106321.diff

16 Files Affected:

• (modified) clang/include/clang/AST/Type.h (+22)
• (modified) clang/include/clang/Basic/DiagnosticSemaKinds.td (+37)
• (modified) clang/include/clang/Sema/Sema.h (+51)
• (modified) clang/lib/AST/Type.cpp (+37)
• (modified) clang/lib/Sema/SemaBoundsSafety.cpp (+399-1)
• (modified) clang/lib/Sema/SemaExpr.cpp (+24)
• (modified) clang/lib/Sema/SemaInit.cpp (+17)
• (modified) clang/test/Sema/attr-counted-by-late-parsed-struct-ptrs.c (+2-3)
• (modified) clang/test/Sema/attr-counted-by-or-null-last-field.c (+2-2)
• (modified) clang/test/Sema/attr-counted-by-or-null-late-parsed-struct-ptrs.c (+2-3)
• (modified) clang/test/Sema/attr-counted-by-or-null-struct-ptrs.c (+4-8)
• (added) clang/test/Sema/attr-counted-by-struct-ptrs-completable-incomplete-pointee.c (+584)
• (modified) clang/test/Sema/attr-counted-by-struct-ptrs.c (+4-8)
• (modified) clang/test/Sema/attr-counted-by-vla.c (+1-2)
• (modified) clang/unittests/Sema/CMakeLists.txt (+1)
• (added) clang/unittests/Sema/GetCountedByAttrSourceRange.cpp (+573)

diff --git a/clang/include/clang/AST/Type.h b/clang/include/clang/AST/Type.h
index a4804e4c6f61ca..fb4638d10f121c 100644
--- a/clang/include/clang/AST/Type.h
+++ b/clang/include/clang/AST/Type.h
@@ -2417,6 +2417,26 @@ class alignas(TypeAlignment) Type : public ExtQualsTypeCommonBase {
     return !isFunctionType();
   }
 
+  /// \returns True if the type is incomplete and it is also a type that
+  /// cannot be completed by a later type definition.
+  ///
+  /// E.g. For `void` this is true but for `struct ForwardDecl;` this is false
+  /// because a definition for `ForwardDecl` could be provided later on in the
+  /// translation unit.
+  ///
+  /// Note even for types that this function returns true for it is still
+  /// possible for the declarations that contain this type to later have a
+  /// complete type in a translation unit. E.g.:
+  ///
+  /// \code{.c}
+  /// // This decl has type 'char[]' which is incomplete and cannot be later
+  /// // completed by another by another type declaration.
+  /// extern char foo[];
+  /// // This decl how has complete type 'char[5]'.
+  /// char foo[5]; // foo has a complete type
+  /// \endcode
+  bool isIncompletableIncompleteType() const;
+
   /// Determine whether this type is an object type.
   bool isObjectType() const {
     // C++ [basic.types]p8:
@@ -3320,6 +3340,8 @@ class CountAttributedType final
   static bool classof(const Type *T) {
     return T->getTypeClass() == CountAttributed;
   }
+
+  StringRef GetAttributeName(bool WithMacroPrefix) const;
 };
 
 /// Represents a type which was implicitly adjusted by the semantic
diff --git a/clang/include/clang/Basic/DiagnosticSemaKinds.td b/clang/include/clang/Basic/DiagnosticSemaKinds.td
index 5cdf36660b2a66..305fab7cd0c8db 100644
--- a/clang/include/clang/Basic/DiagnosticSemaKinds.td
+++ b/clang/include/clang/Basic/DiagnosticSemaKinds.td
@@ -4370,6 +4370,7 @@ def err_mismatched_visibility: Error<"visibility does not match previous declara
 def note_previous_attribute : Note<"previous attribute is here">;
 def note_conflicting_attribute : Note<"conflicting attribute is here">;
 def note_attribute : Note<"attribute is here">;
+def note_named_attribute : Note<"%0 attribute is here">;
 def err_mismatched_ms_inheritance : Error<
   "inheritance model does not match %select{definition|previous declaration}0">;
 def warn_ignored_ms_inheritance : Warning<
@@ -6619,6 +6620,42 @@ def err_counted_by_attr_pointee_unknown_size : Error<
     "%select{|. This will be an error in a future compiler version}3"
     ""
   "}2">;
+def err_counted_by_on_incomplete_type_on_assign : Error <
+  "cannot %select{"
+    "assign to %select{object|'%1'}2 that has|" // AA_Assigning,
+    "pass argument to %select{parameter|parameter '%1'}2 that has|" // AA_Passing,
+    "return|" // AA_Returning,
+    "convert to|" // AA_Converting (UNUSED)
+    "%select{|implicitly }3initialize %select{object|'%1'}2 that has|" // AA_Initializing,
+    "pass argument to parameter that has|" // AA_Sending (UNUSED)
+    "cast to|" // AA_Casting (UNUSED)
+    "pass argument to parameter that has" // AA_Passing_CFAudited (UNUSED)
+  "}0 type %4 because the pointee type %6 is incomplete and the '%5' attribute "
+  "requires the pointee type be complete when %select{"
+    "assigning|" // AA_Assigning,
+    "passing|" // AA_Passing,
+    "returning|" // AA_Returning,
+    "converting|" // AA_Converting (UNUSED)
+    "initializing|" // AA_Initializing,
+    "passing|" // AA_Sending (UNUSED)
+    "casting|" // AA_Casting (UNUSED)
+    "passing" // AA_Passing_CFAudited (UNUSED)
+  "}0; "
+  "consider providing a complete definition for %6 or using the "
+  "'__sized_by%select{|_or_null}7' attribute"
+>;
+def err_counted_by_on_incomplete_type_on_use : Error <
+  "cannot %select{"
+    "use '%1'|" // Generic expr
+    "call '%1'" // CallExpr
+  "}0 with%select{"
+    "|"       // Generic expr
+    " return" // CallExpr
+  "}0 type %2 because the pointee type %3 is incomplete "
+  "and the '%4' attribute requires the pointee type be complete in this context; "
+  "consider providing a complete definition for %3 or using the "
+  "'__sized_by%select{|_or_null}5' attribute"
+>;
 
 def warn_counted_by_attr_elt_type_unknown_size :
   Warning<err_counted_by_attr_pointee_unknown_size.Summary>,
diff --git a/clang/include/clang/Sema/Sema.h b/clang/include/clang/Sema/Sema.h
index b7bd6c2433efd6..b6a194cf32f0a6 100644
--- a/clang/include/clang/Sema/Sema.h
+++ b/clang/include/clang/Sema/Sema.h
@@ -15086,6 +15086,57 @@ class Sema final : public SemaBase {
   bool CheckCountedByAttrOnField(FieldDecl *FD, Expr *E, bool CountInBytes,
                                  bool OrNull);
 
+  /// Perform Bounds Safety Semantic checks for assigning to a `__counted_by` or
+  /// `__counted_by_or_null` pointer type \param LHSTy.
+  ///
+  /// \param LHSTy The type being assigned to. Checks will only be performed if
+  ///              the type is a `counted_by` or `counted_by_or_null ` pointer.
+  /// \param RHSExpr The expression being assigned from.
+  /// \param Action The type assignment being performed
+  /// \param Loc The SourceLocation to use for error diagnostics
+  /// \param ComputeAssignee If provided this function will be called before
+  ///        emitting a diagnostic. The function should return the name of
+  ///        entity being assigned to or an empty string if this cannot be
+  ///        determined.
+  ///
+  /// \returns True iff no diagnostic where emitted, false otherwise.
+  bool BoundsSafetyCheckAssignmentToCountAttrPtr(
+      QualType LHSTy, Expr *RHSExpr, Sema::AssignmentAction Action,
+      SourceLocation Loc,
+      std::function<std::string()> ComputeAssignee = nullptr);
+
+  /// Perform Bounds Safety Semantic checks for initializing a Bounds Safety
+  /// pointer.
+  ///
+  /// \param Entity The entity being initialized
+  /// \param Kind The kind of initialization being performed
+  /// \param Action The type assignment being performed
+  /// \param LHSTy The type being assigned to. Checks will only be performed if
+  ///              the type is a `counted_by` or `counted_by_or_null ` pointer.
+  /// \param RHSExpr The expression being used for initialization.
+  ///
+  /// \returns True iff no diagnostic where emitted, false otherwise.
+  bool BoundsSafetyCheckInitialization(const InitializedEntity &Entity,
+                                       const InitializationKind &Kind,
+                                       Sema::AssignmentAction Action,
+                                       QualType LHSType, Expr *RHSExpr);
+
+  /// Perform Bounds Safety semantic checks for uses of invalid uses counted_by
+  /// or counted_by_or_null pointers in \param E.
+  ///
+  /// \param E the expression to check
+  ///
+  /// \returns True iff no diagnostic where emitted, false otherwise.
+  bool BoundsSafetyCheckUseOfCountAttrPtr(Expr *E);
+
+  /// \returns A SourceRange for \param CATy.
+  ///
+  /// This method tries to compute the most useful SourceRange for diagnostics.
+  /// If this fails the returned SourceRange will be the same as
+  /// `CATy->getCountExpr()->getSourceRange()`.
+  ///
+  SourceRange BoundsSafetySourceRangeFor(const CountAttributedType *CATy);
+
   ///@}
 };
 
diff --git a/clang/lib/AST/Type.cpp b/clang/lib/AST/Type.cpp
index e89ce2e4b38445..137609743fc02c 100644
--- a/clang/lib/AST/Type.cpp
+++ b/clang/lib/AST/Type.cpp
@@ -2438,6 +2438,22 @@ bool Type::isIncompleteType(NamedDecl **Def) const {
   }
 }
 
+bool Type::isIncompletableIncompleteType() const {
+  if (!isIncompleteType())
+    return false;
+
+  // Forward declarations of structs, classes, enums, and unions could be later
+  // completed in a compilation unit by providing a type definition.
+  if (isStructureOrClassType() || isEnumeralType() || isUnionType())
+    return false;
+
+  // Other types are incompletable.
+  //
+  // E.g. `char[]` and `void`. The type is incomplete and no future
+  // type declarations can make the type complete.
+  return true;
+}
+
 bool Type::isSizelessBuiltinType() const {
   if (isSizelessVectorType())
     return true;
@@ -3862,6 +3878,27 @@ CountAttributedType::CountAttributedType(
     DeclSlot[i] = CoupledDecls[i];
 }
 
+StringRef CountAttributedType::GetAttributeName(bool WithMacroPrefix) const {
+#define ENUMERATE_ATTRS(PREFIX)                                                \
+  do {                                                                         \
+    if (isCountInBytes()) {                                                    \
+      if (isOrNull())                                                          \
+        return PREFIX "sized_by_or_null";                                      \
+      return PREFIX "sized_by";                                                \
+    }                                                                          \
+    if (isOrNull())                                                            \
+      return PREFIX "counted_by_or_null";                                      \
+    return PREFIX "counted_by";                                                \
+  } while (0)
+
+  if (WithMacroPrefix)
+    ENUMERATE_ATTRS("__");
+  else
+    ENUMERATE_ATTRS("");
+
+#undef ENUMERATE_ATTRS
+}
+
 TypedefType::TypedefType(TypeClass tc, const TypedefNameDecl *D,
                          QualType Underlying, QualType can)
     : Type(tc, can, toSemanticDependence(can->getDependence())),
diff --git a/clang/lib/Sema/SemaBoundsSafety.cpp b/clang/lib/Sema/SemaBoundsSafety.cpp
index d63a2389ea11de..fa72553317a810 100644
--- a/clang/lib/Sema/SemaBoundsSafety.cpp
+++ b/clang/lib/Sema/SemaBoundsSafety.cpp
@@ -11,6 +11,9 @@
 /// (e.g. `counted_by`)
 ///
 //===----------------------------------------------------------------------===//
+#include "clang/Basic/SourceManager.h"
+#include "clang/Lex/Lexer.h"
+#include "clang/Sema/Initialization.h"
 #include "clang/Sema/Sema.h"
 
 namespace clang {
@@ -102,7 +105,36 @@ bool Sema::CheckCountedByAttrOnField(FieldDecl *FD, Expr *E, bool CountInBytes,
   // only `PointeeTy->isStructureTypeWithFlexibleArrayMember()` is reachable
   // when `FieldTy->isArrayType()`.
   bool ShouldWarn = false;
-  if (PointeeTy->isIncompleteType() && !CountInBytes) {
+  if (PointeeTy->isIncompletableIncompleteType() && !CountInBytes) {
+    // In general using `counted_by` or `counted_by_or_null` on
+    // pointers where the pointee is an incomplete type are problematic. This is
+    // because it isn't possible to compute the pointer's bounds without knowing
+    // the pointee type size. At the same time it is common to forward declare
+    // types in header files.
+    //
+    // E.g.:
+    //
+    // struct Handle;
+    // struct Wrapper {
+    //   size_t size;
+    //   struct Handle* __counted_by(count) handles;
+    // }
+    //
+    // To allow the above code pattern but still prevent the pointee type from
+    // being incomplete in places where bounds checks are needed the following
+    // scheme is used:
+    //
+    // * When the pointee type is a "completable incomplete" type (i.e.
+    // a type that is currently incomplete but might be completed later
+    // on in the translation unit) the attribute is allowed by this method
+    // but later uses of the FieldDecl are checked that the pointee type
+    // is complete see `BoundsSafetyCheckAssignmentToCountAttrPtr`,
+    // `BoundsSafetyCheckInitialization`, and
+    // `BoundsSafetyCheckUseOfCountAttrPtr`
+    //
+    // * When the pointee type is a "incompletable incomplete" type (e.g.
+    // `void`) the attribute is disallowed by this method because we know the
+    // type can never be completed so there's no reason to allow it.
     InvalidTypeKind = CountedByInvalidPointeeTypeKind::INCOMPLETE;
   } else if (PointeeTy->isSizelessType()) {
     InvalidTypeKind = CountedByInvalidPointeeTypeKind::SIZELESS;
@@ -186,4 +218,370 @@ bool Sema::CheckCountedByAttrOnField(FieldDecl *FD, Expr *E, bool CountInBytes,
   return false;
 }
 
+SourceRange Sema::BoundsSafetySourceRangeFor(const CountAttributedType *CATy) {
+  // Note: This implementation relies on `CountAttributedType` being unique.
+  // E.g.:
+  //
+  // struct Foo {
+  //   int count;
+  //   char* __counted_by(count) buffer;
+  //   char* __counted_by(count) buffer2;
+  // };
+  //
+  // The types of `buffer` and `buffer2` are unique. The types being
+  // unique means the SourceLocation of the `counted_by` expression can be used
+  // to find where the attribute was written.
+
+  auto Fallback = CATy->getCountExpr()->getSourceRange();
+  auto CountExprBegin = CATy->getCountExpr()->getBeginLoc();
+
+  // FIXME: We currently don't support the count expression being a macro
+  // itself. E.g.:
+  //
+  // #define ZERO 0
+  // int* __counted_by(ZERO) x;
+  //
+  if (SourceMgr.isMacroBodyExpansion(CountExprBegin))
+    return Fallback;
+
+  auto FetchIdentifierTokenFromOffset =
+      [&](ssize_t Offset) -> std::optional<Token> {
+    SourceLocation OffsetLoc = CountExprBegin.getLocWithOffset(Offset);
+    Token Result;
+    if (Lexer::getRawToken(OffsetLoc, Result, SourceMgr, getLangOpts()))
+      return std::optional<Token>(); // Failed
+
+    if (!Result.isAnyIdentifier())
+      return std::optional<Token>(); // Failed
+
+    return Result; // Success
+  };
+
+  auto CountExprEnd = CATy->getCountExpr()->getEndLoc();
+  auto FindRParenTokenAfter = [&]() -> std::optional<Token> {
+    auto CountExprEndSpelling = SourceMgr.getSpellingLoc(CountExprEnd);
+    auto MaybeRParenTok = Lexer::findNextToken(
+        CountExprEndSpelling, getSourceManager(), getLangOpts());
+
+    if (!MaybeRParenTok.has_value())
+      return std::nullopt;
+
+    if (!MaybeRParenTok->is(tok::r_paren))
+      return std::nullopt;
+
+    return *MaybeRParenTok;
+  };
+
+  // Step back two characters to point at the last character of the attribute
+  // text.
+  //
+  // __counted_by(count)
+  //            ^ ^
+  //            | |
+  //            ---
+  auto MaybeLastAttrCharToken = FetchIdentifierTokenFromOffset(-2);
+  if (!MaybeLastAttrCharToken)
+    return Fallback;
+
+  auto LastAttrCharToken = MaybeLastAttrCharToken.value();
+
+  if (LastAttrCharToken.getLength() > 1) {
+    // Special case: When the character is part of a macro the Token we get
+    // is the whole macro name (e.g. `__counted_by`).
+    if (LastAttrCharToken.getRawIdentifier() !=
+        CATy->GetAttributeName(/*WithMacroPrefix=*/true))
+      return Fallback;
+
+    // Found the beginning of the `__counted_by` macro
+    SourceLocation Begin = LastAttrCharToken.getLocation();
+    // Now try to find the closing `)` of the macro.
+    auto MaybeRParenTok = FindRParenTokenAfter();
+    if (!MaybeRParenTok.has_value())
+      return Fallback;
+
+    return SourceRange(Begin, MaybeRParenTok->getLocation());
+  }
+
+  assert(LastAttrCharToken.getLength() == 1);
+  // The Token we got back is just the last character of the identifier.
+  // This means a macro is not being used and instead the attribute is being
+  // used directly. We need to find the beginning of the identifier. We support
+  // two cases:
+  //
+  // * Non-affixed version. E.g: `counted_by`
+  // * Affixed version. E.g.: `__counted_by__`
+
+  // Try non-affixed version. E.g.:
+  //
+  // __attribute__((counted_by(count)))
+  //                ^          ^
+  //                |          |
+  //                ------------
+
+  // +1 is for `(`
+  const ssize_t NonAffixedSkipCount =
+      CATy->GetAttributeName(/*WithMacroPrefix=*/false).size() + 1;
+  auto MaybeNonAffixedBeginToken =
+      FetchIdentifierTokenFromOffset(-NonAffixedSkipCount);
+  if (!MaybeNonAffixedBeginToken)
+    return Fallback;
+
+  auto NonAffixedBeginToken = MaybeNonAffixedBeginToken.value();
+  if (NonAffixedBeginToken.getRawIdentifier() ==
+      CATy->GetAttributeName(/*WithMacroPrefix=*/false)) {
+    // Found the beginning of the `counted_by`-like attribute
+    auto SL = NonAffixedBeginToken.getLocation();
+
+    // Now try to find the closing `)` of the attribute
+    auto MaybeRParenTok = FindRParenTokenAfter();
+    if (!MaybeRParenTok.has_value())
+      return Fallback;
+
+    return SourceRange(SL, MaybeRParenTok->getLocation());
+  }
+
+  // Try affixed version. E.g.:
+  //
+  // __attribute__((__counted_by__(count)))
+  //                ^              ^
+  //                |              |
+  //                ----------------
+  std::string AffixedTokenStr =
+      (llvm::Twine("__") + CATy->GetAttributeName(/*WithMacroPrefix=*/false) +
+       llvm::Twine("__"))
+          .str();
+  // +1 is for `(`
+  // +4 is for the 4 `_` characters
+  const ssize_t AffixedSkipCount =
+      CATy->GetAttributeName(/*WithMacroPrefix=*/false).size() + 1 + 4;
+  auto MaybeAffixedBeginToken =
+      FetchIdentifierTokenFromOffset(-AffixedSkipCount);
+  if (!MaybeAffixedBeginToken)
+    return Fallback;
+
+  auto AffixedBeginToken = MaybeAffixedBeginToken.value();
+  if (AffixedBeginToken.getRawIdentifier() != AffixedTokenStr)
+    return Fallback;
+
+  // Found the beginning of the `__counted_by__`-like like attribute.
+  auto SL = AffixedBeginToken.getLocation();
+  // Now try to find the closing `)` of the attribute
+  auto MaybeRParenTok = FindRParenTokenAfter();
+  if (!MaybeRParenTok.has_value())
+    return Fallback;
+
+  return SourceRange(SL, MaybeRParenTok->getLocation());
+}
+
+static void EmitIncompleteCountedByPointeeNotes(Sema &S,
+                                                const CountAttributedType *CATy,
+                                                NamedDecl *IncompleteTyDecl,
+                                                bool NoteAttrLocation = true) {
+  assert(IncompleteTyDecl == nullptr || isa<TypeDecl>(IncompleteTyDecl));
+
+  if (NoteAttrLocation) {
+    // Note where the attribute is declared
+    auto AttrSrcRange = S.BoundsSafetySourceRangeFor(CATy);
+    S.Diag(AttrSrcRange.getBegin(), diag::note_named_attribute)
+        << CATy->GetAttributeName(/*WithMacroPrefix=*/true) << AttrSrcRange;
+  }
+
+  if (!IncompleteTyDecl)
+    return;
+
+  // If there's an associated forward declaration display it to emphasize
+  // why the type is incomplete (all we have is a forward declaration).
+
+  // Note the `IncompleteTyDecl` type is the underlying type which might not
+  // be the same as `CATy->getPointeeType()` which could be a typedef.
+  //
+  // The diagnostic printed will be at the location of the underlying type but
+  // the diagnostic text will print the type of `CATy->getPointeeType()` which
+  // could be a typedef name rather than the underlying type. This is ok
+  // though because the diagnostic will print the underlying type name too.
+  // E.g:
+  //
+  // `forward declaration of 'Incomplete_Struct_t'
+  //  (aka 'struct IncompleteStructTy')`
+  //
+  // If this ends up being confusing we could emit a second diagnostic (one
+  // explaining where the typedef is) but that seems overly verbose.
+
+  S.Diag(IncompleteTyDecl->getBeginLoc(), diag::note_forward_declaration)
+      << CATy->getPointeeType();
+}
+
+static bool
+HasCountedByAttrOnIncompletePointee(QualType Ty, NamedDecl **ND,
+                                    const CountAttributedType **CATyOut,
+                                    QualType *PointeeTyOut) {
+  auto *CATy = Ty->getAs<CountAttributedType>();
+  if (!CATy)
+    return false;
+
+  // Incomplete pointee type is only a problem for
+  // counted_by/counted_by_or_null
+  if (CATy->isCountInBytes())
+    return false;
+
+  auto PointeeTy = CATy->getPointeeType();
+  if (PointeeTy.isNull())
+    return false; // Reachable?
+
+  if (!PointeeTy->isIncompleteType(ND))
+    return false;
+
+  if (CATyOut)
+    *CATyOut = CATy;
+  if (PointeeTyOut)
+    *PointeeTyOut = PointeeTy;
+  return true;
+}
+
+/// Perform Checks for assigning to a `__counted_by` or
+/// `__counted_by_or_null` pointer type \param LHSTy where the pointee type
+/// is incomplete which is invalid.
+///
+/// \param S The Sema instance.
+/// \param LHSTy The type being assigned to. Checks will only be performed if
+///              the type is a `counted_by` or `counted_by_or_null ` pointer.
+/// \param RHSExpr The expression being assigned from.
+/// \param Action The type assignment being performed
+/// \param Loc The SourceLocation to use for error diagnostics
+/// \param ComputeAssignee If provided this function will be called...
[truncated]

[delcypher]

@bwendling @kees Apologies for taking so long to get this change upstream. We spent quite a lot of time experimenting with different implementations that ended up not working out. Our internal version of this patch is larger because counted_by is allowed in many more contexts than it is in upstream Clang so some of that has been omitted from this patch. I can upstream the rest once we have support for counted_by is more contexts.

[Sirraide]

I’m pretty sure there are some fundamental problems with the current approach wrt how getting source locations for the attribute is being handled.

The general approach for the rest of it seems fine from a cursory look, though. :)

[erichkeane]

I did a pass, I think @Sirraide is on the right direction in his reviews here, but the general direction of this patch is also completely acceptable.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[delcypher]

@Sirraide @erichkeane Sorry for the delay in addressing your feedback. This PR is ready for another round of reviews. I've tried to keep all my tweaks as separate commits so it's easier to review.

[Sirraide]

I still think the comment is superfluous. This pattern of potentially diagnosing something, but not returning ExprError()/true/whatever so we can still construct a somewhat valid AST is fairly common in Sema, so I don’t think every (or any) call to this function in this fashion really needs to be commented (and the same applies to the call in InitializationSequence::Perform()).

[delcypher]

It would be superfluous if not for the part about struct initialization. To me that is really not very obvious. When possible I think it's very useful to comment why the code is doing something. All too often I come across code where I understand what it's doing but I don't understand why because there are no comments explaining the reason for the behavior.

So I would prefer to keep it. If you prefer I could only keep the bit about struct initialization.

[delcypher]

I've changed the comments to just be.

        // Note the return value isn't used to return a ExprError() when
        // initialization fails . For struct initialization allows all field
        // assignments to be checked rather than bailing on the first error.

which only discusses the non-obvious part of the original comment.

I'm going to mark this as resolved but please re-open if you want to discuss this more.

[Sirraide]

It would be superfluous if not for the part about struct initialization. To me that is really not very obvious.

That’s fair. I might just be to used to this as a pattern. Commenting why we don’t just throw the entire thing out of the window here and still try to keep operating on it even if this fails is probably fine.

I've changed the comments to just be.

        // Note the return value isn't used to return a ExprError() when
        // initialization fails . For struct initialization allows all field
        // assignments to be checked rather than bailing on the first error.

I’d honestly shorten it even more to just:

// Don't return here so we can still check field assingments.

because e.g. the fact that ‘the return value isn't used to return a ExprError() when initialization fails’ is apparent from the fact that... there is no return statement here. ;Þ

[erichkeane]

I think I'm reasonably OK with this, though @Sirraide has a 'requested changes', so he should do the final approval.

[AaronBallman]

It's not clear to me why this is needed at all -- doesn't Attr::getSpelling() suffice on the attribute, rather than asking the type?

[delcypher]

In this context AFAICT we do not have access to the attribute. All we have is the CountAttributedType and we have no easy way to find the corresponding Attr object from this type. This is what #113585 tracks fixing.

There's a related comment about this in this review here.

[AaronBallman]

Perhaps silly question: did you consider a model where you track incomplete types or their checks to ensure they're completed before the end of the TU and issuing a diagnostic for any that remain incomplete at the end of the TU? We do this dance somewhat frequently:

llvm-project/clang/lib/Sema/Sema.cpp

Line 1171 in 0d9cf26

void Sema::ActOnEndOfTranslationUnit() {

[delcypher]

It's not a silly question at all. This is not a model I explored in much depth.

One reason to not do it that way is that our current implementation requires that "uses" of pointers annotated with __counted_by have a complete pointee type at the "use" site because in some cases we represent bounds checks at the AST level and IIRC the creation of the AST node that represents the bounds check needs to know the pointee size.

This is a design decision that could (and probably should be) revisited at some point but I don't think this is the right time to do it.

As noted in the PR description there is a case with tentative definitions that isn't handled correctly right now and would require waiting until the end of the translation unit to handle correctly. It's been omitted because it's not something we've needed so we've simply not implemented it yet.

[Sirraide]

We could revisit that at some point, but I don’t think that needs to be part of this pr yeah.

[AaronBallman]

Yeah, it doesn't seem like we need a function reference here. If we do, this could be changed to a templated function which accepts a callable, and that would still be better than using a function_ref.

[AaronBallman]

IMO, this seems slightly problematic for API designers, which is why I was asking about delaying checks to the end of the TU above. For someone designing an API with poor testing, they can miss the fact that this attribute is incorrect because nothing ever completes the type in this TU; I think both of these cases should still be diagnosed.

WDYT?

[delcypher]

The current design is very much intentional.

Yes it's potentially problematic for API designers but most people are not API designers. For people who are not API designers why should consumers of that header need to care about an incomplete type if they never try to use it?

The current design also allows for something like:

// public_header.h

// Public and  implementation use same struct representation for `Handle`.
struct OpaqueData;
struct Handle {
  struct OpaqueData* __counted_by(count) Objects;
  size_t count;
 };

// library_internals.c

// Only the library internals knows how `OpaqueData` is defined.
struct OpaqueData {
  int real_data;
};

struct Handle alloc_handle(size_t num_objects) {
  struct Handle h;
  h.Objects = malloc(num_objects);
  h.count = num_objects

 if (!h.Objects)
  h.count = 0;

 return h;
}

That said we should try to help out API designers where possible. We could consider potentially emitting warnings (that have a sane suppression mechanism) for __counted_by pointers where the pointee never gets completed in a TU. However, if we really want that I think that belongs in a separate PR.

@rapidsna Do you have anything to add here?

[delcypher]

@AaronBallman I should have also noted the current design also allows for headers that deliberately don't provide complete types to reduce compile times which I believe is fairly common practice.

E.g.

// some_data.h

struct SomeData {
 int value;
};

// list_handle.h

// Forward declare to avoid needing to include `some_data.h`
struct SomeData;

struct ListHandle {
  struct SomeData* __counted_by(count) items;
  size_t count;
};

[Sirraide]

Yeah, that makes sense. Being able to benefit from the attribute in the implementation of a library without requiring a type to be complete at the API level seems like a reasonable use case.

[delcypher]

@Sirraide Thanks for your review and sorry for the delay in getting back to this. I've tried to address as much as possible and left comment on things that still need discussion. Please take a look when time permits.

[delcypher]

@erichkeane @AaronBallman @Sirraide I've addressed everything I can.

I seem to be running into a problem where some of my replies to comments are not showing up in the conversation tab. This seems to be a known problem with GitHub reviews :( . You may have to dig into the Files tab to see some of my replies to comments :(

[delcypher]

This comment is not showing up in the conversation view for some reason.

[Sirraide]

I seem to be running into a problem where some of my replies to comments are not showing up in the conversation tab.

Yeah, I prefer looking at the flood of emails I get from GitHub, because those pretty consistently do include every comment.

[Sirraide]

Alright, there are a few minor things I pointed out here, but nothing major. I think we should be able to land this soon barring any unforeseen complications. ;Þ

(I also replied to a few comment threads, but that was before I did the actual review so it isn’t shown here unfortunately...)

[delcypher]

@Sirraide Apologies for taking so long to get back to this PR. I think I've addressed all your feedback now. Is this ok to land now?

[Sirraide]

I think this should be fine at this point. I’m having a bit of a hard time trying to figure out what changed since my last review because of the force-pushing, but so long as you haven’t made any other major changes since then then this ltgm now.

[Sirraide]

We could revisit that at some point, but I don’t think that needs to be part of this pr yeah.

[delcypher]

I think this should be fine at this point. I’m having a bit of a hard time trying to figure out what changed since my last review because of the force-pushing, but so long as you haven’t made any other major changes since then then this ltgm now.

Sorry about that. This PR had become so stale that there merge conflicts and the easiest way to fix it was to squash because the conflicts were with things I later removed.

Thanks for reviewing and approving. I'm going to cherry-pick this into the Swiftlang fork of llvm-project first and fix everything up (there's stuff that wasn't upstreamed that's going to conflict with this PR). Once I've done that I'll merge this.

[Sirraide]

Sorry about that. This PR had become so stale that there merge conflicts and the easiest way to fix it was to squash because the conflicts were with things I later removed.

I see, yeah, it happens; that’s unfortunate...

[delcypher]

I had to tweak some minor issues with the diagnostic text that I discovered when trying to cherry pick this into our fork. I've fixed them and I'm going to land this now.