[BoundsSafety] get counted_by spelling from source range

This replaces the count expr source location used for
note_counted_by_consider_using_sized_by with the actual attribute
location fetched from TypeLof, except in the rare case where we couldn't
find a TypeLoc. It also attaches a fix-it hint, now that we have the
proper source range.

Fixes #113585

Author: hnrklssn

clang/include/clang/AST/TypeLoc.h

@@ -1304,6 +1304,7 @@ class ObjCInterfaceTypeLoc : public ConcreteTypeLoc<ObjCObjectTypeLoc,
   }
 };
 
+class Sema;
 struct BoundsAttributedLocInfo {
   SourceRange Range;
 };
@@ -1321,6 +1322,9 @@ class BoundsAttributedTypeLoc
   }
   SourceRange getAttrRange() const { return getLocalData()->Range; }
 
+  StringRef getAttrNameAsWritten(Sema &S) const;
+  SourceRange getAttrNameRange(Sema &S) const;
+
   unsigned getLocalDataSize() const { return sizeof(BoundsAttributedLocInfo); }
 };
 

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