[AST] Introduce GenericSelectionExpr::Association

Introduce a new class GenericSelectionExpr::Association which bundle together
an association expression and its TypeSourceInfo.

An iterator GenericSelectionExpr::AssociationIterator is additionally added to
make it possible to iterate over ranges of Associations. This iterator is a
kind of proxy iterator which abstract over how exactly the expressions and the
TypeSourceInfos are stored.

Differential Revision: https://reviews.llvm.org/D57106

Reviewed By: aaron.ballman

Reviewers: aaron.ballman, steveire, dblaikie, mclow.lists



git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@352369 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 28 insertions, 21 deletions

diff --git a/lib/Sema/SemaExprObjC.cpp b/lib/Sema/SemaExprObjC.cpp
index 4ca3b851c1..fdb59874cc 100644
--- a/lib/Sema/SemaExprObjC.cpp
+++ b/lib/Sema/SemaExprObjC.cpp
@@ -4332,14 +4332,16 @@ Expr *Sema::stripARCUnbridgedCast(Expr *e) {
     assert(!gse->isResultDependent());
 
     unsigned n = gse->getNumAssocs();
-    SmallVector<Expr*, 4> subExprs(n);
-    SmallVector<TypeSourceInfo*, 4> subTypes(n);
-    for (unsigned i = 0; i != n; ++i) {
-      subTypes[i] = gse->getAssocTypeSourceInfo(i);
-      Expr *sub = gse->getAssocExpr(i);
-      if (i == gse->getResultIndex())
+    SmallVector<Expr *, 4> subExprs;
+    SmallVector<TypeSourceInfo *, 4> subTypes;
+    subExprs.reserve(n);
+    subTypes.reserve(n);
+    for (const GenericSelectionExpr::Association &assoc : gse->associations()) {
+      subTypes.push_back(assoc.getTypeSourceInfo());
+      Expr *sub = assoc.getAssociationExpr();
+      if (assoc.isSelected())
         sub = stripARCUnbridgedCast(sub);
-      subExprs[i] = sub;
+      subExprs.push_back(sub);
     }
 
     return GenericSelectionExpr::Create(
diff --git a/lib/Sema/SemaPseudoObject.cpp b/lib/Sema/SemaPseudoObject.cpp
index 37dd7b1e8d..28a4d62b03 100644
--- a/lib/Sema/SemaPseudoObject.cpp
+++ b/lib/Sema/SemaPseudoObject.cpp
@@ -140,19 +140,23 @@ namespace {
         unsigned resultIndex = gse->getResultIndex();
         unsigned numAssocs = gse->getNumAssocs();
 
-        SmallVector<Expr*, 8> assocs(numAssocs);
-        SmallVector<TypeSourceInfo*, 8> assocTypes(numAssocs);
-
-        for (unsigned i = 0; i != numAssocs; ++i) {
-          Expr *assoc = gse->getAssocExpr(i);
-          if (i == resultIndex) assoc = rebuild(assoc);
-          assocs[i] = assoc;
-          assocTypes[i] = gse->getAssocTypeSourceInfo(i);
+        SmallVector<Expr *, 8> assocExprs;
+        SmallVector<TypeSourceInfo *, 8> assocTypes;
+        assocExprs.reserve(numAssocs);
+        assocTypes.reserve(numAssocs);
+
+        for (const GenericSelectionExpr::Association &assoc :
+             gse->associations()) {
+          Expr *assocExpr = assoc.getAssociationExpr();
+          if (assoc.isSelected())
+            assocExpr = rebuild(assocExpr);
+          assocExprs.push_back(assocExpr);
+          assocTypes.push_back(assoc.getTypeSourceInfo());
         }
 
         return GenericSelectionExpr::Create(
             S.Context, gse->getGenericLoc(), gse->getControllingExpr(),
-            assocTypes, assocs, gse->getDefaultLoc(), gse->getRParenLoc(),
+            assocTypes, assocExprs, gse->getDefaultLoc(), gse->getRParenLoc(),
             gse->containsUnexpandedParameterPack(), resultIndex);
       }
 
diff --git a/lib/Sema/TreeTransform.h b/lib/Sema/TreeTransform.h
index 389031d491..2e434a3c0e 100644
--- a/lib/Sema/TreeTransform.h
+++ b/lib/Sema/TreeTransform.h
@@ -9071,10 +9071,10 @@ TreeTransform<Derived>::TransformGenericSelectionExpr(GenericSelectionExpr *E) {
 
   SmallVector<Expr *, 4> AssocExprs;
   SmallVector<TypeSourceInfo *, 4> AssocTypes;
-  for (unsigned i = 0; i != E->getNumAssocs(); ++i) {
-    TypeSourceInfo *TS = E->getAssocTypeSourceInfo(i);
-    if (TS) {
-      TypeSourceInfo *AssocType = getDerived().TransformType(TS);
+  for (const GenericSelectionExpr::Association &Assoc : E->associations()) {
+    TypeSourceInfo *TSI = Assoc.getTypeSourceInfo();
+    if (TSI) {
+      TypeSourceInfo *AssocType = getDerived().TransformType(TSI);
       if (!AssocType)
         return ExprError();
       AssocTypes.push_back(AssocType);
@@ -9082,7 +9082,8 @@ TreeTransform<Derived>::TransformGenericSelectionExpr(GenericSelectionExpr *E) {
       AssocTypes.push_back(nullptr);
     }
 
-    ExprResult AssocExpr = getDerived().TransformExpr(E->getAssocExpr(i));
+    ExprResult AssocExpr =
+        getDerived().TransformExpr(Assoc.getAssociationExpr());
     if (AssocExpr.isInvalid())
       return ExprError();
     AssocExprs.push_back(AssocExpr.get());