[BoundsSafety] build TypeLoc for CountAttributedType

This adds attribute SourceRange to CountAttributedTypeLoc and populates
it. The following commit will change diagnostics to use this
information.

Fixes #113582

Author: hnrklssn

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,18 +1304,24 @@ class ObjCInterfaceTypeLoc : public ConcreteTypeLoc<ObjCObjectTypeLoc,
   }
 };
 
-struct BoundsAttributedLocInfo {};
+struct BoundsAttributedLocInfo {
+  SourceRange Range;
+};
 class BoundsAttributedTypeLoc
     : public ConcreteTypeLoc<UnqualTypeLoc, BoundsAttributedTypeLoc,
                              BoundsAttributedType, BoundsAttributedLocInfo> {
 public:
   TypeLoc getInnerLoc() const { return getInnerTypeLoc(); }
   QualType getInnerType() const { return getTypePtr()->desugar(); }
   void initializeLocal(ASTContext &Context, SourceLocation Loc) {
-    // nothing to do
+    setAttrRange({Loc, Loc});
+  }
+  void setAttrRange(SourceRange Range) {
+    getLocalData()->Range = Range;
   }
-  // LocalData is empty and TypeLocBuilder doesn't handle DataSize 1.
-  unsigned getLocalDataSize() const { return 0; }
+  SourceRange getAttrRange() const { return getLocalData()->Range; }
+
+  unsigned getLocalDataSize() const { return sizeof(BoundsAttributedLocInfo); }
 };
 
 class CountAttributedTypeLoc final
@@ -1325,8 +1332,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 {

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

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,