[CodeGen] Revamp counted_by calculations (#70606)

bwendling · web-flow · commit bc09ec696209 · 2023-11-09T10:18:17.000-08:00
Break down the counted_by calculations so that they correctly handle
anonymous structs, which are specified internally as IndirectFieldDecls.

Improves the calculation of __bdos on a different field member in the struct.
And also improves support for __bdos in an index into the FAM. If the index
is further out than the length of the FAM, then we return __bdos's "can't
determine the size" value (zero or negative one, depending on type).

Also simplify the code to use helper methods to get the field referenced
by counted_by and the flexible array member itself, which also had some
issues with FAMs in sub-structs.
diff --git a/clang/lib/CodeGen/CGBuiltin.cpp b/clang/lib/CodeGen/CGBuiltin.cpp
@@ -25,6 +25,7 @@
 #include "clang/AST/Attr.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/OSLog.h"
+#include "clang/AST/OperationKinds.h"
 #include "clang/Basic/TargetBuiltins.h"
 #include "clang/Basic/TargetInfo.h"
 #include "clang/Basic/TargetOptions.h"
@@ -858,63 +859,155 @@ CodeGenFunction::emitBuiltinObjectSize(const Expr *E, unsigned Type,
     }
   }
 
+  // LLVM can't handle Type=3 appropriately, and __builtin_object_size shouldn't
+  // evaluate E for side-effects. In either case, we shouldn't lower to
+  // @llvm.objectsize.
+  if (Type == 3 || (!EmittedE && E->HasSideEffects(getContext())))
+    return getDefaultBuiltinObjectSizeResult(Type, ResType);
+
   if (IsDynamic) {
-    LangOptions::StrictFlexArraysLevelKind StrictFlexArraysLevel =
-        getLangOpts().getStrictFlexArraysLevel();
+    // The code generated here calculates the size of a struct with a flexible
+    // array member that uses the counted_by attribute. There are two instances
+    // we handle:
+    //
+    //       struct s {
+    //         unsigned long flags;
+    //         int count;
+    //         int array[] __attribute__((counted_by(count)));
+    //       }
+    //
+    //   1) bdos of the flexible array itself:
+    //
+    //     __builtin_dynamic_object_size(p->array, 1) ==
+    //         p->count * sizeof(*p->array)
+    //
+    //   2) bdos of a pointer into the flexible array:
+    //
+    //     __builtin_dynamic_object_size(&p->array[42], 1) ==
+    //         (p->count - 42) * sizeof(*p->array)
+    //
+    //   2) bdos of the whole struct, including the flexible array:
+    //
+    //     __builtin_dynamic_object_size(p, 1) ==
+    //        max(sizeof(struct s),
+    //            offsetof(struct s, array) + p->count * sizeof(*p->array))
+    //
     const Expr *Base = E->IgnoreParenImpCasts();
-
-    if (FieldDecl *FD = FindCountedByField(Base, StrictFlexArraysLevel)) {
-      const auto *ME = dyn_cast<MemberExpr>(Base);
-      llvm::Value *ObjectSize = nullptr;
-
-      if (!ME) {
-        const auto *DRE = dyn_cast<DeclRefExpr>(Base);
-        ValueDecl *VD = nullptr;
-
-        ObjectSize = ConstantInt::get(
-            ResType,
-            getContext().getTypeSize(DRE->getType()->getPointeeType()) / 8,
-            true);
-
-        if (auto *RD = DRE->getType()->getPointeeType()->getAsRecordDecl())
-          VD = RD->getLastField();
-
-        Expr *ICE = ImplicitCastExpr::Create(
-            getContext(), DRE->getType(), CK_LValueToRValue,
-            const_cast<Expr *>(cast<Expr>(DRE)), nullptr, VK_PRValue,
-            FPOptionsOverride());
-        ME = MemberExpr::CreateImplicit(getContext(), ICE, true, VD,
-                                        VD->getType(), VK_LValue, OK_Ordinary);
+    const Expr *Idx = nullptr;
+    if (const auto *UO = dyn_cast<UnaryOperator>(Base);
+        UO && UO->getOpcode() == UO_AddrOf) {
+      if (const auto *ASE =
+              dyn_cast<ArraySubscriptExpr>(UO->getSubExpr()->IgnoreParens())) {
+        Base = ASE->getBase();
+        Idx = ASE->getIdx()->IgnoreParenImpCasts();
+
+        if (const auto *IL = dyn_cast<IntegerLiteral>(Idx);
+            IL && !IL->getValue().getSExtValue())
+          Idx = nullptr;
       }
+    }
 
-      // At this point, we know that \p ME is a flexible array member.
-      const auto *ArrayTy = getContext().getAsArrayType(ME->getType());
-      unsigned Size = getContext().getTypeSize(ArrayTy->getElementType());
-
-      llvm::Value *CountField =
-          EmitAnyExprToTemp(MemberExpr::CreateImplicit(
-                                getContext(), const_cast<Expr *>(ME->getBase()),
-                                ME->isArrow(), FD, FD->getType(), VK_LValue,
-                                OK_Ordinary))
-              .getScalarVal();
+    if (const ValueDecl *CountedByFD = FindCountedByField(Base)) {
+      bool IsSigned = CountedByFD->getType()->isSignedIntegerType();
+      const RecordDecl *OuterRD =
+          CountedByFD->getDeclContext()->getOuterLexicalRecordContext();
+      ASTContext &Ctx = getContext();
+
+      // Load the counted_by field.
+      const Expr *CountedByExpr = BuildCountedByFieldExpr(Base, CountedByFD);
+      Value *CountedByInst = EmitAnyExprToTemp(CountedByExpr).getScalarVal();
+      llvm::Type *CountedByTy = CountedByInst->getType();
+
+      if (Idx) {
+        // There's an index into the array. Remove it from the count.
+        bool IdxSigned = Idx->getType()->isSignedIntegerType();
+        Value *IdxInst = EmitAnyExprToTemp(Idx).getScalarVal();
+        IdxInst = IdxSigned ? Builder.CreateSExtOrTrunc(IdxInst, CountedByTy)
+                            : Builder.CreateZExtOrTrunc(IdxInst, CountedByTy);
+
+        // If the index is negative, don't subtract it from the counted_by
+        // value. The pointer is pointing to something before the FAM.
+        IdxInst = Builder.CreateNeg(IdxInst, "", !IdxSigned, IdxSigned);
+        CountedByInst =
+            Builder.CreateAdd(CountedByInst, IdxInst, "", !IsSigned, IsSigned);
+      }
 
-      llvm::Value *Mul = Builder.CreateMul(
-          CountField, llvm::ConstantInt::get(CountField->getType(), Size / 8));
-      Mul = Builder.CreateZExtOrTrunc(Mul, ResType);
+      // Get the size of the flexible array member's base type.
+      const ValueDecl *FAMDecl = nullptr;
+      if (const auto *ME = dyn_cast<MemberExpr>(Base)) {
+        const LangOptions::StrictFlexArraysLevelKind StrictFlexArraysLevel =
+            getLangOpts().getStrictFlexArraysLevel();
+        if (const ValueDecl *MD = ME->getMemberDecl();
+            MD && Decl::isFlexibleArrayMemberLike(
+                      Ctx, MD, MD->getType(), StrictFlexArraysLevel,
+                      /*IgnoreTemplateOrMacroSubstitution=*/true))
+          // Base is referencing the FAM itself.
+          FAMDecl = MD;
+      }
 
-      if (ObjectSize)
-        return Builder.CreateAdd(ObjectSize, Mul);
+      if (!FAMDecl)
+        FAMDecl = FindFlexibleArrayMemberField(Ctx, OuterRD);
+
+      assert(FAMDecl && "Can't find the flexible array member field");
+
+      const ArrayType *ArrayTy = Ctx.getAsArrayType(FAMDecl->getType());
+      CharUnits Size = Ctx.getTypeSizeInChars(ArrayTy->getElementType());
+      llvm::Constant *ElemSize =
+          llvm::ConstantInt::get(CountedByTy, Size.getQuantity(), IsSigned);
+
+      // Calculate how large the flexible array member is in bytes.
+      Value *FAMSize =
+          Builder.CreateMul(CountedByInst, ElemSize, "", !IsSigned, IsSigned);
+      FAMSize = IsSigned ? Builder.CreateSExtOrTrunc(FAMSize, ResType)
+                         : Builder.CreateZExtOrTrunc(FAMSize, ResType);
+      Value *Res = FAMSize;
+
+      if (const auto *DRE = dyn_cast<DeclRefExpr>(Base)) {
+        // The whole struct is specificed in the __bdos.
+        const ASTRecordLayout &Layout = Ctx.getASTRecordLayout(OuterRD);
+
+        // Get the offset of the FAM.
+        CharUnits Offset = Ctx.toCharUnitsFromBits(Ctx.getFieldOffset(FAMDecl));
+        llvm::Constant *FAMOffset =
+            ConstantInt::get(ResType, Offset.getQuantity(), IsSigned);
+
+        // max(sizeof(struct s),
+        //     offsetof(struct s, array) + p->count * sizeof(*p->array))
+        Value *OffsetAndFAMSize =
+            Builder.CreateAdd(FAMOffset, Res, "", !IsSigned, IsSigned);
+
+        // Get the full size of the struct.
+        llvm::Constant *SizeofStruct =
+            ConstantInt::get(ResType, Layout.getSize().getQuantity(), IsSigned);
+
+        Res = IsSigned
+                  ? Builder.CreateBinaryIntrinsic(
+                        llvm::Intrinsic::smax, OffsetAndFAMSize, SizeofStruct)
+                  : Builder.CreateBinaryIntrinsic(
+                        llvm::Intrinsic::umax, OffsetAndFAMSize, SizeofStruct);
+      } else if (const auto *ME = dyn_cast<MemberExpr>(Base)) {
+        // Pointing to a place before the FAM. Add the difference to the FAM's
+        // size.
+        if (const ValueDecl *MD = ME->getMemberDecl(); MD != FAMDecl) {
+          CharUnits Offset = Ctx.toCharUnitsFromBits(Ctx.getFieldOffset(MD));
+          CharUnits FAMOffset =
+              Ctx.toCharUnitsFromBits(Ctx.getFieldOffset(FAMDecl));
+
+          Res = Builder.CreateAdd(
+              Res, ConstantInt::get(ResType, FAMOffset.getQuantity() -
+                                                 Offset.getQuantity()));
+        }
+      }
 
-      return Mul;
+      // A negative 'FAMSize' means that the index was greater than the count,
+      // or an improperly set count field. Return -1 (for types 0 and 1) or 0
+      // (for types 2 and 3).
+      return Builder.CreateSelect(
+          Builder.CreateIsNeg(FAMSize),
+          getDefaultBuiltinObjectSizeResult(Type, ResType), Res);
     }
   }
 
-  // LLVM can't handle Type=3 appropriately, and __builtin_object_size shouldn't
-  // evaluate E for side-effects. In either case, we shouldn't lower to
-  // @llvm.objectsize.
-  if (Type == 3 || (!EmittedE && E->HasSideEffects(getContext())))
-    return getDefaultBuiltinObjectSizeResult(Type, ResType);
-
   Value *Ptr = EmittedE ? EmittedE : EmitScalarExpr(E);
   assert(Ptr->getType()->isPointerTy() &&
          "Non-pointer passed to __builtin_object_size?");
diff --git a/clang/lib/CodeGen/CGExpr.cpp b/clang/lib/CodeGen/CGExpr.cpp
@@ -938,14 +938,10 @@ static llvm::Value *getArrayIndexingBound(CodeGenFunction &CGF,
       // Ignore pass_object_size here. It's not applicable on decayed pointers.
     }
 
-    if (FieldDecl *FD = CGF.FindCountedByField(Base, StrictFlexArraysLevel)) {
-      const auto *ME = dyn_cast<MemberExpr>(CE->getSubExpr());
+    if (const ValueDecl *VD = CGF.FindCountedByField(Base)) {
       IndexedType = Base->getType();
-      return CGF
-          .EmitAnyExprToTemp(MemberExpr::CreateImplicit(
-              CGF.getContext(), const_cast<Expr *>(ME->getBase()),
-              ME->isArrow(), FD, FD->getType(), VK_LValue, OK_Ordinary))
-          .getScalarVal();
+      const Expr *E = CGF.BuildCountedByFieldExpr(Base, VD);
+      return CGF.EmitAnyExprToTemp(E).getScalarVal();
     }
   }
 
@@ -960,46 +956,115 @@ static llvm::Value *getArrayIndexingBound(CodeGenFunction &CGF,
   return nullptr;
 }
 
-FieldDecl *CodeGenFunction::FindCountedByField(
-    const Expr *Base,
-    LangOptions::StrictFlexArraysLevelKind StrictFlexArraysLevel) {
-  const ValueDecl *VD = nullptr;
+const Expr *
+CodeGenFunction::BuildCountedByFieldExpr(const Expr *Base,
+                                         const ValueDecl *CountedByVD) {
+  // Find the outer struct expr (i.e. p in p->a.b.c.d).
+  Expr *CountedByExpr = const_cast<Expr *>(Base)->IgnoreParenImpCasts();
+
+  // Work our way up the expression until we reach the DeclRefExpr.
+  while (!isa<DeclRefExpr>(CountedByExpr))
+    if (const auto *ME = dyn_cast<MemberExpr>(CountedByExpr))
+      CountedByExpr = ME->getBase()->IgnoreParenImpCasts();
+
+  // Add back an implicit cast to create the required pr-value.
+  CountedByExpr = ImplicitCastExpr::Create(
+      getContext(), CountedByExpr->getType(), CK_LValueToRValue, CountedByExpr,
+      nullptr, VK_PRValue, FPOptionsOverride());
+
+  if (const auto *IFD = dyn_cast<IndirectFieldDecl>(CountedByVD)) {
+    // The counted_by field is inside an anonymous struct / union. The
+    // IndirectFieldDecl has the correct order of FieldDecls to build this
+    // easily. (Yay!)
+    for (NamedDecl *ND : IFD->chain()) {
+      auto *VD = cast<ValueDecl>(ND);
+      CountedByExpr =
+          MemberExpr::CreateImplicit(getContext(), CountedByExpr,
+                                     CountedByExpr->getType()->isPointerType(),
+                                     VD, VD->getType(), VK_LValue, OK_Ordinary);
+    }
+  } else {
+    CountedByExpr = MemberExpr::CreateImplicit(
+        getContext(), const_cast<Expr *>(CountedByExpr),
+        CountedByExpr->getType()->isPointerType(),
+        const_cast<ValueDecl *>(CountedByVD), CountedByVD->getType(), VK_LValue,
+        OK_Ordinary);
+  }
 
-  Base = Base->IgnoreParenImpCasts();
+  return CountedByExpr;
+}
+
+const ValueDecl *
+CodeGenFunction::FindFlexibleArrayMemberField(ASTContext &Ctx,
+                                              const RecordDecl *RD) {
+  const LangOptions::StrictFlexArraysLevelKind StrictFlexArraysLevel =
+      getLangOpts().getStrictFlexArraysLevel();
+
+  for (const Decl *D : RD->decls()) {
+    if (const auto *VD = dyn_cast<ValueDecl>(D);
+        VD && Decl::isFlexibleArrayMemberLike(
+                  Ctx, VD, VD->getType(), StrictFlexArraysLevel,
+                  /*IgnoreTemplateOrMacroSubstitution=*/true))
+      return VD;
+
+    if (const auto *Record = dyn_cast<RecordDecl>(D))
+      if (const ValueDecl *VD = FindFlexibleArrayMemberField(Ctx, Record))
+        return VD;
+  }
+
+  return nullptr;
+}
+
+const ValueDecl *CodeGenFunction::FindCountedByField(const Expr *Base) {
+  ASTContext &Ctx = getContext();
+  const RecordDecl *OuterRD = nullptr;
+  const FieldDecl *FD = nullptr;
 
-  if (const auto *ME = dyn_cast<MemberExpr>(Base)) {
-    VD = dyn_cast<ValueDecl>(ME->getMemberDecl());
-  } else if (const auto *DRE = dyn_cast<DeclRefExpr>(Base)) {
-    // Pointing to the full structure.
-    VD = dyn_cast<ValueDecl>(DRE->getDecl());
+  Base = Base->IgnoreParenImpCasts();
 
-    QualType Ty = VD->getType();
+  // Get the outer-most lexical RecordDecl.
+  if (const auto *DRE = dyn_cast<DeclRefExpr>(Base)) {
+    QualType Ty = DRE->getDecl()->getType();
     if (Ty->isPointerType())
       Ty = Ty->getPointeeType();
 
     if (const auto *RD = Ty->getAsRecordDecl())
-      VD = RD->getLastField();
-  } else if (const auto *CE = dyn_cast<CastExpr>(Base)) {
-    if (const auto *ME = dyn_cast<MemberExpr>(CE->getSubExpr()))
-      VD = dyn_cast<ValueDecl>(ME->getMemberDecl());
+      OuterRD = RD->getOuterLexicalRecordContext();
+  } else if (const auto *ME = dyn_cast<MemberExpr>(Base)) {
+    if (const ValueDecl *MD = ME->getMemberDecl()) {
+      OuterRD = MD->getDeclContext()->getOuterLexicalRecordContext();
+
+      const LangOptions::StrictFlexArraysLevelKind StrictFlexArraysLevel =
+          getLangOpts().getStrictFlexArraysLevel();
+      if (Decl::isFlexibleArrayMemberLike(
+              Ctx, MD, MD->getType(), StrictFlexArraysLevel,
+              /*IgnoreTemplateOrMacroSubstitution=*/true))
+        // Base is referencing the FAM itself.
+        FD = dyn_cast<FieldDecl>(MD);
+    }
   }
 
-  const auto *FD = dyn_cast_if_present<FieldDecl>(VD);
-  if (!FD || !FD->getParent() ||
-      !Decl::isFlexibleArrayMemberLike(getContext(), FD, FD->getType(),
-                                       StrictFlexArraysLevel, true))
+  if (!OuterRD)
     return nullptr;
 
+  if (!FD) {
+    const ValueDecl *VD = FindFlexibleArrayMemberField(Ctx, OuterRD);
+    FD = dyn_cast_if_present<FieldDecl>(VD);
+    if (!FD)
+      return nullptr;
+  }
+
   const auto *CBA = FD->getAttr<CountedByAttr>();
   if (!CBA)
     return nullptr;
 
-  StringRef FieldName = CBA->getCountedByField()->getName();
-  auto It =
-      llvm::find_if(FD->getParent()->fields(), [&](const FieldDecl *Field) {
-        return FieldName == Field->getName();
-      });
-  return It != FD->getParent()->field_end() ? *It : nullptr;
+  DeclarationName DName(CBA->getCountedByField());
+  DeclContext::lookup_result Lookup = OuterRD->lookup(DName);
+
+  if (Lookup.empty())
+    return nullptr;
+
+  return dyn_cast<ValueDecl>(Lookup.front());
 }
 
 void CodeGenFunction::EmitBoundsCheck(const Expr *E, const Expr *Base,
diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h
@@ -3022,11 +3022,18 @@ class CodeGenFunction : public CodeGenTypeCache {
   void EmitBoundsCheck(const Expr *E, const Expr *Base, llvm::Value *Index,
                        QualType IndexType, bool Accessed);
 
+  // Find a struct's flexible array member. It may be embedded inside multiple
+  // sub-structs, but must still be the last field.
+  const ValueDecl *FindFlexibleArrayMemberField(ASTContext &Ctx,
+                                                const RecordDecl *RD);
+
   /// Find the FieldDecl specified in a FAM's "counted_by" attribute. Returns
   /// \p nullptr if either the attribute or the field doesn't exist.
-  FieldDecl *FindCountedByField(
-      const Expr *Base,
-      LangOptions::StrictFlexArraysLevelKind StrictFlexArraysLevel);
+  const ValueDecl *FindCountedByField(const Expr *Base);
+
+  /// Build an expression accessing the "counted_by" field.
+  const Expr *BuildCountedByFieldExpr(const Expr *Base,
+                                      const ValueDecl *CountedByVD);
 
   llvm::Value *EmitScalarPrePostIncDec(const UnaryOperator *E, LValue LV,
                                        bool isInc, bool isPre);
diff --git a/clang/test/CodeGen/attr-counted-by.c b/clang/test/CodeGen/attr-counted-by.c
