Support pack iteration in closures

Author: simanerush

lib/AST/ASTWalker.cpp

@@ -1890,11 +1890,9 @@ Stmt *Traversal::visitForEachStmt(ForEachStmt *S) {
   //
   // If for-in is already type-checked, the type-checked version
   // of the sequence is going to be visited as part of `iteratorVar`.
-  if (S->getTypeCheckedSequence()) {
-    if (auto IteratorVar = S->getIteratorVar()) {
-      if (doIt(IteratorVar))
-        return nullptr;
-    }
+  if (auto IteratorVar = S->getIteratorVar()) {
+    if (doIt(IteratorVar))
+      return nullptr;
 
     if (auto NextCall = S->getNextCall()) {
       if ((NextCall = doIt(NextCall)))

lib/AST/Decl.cpp

@@ -6893,14 +6893,6 @@ VarDecl::VarDecl(DeclKind kind, bool isStatic, VarDecl::Introducer introducer,
 Type VarDecl::getTypeInContext() const {
   // If we are performing pack iteration, use the generic environment of the
   // pack expansion expression to get the right context of a local variable.
-  //  if (auto *forEachStmt =
-  //          dyn_cast_or_null<ForEachStmt>(this->getParentPatternStmt())) {
-  //    if (auto *expansion =
-  //            dyn_cast<PackExpansionExpr>(forEachStmt->getParsedSequence())) {
-  //      return GenericEnvironment::mapTypeIntoContext(
-  //          expansion->getGenericEnvironment(), getInterfaceType());
-  //    }
-  //  }
   if (auto *env = getOpenedElementEnvironment())
     return GenericEnvironment::mapTypeIntoContext(env, getInterfaceType());
 

lib/Sema/CSSyntacticElement.cpp

@@ -1874,6 +1874,36 @@ class SyntacticElementSolutionApplication
   }
 
   ASTNode visitForEachStmt(ForEachStmt *forEachStmt) {
+
+    class Walker : public ASTWalker {
+      GenericEnvironment *Environment;
+
+    public:
+      Walker(GenericEnvironment *Environment) {
+        this->Environment = Environment;
+      }
+
+      PreWalkResult<Stmt *> walkToStmtPre(Stmt *S) override {
+        if (isa<ForEachStmt>(S)) {
+          return Action::Stop();
+        }
+        return Action::Continue(S);
+      }
+
+      PreWalkAction walkToDeclPre(Decl *D) override {
+        if (auto *decl = dyn_cast<VarDecl>(D)) {
+          decl->setOpenedElementEnvironment(Environment);
+        }
+        if (isa<AbstractFunctionDecl>(D)) {
+          return Action::Stop();
+        }
+        if (isa<NominalTypeDecl>(D)) {
+          return Action::Stop();
+        }
+        return Action::Continue();
+      }
+    };
+
     ConstraintSystem &cs = solution.getConstraintSystem();
 
     auto forEachTarget = rewriteTarget(*cs.getTargetFor(forEachStmt));
@@ -1884,6 +1914,13 @@ class SyntacticElementSolutionApplication
     auto body = visit(forEachStmt->getBody()).get<Stmt *>();
     forEachStmt->setBody(cast<BraceStmt>(body));
 
+    if (auto *expansion = dyn_cast<PackExpansionExpr>(
+            forEachStmt->getTypeCheckedSequence())) {
+      Walker forEachWalker(expansion->getGenericEnvironment());
+      forEachStmt->getPattern()->walk(forEachWalker);
+      forEachStmt->getBody()->walk(forEachWalker);
+    }
+
     // Check to see if the sequence expr is throwing (in async context),
     // if so require the stmt to have a `try`.
     hadError |= diagnoseUnhandledThrowsInAsyncContext(

lib/Sema/TypeCheckStmt.cpp

@@ -1425,7 +1425,7 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
     RWS->setCond(E);
     return RWS;
   }
-  
+
   Stmt *visitForEachStmt(ForEachStmt *S) {
     // A special walker to record opened element environment for var decls in a
     // for-each loop.