[SE-0408] Enable nested iteration

Author: simanerush

include/swift/Sema/ConstraintSystem.h

@@ -2301,6 +2301,8 @@ class ConstraintSystem {
   llvm::SmallMapVector<PackElementExpr *, PackExpansionExpr *, 2>
       PackEnvironments;
 
+  llvm::SmallVector<GenericEnvironment * , 4> PackElementGenericEnvironments;
+
   /// The set of functions that have been transformed by a result builder.
   llvm::MapVector<AnyFunctionRef, AppliedBuilderTransform>
       resultBuilderTransformed;

include/swift/Sema/SyntacticElementTarget.h

@@ -162,6 +162,7 @@ class SyntacticElementTarget {
       DeclContext *dc;
       Pattern *pattern;
       bool ignoreWhereClause;
+      GenericEnvironment *packElementEnv;
       ForEachStmtInfo info;
     } forEachStmt;
 
@@ -241,11 +242,12 @@ class SyntacticElementTarget {
   }
 
   SyntacticElementTarget(ForEachStmt *stmt, DeclContext *dc,
-                         bool ignoreWhereClause)
+                         bool ignoreWhereClause, GenericEnvironment *packElementEnv)
       : kind(Kind::forEachStmt) {
     forEachStmt.stmt = stmt;
     forEachStmt.dc = dc;
     forEachStmt.ignoreWhereClause = ignoreWhereClause;
+    forEachStmt.packElementEnv = packElementEnv;
   }
 
   /// Form a target for the initialization of a pattern from an expression.
@@ -263,7 +265,8 @@ class SyntacticElementTarget {
   /// Form a target for a for-in loop.
   static SyntacticElementTarget forForEachStmt(ForEachStmt *stmt,
                                                DeclContext *dc,
-                                               bool ignoreWhereClause = false);
+                                               bool ignoreWhereClause = false,
+                                               GenericEnvironment *packElementEnv = nullptr);
 
   /// Form a target for a property with an attached property wrapper that is
   /// initialized out-of-line.
@@ -538,6 +541,11 @@ class SyntacticElementTarget {
     return forEachStmt.ignoreWhereClause;
   }
 
+  GenericEnvironment *getPackElementEnv() const {
+    assert(isForEachStmt());
+    return forEachStmt.packElementEnv;
+  }
+
   const ForEachStmtInfo &getForEachStmtInfo() const {
     assert(isForEachStmt());
     return forEachStmt.info;

lib/AST/GenericEnvironment.cpp

@@ -734,10 +734,11 @@ GenericEnvironment::mapElementTypeIntoPackContext(Type type) const {
 
   type = type->mapTypeOutOfContext();
 
+  auto interfaceType = element->getInterfaceType();
+
   llvm::SmallDenseMap<GenericParamKey, GenericTypeParamType *>
       packParamForElement;
-  auto elementDepth =
-      sig.getInnermostGenericParams().front()->getDepth() + 1;
+  auto elementDepth = interfaceType->getRootGenericParam()->getDepth();
 
   for (auto *genericParam : sig.getGenericParams()) {
     if (!genericParam->isParameterPack())
@@ -792,6 +793,8 @@ Type BuildForwardingSubstitutions::operator()(SubstitutableType *type) const {
     auto param = type->castTo<GenericTypeParamType>();
     if (!param->isParameterPack())
       return resultType;
+    if (resultType->is<PackType>())
+      return resultType;
     return PackType::getSingletonPackExpansion(resultType);
   }
   return Type();

lib/Sema/CSGen.cpp

@@ -4892,6 +4892,12 @@ bool ConstraintSystem::generateConstraints(
   }
 
   case SyntacticElementTarget::Kind::forEachStmt: {
+
+    // Cache the outer generic environment, if it exists.
+    if (target.getPackElementEnv()) {
+      PackElementGenericEnvironments.push_back(target.getPackElementEnv());
+    }
+
     // For a for-each statement, generate constraints for the pattern, where
     // clause, and sequence traversal.
     auto resultTarget = generateForEachStmtConstraints(*this, target);

lib/Sema/ConstraintSystem.cpp

@@ -680,10 +680,16 @@ ConstraintSystem::getPackElementEnvironment(ConstraintLocator *locator,
       shapeClass->mapTypeOutOfContext()->getCanonicalType());
 
   auto &ctx = getASTContext();
+  GenericEnvironment *outerGenericEnvironment = PackElementGenericEnvironments.empty() ? nullptr : PackElementGenericEnvironments.back();
+
+  auto genericSignature = outerGenericEnvironment ?
+                  outerGenericEnvironment->getGenericSignature().getCanonicalSignature() :
+                  DC->getGenericSignatureOfContext().getCanonicalSignature();
+
   auto elementSig = ctx.getOpenedElementSignature(
-      DC->getGenericSignatureOfContext().getCanonicalSignature(), shapeParam);
+      genericSignature, shapeParam);
   auto *contextEnv = DC->getGenericEnvironmentOfContext();
-  auto contextSubs = contextEnv->getForwardingSubstitutionMap();
+  auto contextSubs = outerGenericEnvironment ? outerGenericEnvironment->getForwardingSubstitutionMap() : contextEnv->getForwardingSubstitutionMap();
   return GenericEnvironment::forOpenedElement(elementSig, uuidAndShape.first,
                                               shapeParam, contextSubs);
 }

lib/Sema/SyntacticElementTarget.cpp

@@ -178,8 +178,9 @@ SyntacticElementTarget SyntacticElementTarget::forInitialization(
 
 SyntacticElementTarget
 SyntacticElementTarget::forForEachStmt(ForEachStmt *stmt, DeclContext *dc,
-                                       bool ignoreWhereClause) {
-  SyntacticElementTarget target(stmt, dc, ignoreWhereClause);
+                                       bool ignoreWhereClause,
+                                       GenericEnvironment *packElementEnv) {
+  SyntacticElementTarget target(stmt, dc, ignoreWhereClause, packElementEnv);
   return target;
 }
 

lib/Sema/TypeCheckConstraints.cpp

@@ -896,7 +896,7 @@ bool TypeChecker::typeCheckPatternBinding(PatternBindingDecl *PBD,
   return hadError;
 }
 
-bool TypeChecker::typeCheckForEachBinding(DeclContext *dc, ForEachStmt *stmt) {
+bool TypeChecker::typeCheckForEachBinding(DeclContext *dc, ForEachStmt *stmt, GenericEnvironment *packElementEnv) {
   auto &Context = dc->getASTContext();
   FrontendStatsTracer statsTracer(Context.Stats, "typecheck-for-each", stmt);
   PrettyStackTraceStmt stackTrace(Context, "type-checking-for-each", stmt);
@@ -912,7 +912,7 @@ bool TypeChecker::typeCheckForEachBinding(DeclContext *dc, ForEachStmt *stmt) {
     return true;
   };
 
-  auto target = SyntacticElementTarget::forForEachStmt(stmt, dc);
+  auto target = SyntacticElementTarget::forForEachStmt(stmt, dc, false, packElementEnv);
   if (!typeCheckTarget(target))
     return failed();
 

lib/Sema/TypeCheckStmt.cpp

@@ -1000,6 +1000,9 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
 
   StmtChecker(DeclContext *DC) : Ctx(DC->getASTContext()), DC(DC) { }
 
+  // TODO: store a pointer union between a foreach and a decl context instead
+  llvm::SmallVector<GenericEnvironment*, 4> genericSigStack;
+
   //===--------------------------------------------------------------------===//
   // Helper Functions.
   //===--------------------------------------------------------------------===//
@@ -1434,17 +1437,27 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
   }
 
   Stmt *visitForEachStmt(ForEachStmt *S) {
-    if (TypeChecker::typeCheckForEachBinding(DC, S))
+    GenericEnvironment *genericSignature = genericSigStack.empty() ? nullptr : genericSigStack.back();
+
+    if (TypeChecker::typeCheckForEachBinding(DC, S, genericSignature))
       return nullptr;
 
     // Type-check the body of the loop.
     auto sourceFile = DC->getParentSourceFile();
     checkLabeledStmtShadowing(getASTContext(), sourceFile, S);
 
     BraceStmt *Body = S->getBody();
+
+    if (auto packExpansion = dyn_cast<PackExpansionExpr>(S->getParsedSequence()))
+      genericSigStack.push_back(packExpansion->getGenericEnvironment());
+
+    
     typeCheckStmt(Body);
     S->setBody(Body);
-    
+
+    if (isa<PackExpansionExpr>(S->getParsedSequence()))
+      genericSigStack.pop_back();
+
     return S;
   }
 

lib/Sema/TypeChecker.h

@@ -759,7 +759,7 @@ bool typeCheckPatternBinding(PatternBindingDecl *PBD, unsigned patternNumber,
 /// Type-check a for-each loop's pattern binding and sequence together.
 ///
 /// \returns true if a failure occurred.
-bool typeCheckForEachBinding(DeclContext *dc, ForEachStmt *stmt);
+bool typeCheckForEachBinding(DeclContext *dc, ForEachStmt *stmt, GenericEnvironment *packElementEnv);
 
 /// Compute the set of captures for the given function or closure.
 void computeCaptures(AnyFunctionRef AFR);

test/stmt/foreach.swift

@@ -330,4 +330,14 @@ do {
       // expected-error@-1 {{'where' clause in pack iteration is not supported}}
     }
   }
+  
+  func nested<each T, each U>(value: repeat each T, value1: repeat each U) {
+    for e1 in repeat each value {
+      for _ in [] {}
+      for e2 in repeat each value1 {
+        let y = e1 // Ok
+      }
+      let x = e1 // Ok
+    }
+  }
 }