[Clang] Use correct base expression for counted_by field (#73168)

clang/lib/CodeGen/CGBuiltin.cpp

@@ -945,7 +945,7 @@ CodeGenFunction::emitFlexibleArrayMemberSize(const Expr *E, unsigned Type,
                      : Builder.CreateZExtOrTrunc(FAMSize, ResType);
   Value *Res = FAMSize;
 
-  if (const auto *DRE = dyn_cast<DeclRefExpr>(Base)) {
+  if (isa<DeclRefExpr>(Base)) {
     // The whole struct is specificed in the __bdos.
     const RecordDecl *OuterRD =
         CountedByFD->getDeclContext()->getOuterLexicalRecordContext();

clang/lib/CodeGen/CGExpr.cpp

@@ -962,31 +962,50 @@ CodeGenFunction::BuildCountedByFieldExpr(const Expr *Base,
   // Find the outer struct expr (i.e. p in p->a.b.c.d).
   Expr *CountedByExpr = const_cast<Expr *>(Base)->IgnoreParenImpCasts();
 
+  // Get the enclosing struct, but not the outermost enclosing struct.
+  const DeclContext *DC = CountedByVD->getLexicalDeclContext();
+  const auto *CountedByRD = dyn_cast<RecordDecl>(DC);
+  if (!CountedByRD)
+    return nullptr;
+
   // 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();
+  while (auto *ME = dyn_cast<MemberExpr>(CountedByExpr)) {
+    CountedByExpr = ME->getBase()->IgnoreImpCasts();
+
+    // Use the base of an ArraySubscriptExpr.
+    while (auto *ASE = dyn_cast<ArraySubscriptExpr>(CountedByExpr))
+      CountedByExpr = ASE->getBase()->IgnoreImpCasts();
+
+    QualType Ty = CountedByExpr->getType();
+    if (Ty->isPointerType())
+      Ty = Ty->getPointeeType();
+
+    // Stop when we reach the struct containing the counted_by field.
+    if (CountedByRD == Ty->getAsRecordDecl())
+      break;
+  }
+
+  ASTContext &Ctx = getContext();
 
-  // 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 (!CountedByExpr->isPRValue())
+    // Add an implicit cast to create the required pr-value.
+    CountedByExpr = ImplicitCastExpr::Create(
+        Ctx, 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);
+      CountedByExpr = MemberExpr::CreateImplicit(
+          Ctx, CountedByExpr, CountedByExpr->getType()->isPointerType(), VD,
+          VD->getType(), VK_LValue, OK_Ordinary);
     }
   } else {
     CountedByExpr = MemberExpr::CreateImplicit(
-        getContext(), const_cast<Expr *>(CountedByExpr),
-        CountedByExpr->getType()->isPointerType(),
+        Ctx, CountedByExpr, CountedByExpr->getType()->isPointerType(),
         const_cast<ValueDecl *>(CountedByVD), CountedByVD->getType(), VK_LValue,
         OK_Ordinary);
   }