Implement TreeTransform for CountAttributedType; address other style suggestions

Author: rapidsna

clang/include/clang/AST/Type.h

@@ -2007,14 +2007,14 @@ class alignas(TypeAlignment) Type : public ExtQualsTypeCommonBase {
     LLVM_PREFERRED_TYPE(TypeBitfields)
     unsigned : NumTypeBits;
 
-    /// The limit is 15.
     static constexpr unsigned NumCoupledDeclsBits = 4;
     unsigned NumCoupledDecls : NumCoupledDeclsBits;
     LLVM_PREFERRED_TYPE(bool)
     unsigned CountInBytes : 1;
     LLVM_PREFERRED_TYPE(bool)
     unsigned OrNull : 1;
   };
+  static_assert(sizeof(CountAttributedTypeBitfields) <= sizeof(unsigned));
 
   union {
     TypeBitfields TypeBits;
@@ -2279,6 +2279,7 @@ class alignas(TypeAlignment) Type : public ExtQualsTypeCommonBase {
   bool isFunctionProtoType() const { return getAs<FunctionProtoType>(); }
   bool isPointerType() const;
   bool isAnyPointerType() const;   // Any C pointer or ObjC object pointer
+  bool isCountAttributedType() const;
   bool isBlockPointerType() const;
   bool isVoidPointerType() const;
   bool isReferenceType() const;

clang/include/clang/AST/TypeLoc.h

@@ -1125,13 +1125,11 @@ class BoundsAttributedTypeLoc
     : public ConcreteTypeLoc<UnqualTypeLoc, BoundsAttributedTypeLoc,
                              BoundsAttributedType, BoundsAttributedLocInfo> {
 public:
-  TypeLoc getInnerLoc() const { return this->getInnerTypeLoc(); }
-  QualType getInnerType() const { return this->getTypePtr()->desugar(); }
+  TypeLoc getInnerLoc() const { return getInnerTypeLoc(); }
+  QualType getInnerType() const { return getTypePtr()->desugar(); }
   void initializeLocal(ASTContext &Context, SourceLocation Loc) {
     // nothing to do
   }
-  unsigned getLocalDataAlignment() const { return 1; }
-  unsigned getLocalDataSize() const { return 0; }
 };
 
 class CountAttributedTypeLoc final

clang/lib/AST/ASTContext.cpp

@@ -13270,9 +13270,9 @@ static QualType getCommonSugarTypeNode(ASTContext &Ctx, const Type *X,
       return QualType();
     if (DX->isOrNull() != DY->isOrNull())
       return QualType();
-    const auto CEX = DX->getCountExpr();
-    const auto CEY = DY->getCountExpr();
-    const auto CDX = DX->getCoupledDecls();
+    Expr *CEX = DX->getCountExpr();
+    Expr *CEY = DY->getCountExpr();
+    llvm::ArrayRef<clang::TypeCoupledDeclRefInfo> CDX = DX->getCoupledDecls();
     if (Ctx.hasSameExpr(CEX, CEY))
       return Ctx.getCountAttributedType(Ctx.getQualifiedType(Underlying), CEX,
                                         DX->isCountInBytes(), DX->isOrNull(),
@@ -13281,9 +13281,9 @@ static QualType getCommonSugarTypeNode(ASTContext &Ctx, const Type *X,
       return QualType();
     // Two declarations with the same integer constant may still differ in their
     // expression pointers, so we need to evaluate them.
-    const auto VX = CEX->getIntegerConstantExpr(Ctx);
-    const auto VY = CEY->getIntegerConstantExpr(Ctx);
-    if (*VX != *VY)
+    llvm::APSInt VX = *CEX->getIntegerConstantExpr(Ctx);
+    llvm::APSInt VY = *CEY->getIntegerConstantExpr(Ctx);
+    if (VX != VY)
       return QualType();
     return Ctx.getCountAttributedType(Ctx.getQualifiedType(Underlying), CEX,
                                       DX->isCountInBytes(), DX->isOrNull(),

clang/lib/AST/Type.cpp

@@ -386,10 +386,8 @@ void DependentBitIntType::Profile(llvm::FoldingSetNodeID &ID,
 }
 
 bool BoundsAttributedType::referencesFieldDecls() const {
-  for (const auto &Decl : dependent_decls())
-    if (isa<FieldDecl>(Decl.getDecl()))
-      return true;
-  return false;
+  return llvm::any_of(dependent_decls(), [](const TypeCoupledDeclRefInfo &Info) {
+    return isa<FieldDecl>(Info.getDecl()); });
 }
 
 void CountAttributedType::Profile(llvm::FoldingSetNodeID &ID,
@@ -669,6 +667,10 @@ bool Type::isScopedEnumeralType() const {
   return false;
 }
 
+bool Type::isCountAttributedType() const {
+  return getAs<CountAttributedType>();
+}
+
 const ComplexType *Type::getAsComplexIntegerType() const {
   if (const auto *Complex = getAs<ComplexType>())
     if (Complex->getElementType()->isIntegerType())

clang/lib/CodeGen/CGBuiltin.cpp

@@ -858,7 +858,7 @@ static unsigned CountCountedByAttrs(const RecordDecl *RD) {
 
   for (const Decl *D : RD->decls()) {
     if (const auto *FD = dyn_cast<FieldDecl>(D);
-        FD && FD->getType()->getAs<CountAttributedType>()) {
+        FD && FD->getType()->isCountAttributedType()) {
       return ++Num;
     }
 
@@ -975,7 +975,7 @@ CodeGenFunction::emitFlexibleArrayMemberSize(const Expr *E, unsigned Type,
       FindFlexibleArrayMemberField(Ctx, OuterRD, FAMName, Offset);
   Offset = Ctx.toCharUnitsFromBits(Offset).getQuantity();
 
-  if (!FAMDecl || !FAMDecl->getType()->getAs<CountAttributedType>())
+  if (!FAMDecl || !FAMDecl->getType()->isCountAttributedType())
     // No flexible array member found or it doesn't have the "counted_by"
     // attribute.
     return nullptr;

clang/lib/CodeGen/CGExpr.cpp

@@ -1139,12 +1139,12 @@ const FieldDecl *CodeGenFunction::FindCountedByField(const FieldDecl *FD) {
   if (!FD)
     return nullptr;
 
-  auto *CAT = FD->getType()->getAs<CountAttributedType>();
+  const auto *CAT = FD->getType()->getAs<CountAttributedType>();
   if (!CAT)
     return nullptr;
 
-  auto *CountDRE = cast<DeclRefExpr>(CAT->getCountExpr());
-  auto *CountDecl = CountDRE->getDecl();
+  const auto *CountDRE = cast<DeclRefExpr>(CAT->getCountExpr());
+  const auto *CountDecl = CountDRE->getDecl();
   if (const auto *IFD = dyn_cast<IndirectFieldDecl>(CountDecl))
     CountDecl = IFD->getAnonField();
 
@@ -4240,7 +4240,7 @@ LValue CodeGenFunction::EmitArraySubscriptExpr(const ArraySubscriptExpr *E,
       if (const auto *ME = dyn_cast<MemberExpr>(Array);
           ME &&
           ME->isFlexibleArrayMemberLike(getContext(), StrictFlexArraysLevel) &&
-          ME->getMemberDecl()->getType()->getAs<CountAttributedType>()) {
+          ME->getMemberDecl()->getType()->isCountAttributedType()) {
         const FieldDecl *FAMDecl = dyn_cast<FieldDecl>(ME->getMemberDecl());
         if (const FieldDecl *CountFD = FindCountedByField(FAMDecl)) {
           if (std::optional<int64_t> Diff =

clang/lib/Sema/SemaDeclAttr.cpp

@@ -8578,8 +8578,7 @@ static void handleCountedByAttrField(Sema &S, Decl *D, const ParsedAttr &AL) {
   if (CheckCountExpr(S, FD, CountExpr, Decls))
     return;
 
-  QualType CAT = S.BuildCountAttributedArrayType(
-      FD->getType(), CountExpr, llvm::ArrayRef(Decls.begin(), Decls.end()));
+  QualType CAT = S.BuildCountAttributedArrayType(FD->getType(), CountExpr);
   FD->setType(CAT);
 }
 

clang/lib/Sema/SemaType.cpp

@@ -6494,6 +6494,9 @@ namespace {
     void VisitAttributedTypeLoc(AttributedTypeLoc TL) {
       fillAttributedTypeLoc(TL, State);
     }
+    void VisitCountAttributedTypeLoc(CountAttributedTypeLoc TL) {
+      // nothing
+    }
     void VisitBTFTagAttributedTypeLoc(BTFTagAttributedTypeLoc TL) {
       // nothing
     }
@@ -9745,11 +9748,18 @@ QualType Sema::BuildTypeofExprType(Expr *E, TypeOfKind Kind) {
   return Context.getTypeOfExprType(E, Kind);
 }
 
-QualType Sema::BuildCountAttributedArrayType(
-    QualType WrappedTy, Expr *CountExpr,
-    llvm::ArrayRef<TypeCoupledDeclRefInfo> Decls) {
+static void BuildTypeCoupledDecls(Expr *E,
+                                  llvm::SmallVectorImpl<TypeCoupledDeclRefInfo> &Decls) {
+  // Currently, 'counted_by' only allows direct DeclRefExpr to FieldDecl.
+  auto *CountDecl = cast<DeclRefExpr>(E)->getDecl();
+  Decls.push_back(TypeCoupledDeclRefInfo(CountDecl, /*IsDref*/ false));
+}
+
+QualType Sema::BuildCountAttributedArrayType(QualType WrappedTy, Expr *CountExpr) {
   assert(WrappedTy->isIncompleteArrayType());
 
+  llvm::SmallVector<TypeCoupledDeclRefInfo, 1> Decls;
+  BuildTypeCoupledDecls(CountExpr, Decls);
   /// When the resulting expression is invalid, we still create the AST using
   /// the original count expression for the sake of AST dump.
   return Context.getCountAttributedType(

clang/lib/Sema/TreeTransform.h

@@ -7264,8 +7264,29 @@ QualType TreeTransform<Derived>::TransformAttributedType(TypeLocBuilder &TLB,
 template <typename Derived>
 QualType TreeTransform<Derived>::TransformCountAttributedType(
     TypeLocBuilder &TLB, CountAttributedTypeLoc TL) {
-  // TODO
-  llvm_unreachable("Unexpected TreeTransform for CountAttributedType");
+  const CountAttributedType *OldTy = TL.getTypePtr();
+  QualType InnerTy = getDerived().TransformType(TLB, TL.getInnerLoc());
+  if (InnerTy.isNull())
+    return QualType();
+
+  Expr *OldCount = TL.getCountExpr();
+  Expr *NewCount = nullptr;
+  if (OldCount) {
+    ExprResult CountResult = getDerived().TransformExpr(OldCount);
+    if (CountResult.isInvalid())
+      return QualType();
+    NewCount = CountResult.get();
+  }
+
+  QualType Result = TL.getType();
+  if (getDerived().AlwaysRebuild() || InnerTy != OldTy->desugar() ||
+      OldCount != NewCount) {
+    // Currently, CountAttributedType can only wrap incomplete array types.
+    Result = SemaRef.BuildCountAttributedArrayType(InnerTy, NewCount);
+  }
+
+  TLB.push<CountAttributedTypeLoc>(Result);
+  return Result;
 }
 
 template <typename Derived>

clang/lib/Serialization/ASTReader.cpp

@@ -6996,7 +6996,7 @@ void TypeLocReader::VisitAttributedTypeLoc(AttributedTypeLoc TL) {
 }
 
 void TypeLocReader::VisitCountAttributedTypeLoc(CountAttributedTypeLoc TL) {
-  // nothing to do
+  // Nothing to do
 }
 
 void TypeLocReader::VisitBTFTagAttributedTypeLoc(BTFTagAttributedTypeLoc TL) {

clang/lib/Serialization/ASTWriter.cpp

@@ -515,7 +515,7 @@ void TypeLocWriter::VisitAttributedTypeLoc(AttributedTypeLoc TL) {
 }
 
 void TypeLocWriter::VisitCountAttributedTypeLoc(CountAttributedTypeLoc TL) {
-  // nothing to do
+  // Nothing to do
 }
 
 void TypeLocWriter::VisitBTFTagAttributedTypeLoc(BTFTagAttributedTypeLoc TL) {