Flatten the AST so that we dont' get ASan issues

Author: erichkeane

clang/include/clang/AST/OpenACCClause.h

@@ -1277,7 +1277,7 @@ class OpenACCCreateClause final
 };
 
 // A structure to stand in for the recipe on a reduction.  RecipeDecl is the
-// 'main' declaration used for initializaiton, which is fixed. 
+// 'main' declaration used for initializaiton, which is fixed.
 struct OpenACCReductionRecipe {
   VarDecl *AllocaDecl;
 
@@ -1297,36 +1297,69 @@ struct OpenACCReductionRecipe {
   // -For a struct without the operator, this will be 1 element per field, which
   //  should be the combiner for that element.
   // -For an array of any of the above, it will be the above for the element.
-  llvm::SmallVector<CombinerRecipe, 1> CombinerRecipes;
+  // Note: These are necessarily stored in either Trailing Storage (when in the
+  // AST), or in a separate collection when being semantically analyzed.
+  llvm::ArrayRef<CombinerRecipe> CombinerRecipes;
 
   OpenACCReductionRecipe(VarDecl *A, llvm::ArrayRef<CombinerRecipe> Combiners)
       : AllocaDecl(A), CombinerRecipes(Combiners) {}
 
   bool isSet() const { return AllocaDecl; }
-  static OpenACCReductionRecipe Empty() {
-    return OpenACCReductionRecipe(/*AllocaDecl=*/nullptr, {});
+};
+
+// A version of the above that is used for semantic analysis, at a time before
+// the OpenACCReductionClause node has been created.  This one has storage for
+// the CombinerRecipe, since Trailing storage for it doesn't exist yet.
+struct OpenACCReductionRecipeWithStorage : OpenACCReductionRecipe {
+  llvm::SmallVector<CombinerRecipe, 1> CombinerRecipeStorage;
+
+  OpenACCReductionRecipeWithStorage(VarDecl *A,
+                                    llvm::ArrayRef<CombinerRecipe> Combiners)
+      : OpenACCReductionRecipe(A, {}), CombinerRecipeStorage(Combiners) {
+    CombinerRecipes = CombinerRecipeStorage;
+  }
+  static OpenACCReductionRecipeWithStorage Empty() {
+    return OpenACCReductionRecipeWithStorage(/*AllocaDecl=*/nullptr, {});
   }
 };
 
 class OpenACCReductionClause final
     : public OpenACCClauseWithVarList,
       private llvm::TrailingObjects<OpenACCReductionClause, Expr *,
-                                    OpenACCReductionRecipe> {
+                                    OpenACCReductionRecipe,
+                                    OpenACCReductionRecipe::CombinerRecipe> {
   friend TrailingObjects;
   OpenACCReductionOperator Op;
 
   OpenACCReductionClause(SourceLocation BeginLoc, SourceLocation LParenLoc,
                          OpenACCReductionOperator Operator,
                          ArrayRef<Expr *> VarList,
-                         ArrayRef<OpenACCReductionRecipe> Recipes,
+                         ArrayRef<OpenACCReductionRecipeWithStorage> Recipes,
                          SourceLocation EndLoc)
       : OpenACCClauseWithVarList(OpenACCClauseKind::Reduction, BeginLoc,
                                  LParenLoc, EndLoc),
         Op(Operator) {
-          assert(VarList.size() == Recipes.size());
+    assert(VarList.size() == Recipes.size());
     setExprs(getTrailingObjects<Expr *>(VarList.size()), VarList);
-    llvm::uninitialized_copy(Recipes, getTrailingObjects<
-                             OpenACCReductionRecipe > ());
+
+    // Since we're using trailing storage on this node to store the 'combiner'
+    // recipes of the Reduction Recipes (which have a 1:M relationship), we need
+    // to ensure we get the ArrayRef of each of our combiner 'correct'.
+    OpenACCReductionRecipe::CombinerRecipe *CurCombinerLoc =
+        getTrailingObjects<OpenACCReductionRecipe::CombinerRecipe>();
+    for (const auto &[Idx, R] : llvm::enumerate(Recipes)) {
+
+      // ArrayRef to the 'correct' data location in trailing storage.
+      llvm::MutableArrayRef<OpenACCReductionRecipe::CombinerRecipe>
+          NewCombiners{CurCombinerLoc, R.CombinerRecipes.size()};
+      CurCombinerLoc += R.CombinerRecipes.size();
+
+      llvm::uninitialized_copy(R.CombinerRecipes, NewCombiners.begin());
+
+      // Placement new into the correct location in trailng storage.
+      new (&getTrailingObjects<OpenACCReductionRecipe>()[Idx])
+          OpenACCReductionRecipe(R.AllocaDecl, NewCombiners);
+    }
   }
 
 public:
@@ -1347,13 +1380,17 @@ class OpenACCReductionClause final
   static OpenACCReductionClause *
   Create(const ASTContext &C, SourceLocation BeginLoc, SourceLocation LParenLoc,
          OpenACCReductionOperator Operator, ArrayRef<Expr *> VarList,
-         ArrayRef<OpenACCReductionRecipe> Recipes, SourceLocation EndLoc);
+         ArrayRef<OpenACCReductionRecipeWithStorage> Recipes,
+         SourceLocation EndLoc);
 
   OpenACCReductionOperator getReductionOp() const { return Op; }
 
   size_t numTrailingObjects(OverloadToken<Expr *>) const {
     return getExprs().size();
   }
+  size_t numTrailingObjects(OverloadToken<OpenACCReductionRecipe>) const {
+    return getExprs().size();
+  }
 };
 
 class OpenACCLinkClause final

clang/include/clang/Sema/SemaOpenACC.h

@@ -245,7 +245,7 @@ class SemaOpenACC : public SemaBase {
 
   OpenACCPrivateRecipe CreatePrivateInitRecipe(const Expr *VarExpr);
   OpenACCFirstPrivateRecipe CreateFirstPrivateInitRecipe(const Expr *VarExpr);
-  OpenACCReductionRecipe
+  OpenACCReductionRecipeWithStorage
   CreateReductionInitRecipe(OpenACCReductionOperator ReductionOperator,
                             const Expr *VarExpr);
 
@@ -951,12 +951,14 @@ class SemaOpenACC : public SemaBase {
                   ArrayRef<Expr *> IntExprs, SourceLocation EndLoc);
   // Does the checking for a 'reduction ' clause that needs to be done in
   // dependent and not dependent cases.
-  OpenACCClause *CheckReductionClause(
-      ArrayRef<const OpenACCClause *> ExistingClauses,
-      OpenACCDirectiveKind DirectiveKind, SourceLocation BeginLoc,
-      SourceLocation LParenLoc, OpenACCReductionOperator ReductionOp,
-      ArrayRef<Expr *> Vars, ArrayRef<OpenACCReductionRecipe> Recipes,
-      SourceLocation EndLoc);
+  OpenACCClause *
+  CheckReductionClause(ArrayRef<const OpenACCClause *> ExistingClauses,
+                       OpenACCDirectiveKind DirectiveKind,
+                       SourceLocation BeginLoc, SourceLocation LParenLoc,
+                       OpenACCReductionOperator ReductionOp,
+                       ArrayRef<Expr *> Vars,
+                       ArrayRef<OpenACCReductionRecipeWithStorage> Recipes,
+                       SourceLocation EndLoc);
 
   ExprResult BuildOpenACCAsteriskSizeExpr(SourceLocation AsteriskLoc);
   ExprResult ActOnOpenACCAsteriskSizeExpr(SourceLocation AsteriskLoc);

clang/lib/AST/OpenACCClause.cpp

@@ -506,11 +506,17 @@ OpenACCDeviceTypeClause *OpenACCDeviceTypeClause::Create(
 OpenACCReductionClause *OpenACCReductionClause::Create(
     const ASTContext &C, SourceLocation BeginLoc, SourceLocation LParenLoc,
     OpenACCReductionOperator Operator, ArrayRef<Expr *> VarList,
-    ArrayRef<OpenACCReductionRecipe> Recipes,
+    ArrayRef<OpenACCReductionRecipeWithStorage> Recipes,
     SourceLocation EndLoc) {
-  void *Mem = C.Allocate(
-      OpenACCReductionClause::totalSizeToAlloc<Expr *, OpenACCReductionRecipe>(
-          VarList.size(), Recipes.size()));
+  size_t NumCombiners = llvm::accumulate(
+      Recipes, 0, [](size_t Num, const OpenACCReductionRecipe &R) {
+        return Num + R.CombinerRecipes.size();
+      });
+
+  void *Mem = C.Allocate(OpenACCReductionClause::totalSizeToAlloc<
+                         Expr *, OpenACCReductionRecipe,
+                         OpenACCReductionRecipe::CombinerRecipe>(
+      VarList.size(), Recipes.size(), NumCombiners));
   return new (Mem) OpenACCReductionClause(BeginLoc, LParenLoc, Operator,
                                           VarList, Recipes, EndLoc);
 }

clang/lib/Sema/SemaOpenACC.cpp

@@ -2883,12 +2883,12 @@ SemaOpenACC::CreateFirstPrivateInitRecipe(const Expr *VarExpr) {
   return OpenACCFirstPrivateRecipe(AllocaDecl, Temporary);
 }
 
-OpenACCReductionRecipe SemaOpenACC::CreateReductionInitRecipe(
+OpenACCReductionRecipeWithStorage SemaOpenACC::CreateReductionInitRecipe(
     OpenACCReductionOperator ReductionOperator, const Expr *VarExpr) {
   // We don't strip bounds here, so that we are doing our recipe init at the
   // 'lowest' possible level.  Codegen is going to have to do its own 'looping'.
   if (!VarExpr || VarExpr->getType()->isDependentType())
-    return OpenACCReductionRecipe::Empty();
+    return OpenACCReductionRecipeWithStorage::Empty();
 
   QualType VarTy =
       VarExpr->getType().getNonReferenceType().getUnqualifiedType();
@@ -2905,7 +2905,7 @@ OpenACCReductionRecipe SemaOpenACC::CreateReductionInitRecipe(
   // at any of the combiners.
   if (CreateReductionCombinerRecipe(VarExpr->getBeginLoc(), ReductionOperator,
                                     VarTy, CombinerRecipes))
-    return OpenACCReductionRecipe::Empty();
+    return OpenACCReductionRecipeWithStorage::Empty();
 
   VarDecl *AllocaDecl = CreateAllocaDecl(
       getASTContext(), SemaRef.getCurContext(), VarExpr->getBeginLoc(),
@@ -2956,7 +2956,7 @@ OpenACCReductionRecipe SemaOpenACC::CreateReductionInitRecipe(
     AllocaDecl->setInitStyle(VarDecl::CallInit);
   }
 
-  return OpenACCReductionRecipe(AllocaDecl, CombinerRecipes);
+  return OpenACCReductionRecipeWithStorage(AllocaDecl, CombinerRecipes);
 }
 
 bool SemaOpenACC::CreateReductionCombinerRecipe(

clang/lib/Sema/SemaOpenACCClause.cpp

@@ -1772,7 +1772,7 @@ OpenACCClause *SemaOpenACCClauseVisitor::VisitReductionClause(
   }
 
   SmallVector<Expr *> ValidVars;
-  SmallVector<OpenACCReductionRecipe> Recipes;
+  SmallVector<OpenACCReductionRecipeWithStorage> Recipes;
 
   for (Expr *Var : Clause.getVarList()) {
     ExprResult Res = SemaRef.CheckReductionVar(Clause.getDirectiveKind(),
@@ -2196,7 +2196,7 @@ OpenACCClause *SemaOpenACC::CheckReductionClause(
     ArrayRef<const OpenACCClause *> ExistingClauses,
     OpenACCDirectiveKind DirectiveKind, SourceLocation BeginLoc,
     SourceLocation LParenLoc, OpenACCReductionOperator ReductionOp,
-    ArrayRef<Expr *> Vars, ArrayRef<OpenACCReductionRecipe> Recipes,
+    ArrayRef<Expr *> Vars, ArrayRef<OpenACCReductionRecipeWithStorage> Recipes,
     SourceLocation EndLoc) {
   if (DirectiveKind == OpenACCDirectiveKind::Loop ||
       isOpenACCCombinedDirectiveKind(DirectiveKind)) {

clang/lib/Sema/TreeTransform.h

@@ -12374,7 +12374,7 @@ void OpenACCClauseTransform<Derived>::VisitReductionClause(
     const OpenACCReductionClause &C) {
   SmallVector<Expr *> TransformedVars = VisitVarList(C.getVarList());
   SmallVector<Expr *> ValidVars;
-  llvm::SmallVector<OpenACCReductionRecipe> Recipes;
+  llvm::SmallVector<OpenACCReductionRecipeWithStorage> Recipes;
 
   for (const auto [Var, OrigRecipe] :
        llvm::zip(TransformedVars, C.getRecipes())) {
@@ -12384,7 +12384,7 @@ void OpenACCClauseTransform<Derived>::VisitReductionClause(
       ValidVars.push_back(Res.get());
 
       if (OrigRecipe.isSet())
-        Recipes.push_back(OrigRecipe);
+        Recipes.emplace_back(OrigRecipe.AllocaDecl, OrigRecipe.CombinerRecipes);
       else
         Recipes.push_back(Self.getSema().OpenACC().CreateReductionInitRecipe(
             C.getReductionOp(), Res.get()));

clang/lib/Serialization/ASTReader.cpp

@@ -13006,7 +13006,7 @@ OpenACCClause *ASTRecordReader::readOpenACCClause() {
     SourceLocation LParenLoc = readSourceLocation();
     OpenACCReductionOperator Op = readEnum<OpenACCReductionOperator>();
     llvm::SmallVector<Expr *> VarList = readOpenACCVarList();
-    llvm::SmallVector<OpenACCReductionRecipe> RecipeList;
+    llvm::SmallVector<OpenACCReductionRecipeWithStorage> RecipeList;
 
     for (unsigned I = 0; I < VarList.size(); ++I) {
       VarDecl *Recipe = readDeclAs<VarDecl>();