[BoundsSafety] build TypeLoc for CountAttributedType

[hnrklssn]

This adds attribute SourceRange to CountAttributedTypeLoc and populates it, then uses it to replace a not-quite-correct source location fetched from the count expression.

Fixes #113582
Fixes #113585

[llvmbot]

@llvm/pr-subscribers-clang

Author: Henrik G. Olsson (hnrklssn)

Changes

This adds attribute SourceRange to CountAttributedTypeLoc and populates it, then uses it to replace a not-quite-correct source location fetched from the count expression.

---

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

4 Files Affected:

• (modified) clang/include/clang/AST/TypeLoc.h (+15-6)
• (modified) clang/lib/AST/TypeLoc.cpp (-4)
• (modified) clang/lib/Sema/SemaBoundsSafety.cpp (+105-16)
• (modified) clang/lib/Sema/SemaDeclAttr.cpp (+9)

diff --git a/clang/include/clang/AST/TypeLoc.h b/clang/include/clang/AST/TypeLoc.h
index 2cefaa9611c98..0cf17c463d6cd 100644
--- a/clang/include/clang/AST/TypeLoc.h
+++ b/clang/include/clang/AST/TypeLoc.h
@@ -19,6 +19,7 @@
 #include "clang/AST/NestedNameSpecifierBase.h"
 #include "clang/AST/TemplateBase.h"
 #include "clang/AST/TypeBase.h"
+#include "clang/Basic/IdentifierTable.h"
 #include "clang/Basic/LLVM.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/Specifiers.h"
@@ -1303,7 +1304,10 @@ class ObjCInterfaceTypeLoc : public ConcreteTypeLoc<ObjCObjectTypeLoc,
   }
 };
 
-struct BoundsAttributedLocInfo {};
+class Sema;
+struct BoundsAttributedLocInfo {
+  SourceRange Range;
+};
 class BoundsAttributedTypeLoc
     : public ConcreteTypeLoc<UnqualTypeLoc, BoundsAttributedTypeLoc,
                              BoundsAttributedType, BoundsAttributedLocInfo> {
@@ -1311,10 +1315,17 @@ class BoundsAttributedTypeLoc
   TypeLoc getInnerLoc() const { return getInnerTypeLoc(); }
   QualType getInnerType() const { return getTypePtr()->desugar(); }
   void initializeLocal(ASTContext &Context, SourceLocation Loc) {
-    // nothing to do
+    setAttrRange({Loc, Loc});
   }
-  // LocalData is empty and TypeLocBuilder doesn't handle DataSize 1.
-  unsigned getLocalDataSize() const { return 0; }
+  void setAttrRange(SourceRange Range) {
+    getLocalData()->Range = Range;
+  }
+  SourceRange getAttrRange() const { return getLocalData()->Range; }
+
+  StringRef getAttrNameAsWritten(Sema &S) const;
+  SourceRange getAttrNameRange(Sema &S) const;
+
+  unsigned getLocalDataSize() const { return sizeof(BoundsAttributedLocInfo); }
 };
 
 class CountAttributedTypeLoc final
@@ -1325,8 +1336,6 @@ class CountAttributedTypeLoc final
   Expr *getCountExpr() const { return getTypePtr()->getCountExpr(); }
   bool isCountInBytes() const { return getTypePtr()->isCountInBytes(); }
   bool isOrNull() const { return getTypePtr()->isOrNull(); }
-
-  SourceRange getLocalSourceRange() const;
 };
 
 struct MacroQualifiedLocInfo {
diff --git a/clang/lib/AST/TypeLoc.cpp b/clang/lib/AST/TypeLoc.cpp
index f54ccf0932bc7..66f9e3c67df84 100644
--- a/clang/lib/AST/TypeLoc.cpp
+++ b/clang/lib/AST/TypeLoc.cpp
@@ -590,10 +590,6 @@ SourceRange AttributedTypeLoc::getLocalSourceRange() const {
   return getAttr() ? getAttr()->getRange() : SourceRange();
 }
 
-SourceRange CountAttributedTypeLoc::getLocalSourceRange() const {
-  return getCountExpr() ? getCountExpr()->getSourceRange() : SourceRange();
-}
-
 SourceRange BTFTagAttributedTypeLoc::getLocalSourceRange() const {
   return getAttr() ? getAttr()->getRange() : SourceRange();
 }
diff --git a/clang/lib/Sema/SemaBoundsSafety.cpp b/clang/lib/Sema/SemaBoundsSafety.cpp
index de9adf8ef5a1b..a699c9d116b15 100644
--- a/clang/lib/Sema/SemaBoundsSafety.cpp
+++ b/clang/lib/Sema/SemaBoundsSafety.cpp
@@ -11,6 +11,14 @@
 /// (e.g. `counted_by`)
 ///
 //===----------------------------------------------------------------------===//
+#include "clang/AST/Decl.h"
+#include "clang/AST/Expr.h"
+#include "clang/AST/StmtVisitor.h"
+#include "clang/AST/TypeBase.h"
+#include "clang/AST/TypeLoc.h"
+#include "clang/Basic/Diagnostic.h"
+#include "clang/Basic/SourceLocation.h"
+#include "clang/Basic/SourceManager.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Sema/Initialization.h"
 #include "clang/Sema/Sema.h"
@@ -231,9 +239,69 @@ bool Sema::CheckCountedByAttrOnField(FieldDecl *FD, Expr *E, bool CountInBytes,
   return false;
 }
 
+// FIXME: for some reason diagnostics highlight the end character, while
+// getSourceText() does not include the end character.
+static SourceRange getAttrNameRangeImpl(Sema &S, SourceLocation Begin,
+                                        bool IsForDiagnostics) {
+  SourceManager &SM = S.getSourceManager();
+  SourceLocation TokenStart = Begin;
+  while (TokenStart.isMacroID())
+    TokenStart = SM.getImmediateExpansionRange(TokenStart).getBegin();
+  unsigned Offset = IsForDiagnostics ? 1 : 0;
+  SourceLocation End = S.getLocForEndOfToken(TokenStart, Offset);
+  return {TokenStart, End};
+}
+
+StringRef BoundsAttributedTypeLoc::getAttrNameAsWritten(Sema &S) const {
+  SourceRange Range = getAttrNameRangeImpl(S, getAttrRange().getBegin(), false);
+  CharSourceRange NameRange = CharSourceRange::getCharRange(Range);
+  return Lexer::getSourceText(NameRange, S.getSourceManager(), S.getLangOpts());
+}
+
+SourceRange BoundsAttributedTypeLoc::getAttrNameRange(Sema &S) const {
+  return getAttrNameRangeImpl(S, getAttrRange().getBegin(), true);
+}
+
+static TypeSourceInfo *getTSI(const Decl *D) {
+  if (const auto* DD = dyn_cast<DeclaratorDecl>(D)) {
+    return DD->getTypeSourceInfo();
+  }
+  return nullptr;
+}
+
+struct TypeLocFinder : public ConstStmtVisitor<TypeLocFinder, TypeLoc> {
+  TypeLoc VisitParenExpr(const ParenExpr* E) {
+    return Visit(E->getSubExpr());
+  }
+
+  TypeLoc VisitDeclRefExpr(const DeclRefExpr *E) {
+    return getTSI(E->getDecl())->getTypeLoc();
+  }
+
+  TypeLoc VisitMemberExpr(const MemberExpr *E) {
+    return getTSI(E->getMemberDecl())->getTypeLoc();
+  }
+
+  TypeLoc VisitExplicitCastExpr(const ExplicitCastExpr *E) {
+    return E->getTypeInfoAsWritten()->getTypeLoc();
+  }
+
+  TypeLoc VisitCallExpr(const CallExpr *E) {
+    if (const auto *D = E->getCalleeDecl()) {
+      FunctionTypeLoc FTL = getTSI(D)->getTypeLoc().getAs<FunctionTypeLoc>();
+      if (FTL.isNull()) {
+        return FTL;
+      }
+      return FTL.getReturnLoc();
+    }
+    return {};
+  }
+};
+
 static void EmitIncompleteCountedByPointeeNotes(Sema &S,
                                                 const CountAttributedType *CATy,
-                                                NamedDecl *IncompleteTyDecl) {
+                                                NamedDecl *IncompleteTyDecl,
+                                                TypeLoc TL) {
   assert(IncompleteTyDecl == nullptr || isa<TypeDecl>(IncompleteTyDecl));
 
   if (IncompleteTyDecl) {
@@ -253,20 +321,36 @@ static void EmitIncompleteCountedByPointeeNotes(Sema &S,
         << CATy->getPointeeType();
   }
 
-  // Suggest using __sized_by(_or_null) instead of __counted_by(_or_null) as
-  // __sized_by(_or_null) doesn't have the complete type restriction.
-  //
-  // We use the source range of the expression on the CountAttributedType as an
-  // approximation for the source range of the attribute. This isn't quite right
-  // but isn't easy to fix right now.
-  //
-  // TODO: Implement logic to find the relevant TypeLoc for the attribute and
-  // get the SourceRange from that (#113582).
-  //
-  // TODO: We should emit a fix-it here.
-  SourceRange AttrSrcRange = CATy->getCountExpr()->getSourceRange();
+  CountAttributedTypeLoc CATL;
+  if (!TL.isNull())
+    CATL = TL.getAs<CountAttributedTypeLoc>();
+
+  if (CATL.isNull()) {
+    // Fall back to pointing to the count expr - not great, but close enough.
+    // This should happen rarely, if ever.
+    S.Diag(CATy->getCountExpr()->getExprLoc(),
+           diag::note_counted_by_consider_using_sized_by)
+        << CATy->isOrNull();
+    return;
+  }
+  SourceRange AttrSrcRange = CATL.getAttrNameRange(S);
+
+  StringRef Spelling = CATL.getAttrNameAsWritten(S);
+  StringRef FixedSpelling;
+  if (Spelling == "__counted_by")
+    FixedSpelling = "__sized_by";
+  else if (Spelling == "counted_by")
+    FixedSpelling = "sized_by";
+  else if (Spelling == "__counted_by_or_null")
+    FixedSpelling = "__sized_by_or_null";
+  else if (Spelling == "counted_by_or_null")
+    FixedSpelling = "sized_by_or_null";
+  FixItHint Fix;
+  if (!FixedSpelling.empty())
+    Fix = FixItHint::CreateReplacement(AttrSrcRange, FixedSpelling);
+
   S.Diag(AttrSrcRange.getBegin(), diag::note_counted_by_consider_using_sized_by)
-      << CATy->isOrNull() << AttrSrcRange;
+      << CATy->isOrNull() << AttrSrcRange << Fix;
 }
 
 static std::tuple<const CountAttributedType *, QualType>
@@ -335,7 +419,11 @@ static bool CheckAssignmentToCountAttrPtrWithIncompletePointeeTy(
       << CATy->getAttributeName(/*WithMacroPrefix=*/true) << PointeeTy
       << CATy->isOrNull() << RHSExpr->getSourceRange();
 
-  EmitIncompleteCountedByPointeeNotes(S, CATy, IncompleteTyDecl);
+  TypeLoc TL;
+  if (TypeSourceInfo *TSI = getTSI(Assignee))
+    TL = TSI->getTypeLoc();
+
+  EmitIncompleteCountedByPointeeNotes(S, CATy, IncompleteTyDecl, TL);
   return false; // check failed
 }
 
@@ -408,7 +496,8 @@ bool Sema::BoundsSafetyCheckUseOfCountAttrPtr(const Expr *E) {
       << CATy->getAttributeName(/*WithMacroPrefix=*/true) << CATy->isOrNull()
       << E->getSourceRange();
 
-  EmitIncompleteCountedByPointeeNotes(*this, CATy, IncompleteTyDecl);
+  TypeLoc TL = TypeLocFinder().Visit(E);
+  EmitIncompleteCountedByPointeeNotes(*this, CATy, IncompleteTyDecl, TL);
   return false;
 }
 
diff --git a/clang/lib/Sema/SemaDeclAttr.cpp b/clang/lib/Sema/SemaDeclAttr.cpp
index a9e7b44ac9d73..1753261822755 100644
--- a/clang/lib/Sema/SemaDeclAttr.cpp
+++ b/clang/lib/Sema/SemaDeclAttr.cpp
@@ -10,10 +10,12 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "TypeLocBuilder.h"
 #include "clang/AST/APValue.h"
 #include "clang/AST/ASTConsumer.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/ASTMutationListener.h"
+#include "clang/AST/Attr.h"
 #include "clang/AST/CXXInheritance.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclCXX.h"
@@ -24,6 +26,8 @@
 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/Mangle.h"
 #include "clang/AST/Type.h"
+#include "clang/AST/TypeBase.h"
+#include "clang/AST/TypeLoc.h"
 #include "clang/Basic/CharInfo.h"
 #include "clang/Basic/Cuda.h"
 #include "clang/Basic/DarwinSDKInfo.h"
@@ -6580,9 +6584,14 @@ static void handleCountedByAttrField(Sema &S, Decl *D, const ParsedAttr &AL) {
   if (S.CheckCountedByAttrOnField(FD, CountExpr, CountInBytes, OrNull))
     return;
 
+  TypeLocBuilder TLB;
   QualType CAT = S.BuildCountAttributedArrayOrPointerType(
       FD->getType(), CountExpr, CountInBytes, OrNull);
+  TLB.pushFullCopy(FD->getTypeSourceInfo()->getTypeLoc());
+  CountAttributedTypeLoc CATL = TLB.push<CountAttributedTypeLoc>(CAT);
+  CATL.setAttrRange(AL.getRange());
   FD->setType(CAT);
+  FD->setTypeSourceInfo(TLB.getTypeSourceInfo(S.getASTContext(), CAT));
 }
 
 static void handleFunctionReturnThunksAttr(Sema &S, Decl *D,

[github-actions]

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

[hnrklssn]

Just found out about -fdiagnostics-print-source-range-info. I'll add some test cases for source ranges & fix-its tomorrow.

[hnrklssn]

Should be ready for review now :)

[zmodem]

Just nits from my side. Otherwise lgtm.

[rapidsna]

Thanks @hnrklssn! This looks great. Not that I'm suggesting to do right now, but I wonder if we want to embed Attr * in bounds-attributed types eventually to preserve all attribute related stuffs, and whether that would also remove the needs for custom handing of ranges like this.

[hnrklssn]

I think Attr would have the same range (and MacroII seems to have basically the same range as well), since it would also get its range from ParsedAttr. I'm not sure whether there's really that much else we would want in Attr, and since we don't currently allocate an Attr for counted_by it felt more sensible to just grab the source range directly. If we figure out a use case for Attr we can always change it later.