Keep track of the set of array index variables we use when we

synthesize a by-copy captured array in a lambda. This information will
be needed by IR generation.


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

4 files changed, 49 insertions, 4 deletions

diff --git a/lib/AST/DeclCXX.cpp b/lib/AST/DeclCXX.cpp
index 257a507b23..fcd45088c2 100644
--- a/lib/AST/DeclCXX.cpp
+++ b/lib/AST/DeclCXX.cpp
@@ -38,13 +38,23 @@ AccessSpecDecl *AccessSpecDecl::CreateDeserialized(ASTContext &C, unsigned ID) {
 void CXXRecordDecl::LambdaDefinitionData::allocateExtra(
        ArrayRef<LambdaExpr::Capture> Captures,
        ArrayRef<Expr *> CaptureInits,
+       ArrayRef<VarDecl *> ArrayIndexVars,
+       ArrayRef<unsigned> ArrayIndexStarts,
        Stmt *Body) {
   NumCaptures = Captures.size();
   NumExplicitCaptures = 0;
   
   ASTContext &Context = Definition->getASTContext();
+  unsigned ArrayIndexSize = 0;
+  if (ArrayIndexVars.size() > 0) {
+    HasArrayIndexVars = true;
+    ArrayIndexSize = sizeof(unsigned) * (Captures.size() + 1)
+                   + sizeof(VarDecl *) * ArrayIndexVars.size();
+  }
+  
   this->Extra = Context.Allocate(sizeof(Capture) * Captures.size() +
-                                 sizeof(Stmt*) * (Captures.size() + 1));
+                                 sizeof(Stmt*) * (Captures.size() + 1) +
+                                 ArrayIndexSize);
   
   // Copy captures.
   Capture *ToCapture = getCaptures();
@@ -62,6 +72,15 @@ void CXXRecordDecl::LambdaDefinitionData::allocateExtra(
   
   // Copy the body of the lambda.
   *Stored++ = Body;
+  
+  if (ArrayIndexVars.size() > 0) {
+    assert(ArrayIndexStarts.size() == Captures.size());
+    memcpy(getArrayIndexVars(), ArrayIndexVars.data(),
+           sizeof(VarDecl *) * ArrayIndexVars.size());
+    memcpy(getArrayIndexStarts(), ArrayIndexStarts.data(), 
+           sizeof(unsigned) * Captures.size());
+    getArrayIndexStarts()[Captures.size()] = ArrayIndexVars.size();
+  }
 }
 
 
diff --git a/lib/AST/ExprCXX.cpp b/lib/AST/ExprCXX.cpp
index 1f8a57a4e3..1e766ada23 100644
--- a/lib/AST/ExprCXX.cpp
+++ b/lib/AST/ExprCXX.cpp
@@ -753,6 +753,8 @@ LambdaExpr::LambdaExpr(QualType T,
                        ArrayRef<Capture> Captures, 
                        bool ExplicitParams,
                        ArrayRef<Expr *> CaptureInits,
+                       ArrayRef<VarDecl *> ArrayElementVars,
+                       ArrayRef<unsigned> ArrayElementStarts,
                        SourceLocation ClosingBrace)
   : Expr(LambdaExprClass, T, VK_RValue, OK_Ordinary,
          T->isDependentType(), T->isDependentType(), T->isDependentType(),
@@ -765,7 +767,8 @@ LambdaExpr::LambdaExpr(QualType T,
   assert(CaptureInits.size() == Captures.size() && "Wrong number of arguments");
   CXXRecordDecl *Class = getLambdaClass();
   CXXRecordDecl::LambdaDefinitionData &Data = Class->getLambdaData();
-  Data.allocateExtra(Captures, CaptureInits, getCallOperator()->getBody());
+  Data.allocateExtra(Captures, CaptureInits, ArrayElementVars, 
+                     ArrayElementStarts, getCallOperator()->getBody());
   
   // FIXME: Propagate "has unexpanded parameter pack" bit.
 }
@@ -777,6 +780,8 @@ LambdaExpr *LambdaExpr::Create(ASTContext &Context,
                                ArrayRef<Capture> Captures, 
                                bool ExplicitParams,
                                ArrayRef<Expr *> CaptureInits,
+                               ArrayRef<VarDecl *> ArrayElementVars,
+                               ArrayRef<unsigned> ArrayElementStarts,
                                SourceLocation ClosingBrace) {
   // Determine the type of the expression (i.e., the type of the
   // function object we're creating).
@@ -784,6 +789,7 @@ LambdaExpr *LambdaExpr::Create(ASTContext &Context,
 
   return new (Context) LambdaExpr(T, IntroducerRange, CaptureDefault, 
                                   Captures, ExplicitParams, CaptureInits,
+                                  ArrayElementVars, ArrayElementStarts,
                                   ClosingBrace);
 }
 
@@ -826,6 +832,19 @@ LambdaExpr::capture_init_iterator LambdaExpr::capture_init_end() const {
   return reinterpret_cast<Expr **>(Data.getStoredStmts() + Data.NumCaptures);
 }
 
+ArrayRef<VarDecl *> 
+LambdaExpr::getCaptureInitIndexVars(capture_init_iterator Iter) const {
+  CXXRecordDecl::LambdaDefinitionData &Data = getLambdaClass()->getLambdaData();
+  assert(Data.HasArrayIndexVars && "No array index-var data?");
+  
+  unsigned Index = Iter - capture_init_begin();
+  assert(Index < Data.NumCaptures && "Capture index out-of-range");
+  VarDecl **IndexVars = Data.getArrayIndexVars();
+  unsigned *IndexStarts = Data.getArrayIndexStarts();
+  return ArrayRef<VarDecl *>(IndexVars + IndexStarts[Index],
+                             IndexVars + IndexStarts[Index + 1]);
+}
+
 CXXRecordDecl *LambdaExpr::getLambdaClass() const {
   return getType()->getAsCXXRecordDecl();
 }
diff --git a/lib/Sema/SemaExpr.cpp b/lib/Sema/SemaExpr.cpp
index 9bd1074116..0d43afcd81 100644
--- a/lib/Sema/SemaExpr.cpp
+++ b/lib/Sema/SemaExpr.cpp
@@ -9642,6 +9642,7 @@ static ExprResult captureInLambda(Sema &S, LambdaScopeInfo *LSI,
   bool InitializingArray = false;
   QualType BaseType = FieldType;
   QualType SizeType = S.Context.getSizeType();
+  LSI->ArrayIndexStarts.push_back(LSI->ArrayIndexVars.size());
   while (const ConstantArrayType *Array
                         = S.Context.getAsConstantArrayType(BaseType)) {
     InitializingArray = true;
@@ -9660,7 +9661,8 @@ static ExprResult captureInLambda(Sema &S, LambdaScopeInfo *LSI,
                         S.Context.getTrivialTypeSourceInfo(SizeType, Loc),
                         SC_None, SC_None);
     IndexVariables.push_back(IterationVar);
-
+    LSI->ArrayIndexVars.push_back(IterationVar);
+    
     // Create a reference to the iteration variable.
     ExprResult IterationVarRef
       = S.BuildDeclRefExpr(IterationVar, SizeType, VK_LValue, Loc);
diff --git a/lib/Sema/SemaLambda.cpp b/lib/Sema/SemaLambda.cpp
index 44b388364a..f91b93f957 100644
--- a/lib/Sema/SemaLambda.cpp
+++ b/lib/Sema/SemaLambda.cpp
@@ -306,6 +306,8 @@ ExprResult Sema::ActOnLambdaExpr(SourceLocation StartLoc,
   SourceRange IntroducerRange;
   bool ExplicitParams;
   bool LambdaExprNeedsCleanups;
+  llvm::SmallVector<VarDecl *, 4> ArrayIndexVars;
+  llvm::SmallVector<unsigned, 4> ArrayIndexStarts;
   {
     LambdaScopeInfo *LSI = getCurLambda();
     CallOperator = LSI->CallOperator;
@@ -313,7 +315,9 @@ ExprResult Sema::ActOnLambdaExpr(SourceLocation StartLoc,
     IntroducerRange = LSI->IntroducerRange;
     ExplicitParams = LSI->ExplicitParams;
     LambdaExprNeedsCleanups = LSI->ExprNeedsCleanups;
-
+    ArrayIndexVars.swap(LSI->ArrayIndexVars);
+    ArrayIndexStarts.swap(LSI->ArrayIndexStarts);
+    
     // Translate captures.
     for (unsigned I = 0, N = LSI->Captures.size(); I != N; ++I) {
       LambdaScopeInfo::Capture From = LSI->Captures[I];
@@ -467,6 +471,7 @@ ExprResult Sema::ActOnLambdaExpr(SourceLocation StartLoc,
   LambdaExpr *Lambda = LambdaExpr::Create(Context, Class, IntroducerRange, 
                                           CaptureDefault, Captures, 
                                           ExplicitParams, CaptureInits, 
+                                          ArrayIndexVars, ArrayIndexStarts,
                                           Body->getLocEnd());
 
   // C++11 [expr.prim.lambda]p2: