[SE-0408] Enable Pack Iteration

include/swift/AST/Decl.h

@@ -5855,6 +5855,7 @@ enum class PropertyWrapperSynthesizedPropertyKind {
 class VarDecl : public AbstractStorageDecl {
   friend class NamingPatternRequest;
   NamedPattern *NamingPattern = nullptr;
+  GenericEnvironment *OpenedElementEnvironment = nullptr;
 
 public:
   enum class Introducer : uint8_t {
@@ -5982,6 +5983,13 @@ class VarDecl : public AbstractStorageDecl {
   NamedPattern *getNamingPattern() const;
   void setNamingPattern(NamedPattern *Pat);
 
+  GenericEnvironment *getOpenedElementEnvironment() const {
+    return OpenedElementEnvironment;
+  }
+  void setOpenedElementEnvironment(GenericEnvironment *Env) {
+    OpenedElementEnvironment = Env;
+  }
+
   /// If this is a VarDecl that does not belong to a CaseLabelItem's pattern,
   /// return this. Otherwise, this VarDecl must belong to a CaseStmt's
   /// CaseLabelItem. In that case, return the first case label item of the first

include/swift/AST/DiagnosticsSema.def

@@ -5834,8 +5834,8 @@ ERROR(expansion_not_allowed,none,
       "pack expansion %0 can only appear in a function parameter list, "
       "tuple element, or generic argument of a variadic type", (Type))
 ERROR(expansion_expr_not_allowed,none,
-      "value pack expansion can only appear inside a function argument list "
-      "or tuple element", ())
+      "value pack expansion can only appear inside a function argument list, "
+      "tuple element, or as the expression of a for-in loop", ())
 ERROR(invalid_expansion_argument,none,
       "cannot pass value pack expansion to non-pack parameter of type %0",
       (Type))
@@ -7715,5 +7715,12 @@ ERROR(referencebindings_binding_must_be_to_lvalue,none,
 ERROR(result_depends_on_no_result,none,
       "Incorrect use of %0 with no result", (StringRef))
 
+//------------------------------------------------------------------------------
+// MARK: Pack Iteration Diagnostics
+//------------------------------------------------------------------------------
+
+ERROR(pack_iteration_where_clause_not_supported, none,
+      "'where' clause in pack iteration is not supported", ())
+
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

include/swift/Sema/CSFix.h

@@ -440,6 +440,9 @@ enum class FixKind : uint8_t {
   /// Allow pack expansion expressions in a context that does not support them.
   AllowInvalidPackExpansion,
 
+  /// Ignore `where` clause in a for-in loop with a pack expansion expression.
+  IgnoreWhereClauseInPackIteration,
+
   /// Allow a pack expansion parameter of N elements to be matched
   /// with a single tuple literal argument of the same arity.
   DestructureTupleToMatchPackExpansionParameter,
@@ -2223,6 +2226,26 @@ class AllowInvalidPackExpansion final : public ConstraintFix {
   }
 };
 
+class IgnoreWhereClauseInPackIteration final : public ConstraintFix {
+  IgnoreWhereClauseInPackIteration(ConstraintSystem &cs,
+                                   ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::IgnoreWhereClauseInPackIteration, locator) {}
+
+public:
+  std::string getName() const override {
+    return "ignore where clause in pack iteration";
+  }
+
+  bool diagnose(const Solution &solution, bool asNote = false) const override;
+
+  static IgnoreWhereClauseInPackIteration *create(ConstraintSystem &cs,
+                                                  ConstraintLocator *locator);
+
+  static bool classof(const ConstraintFix *fix) {
+    return fix->getKind() == FixKind::IgnoreWhereClauseInPackIteration;
+  }
+};
+
 class CollectionElementContextualMismatch final
     : public ContextualMismatch,
       private llvm::TrailingObjects<CollectionElementContextualMismatch,

include/swift/Sema/ConstraintSystem.h

@@ -3759,6 +3759,11 @@ class ConstraintSystem {
                                        RememberChoice_t rememberChoice,
                                        ConstraintLocatorBuilder locator,
                                        ConstraintFix *compatFix = nullptr);
+
+  /// Add a materialize constraint for a pack expansion.
+  TypeVariableType *
+  addMaterializePackExpansionConstraint(Type patternType,
+                                        ConstraintLocatorBuilder locator);
 
   /// Add a disjunction constraint.
   void

include/swift/Sema/SyntacticElementTarget.h

@@ -22,15 +22,16 @@
 #include "swift/AST/Pattern.h"
 #include "swift/AST/Stmt.h"
 #include "swift/AST/TypeLoc.h"
+#include "swift/Basic/TaggedUnion.h"
 #include "swift/Sema/ConstraintLocator.h"
 #include "swift/Sema/ContextualTypeInfo.h"
 
 namespace swift {
 
 namespace constraints {
-/// Describes information about a for-each loop that needs to be tracked
-/// within the constraint system.
-struct ForEachStmtInfo {
+/// Describes information specific to a sequence
+/// in a for-each loop.
+struct SequenceIterationInfo {
   /// The type of the sequence.
   Type sequenceType;
 
@@ -47,6 +48,19 @@ struct ForEachStmtInfo {
   Expr *nextCall;
 };
 
+/// Describes information specific to a pack expansion expression
+/// in a for-each loop.
+struct PackIterationInfo {
+  /// The type of the pattern that matches the elements.
+  Type patternType;
+};
+
+/// Describes information about a for-each loop that needs to be tracked
+/// within the constraint system.
+struct ForEachStmtInfo : TaggedUnion<SequenceIterationInfo, PackIterationInfo> {
+  using TaggedUnion::TaggedUnion;
+};
+
 /// Describes the target to which a constraint system's solution can be
 /// applied.
 class SyntacticElementTarget {

lib/AST/ASTVerifier.cpp

@@ -795,6 +795,13 @@ class Verifier : public ASTWalker {
       if (!shouldVerify(cast<Stmt>(S)))
         return false;
 
+      if (auto *expansion =
+              dyn_cast<PackExpansionExpr>(S->getParsedSequence())) {
+        if (!shouldVerify(expansion)) {
+          return false;
+        }
+      }
+
       if (!S->getElementExpr())
         return true;
 
@@ -804,6 +811,11 @@ class Verifier : public ASTWalker {
     }
 
     void cleanup(ForEachStmt *S) {
+      if (auto *expansion =
+              dyn_cast<PackExpansionExpr>(S->getParsedSequence())) {
+        cleanup(expansion);
+      }
+
       if (!S->getElementExpr())
         return;
 
@@ -2605,6 +2617,16 @@ class Verifier : public ASTWalker {
         abort();
       }
 
+      // If we are performing pack iteration, variables have to carry the
+      // generic environment. Catching the missing environment here will prevent
+      // the code from being lowered.
+      if (var->getTypeInContext()->is<ErrorType>()) {
+        Out << "VarDecl is missing a Generic Environment: ";
+        var->getInterfaceType().print(Out);
+        Out << "\n";
+        abort();
+      }
+
       // The fact that this is *directly* be a reference storage type
       // cuts the code down quite a bit in getTypeOfReference.
       if (var->getAttrs().hasAttribute<ReferenceOwnershipAttr>() !=

lib/AST/ASTWalker.cpp

@@ -1905,11 +1905,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

@@ -7135,6 +7135,11 @@ 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 *env = getOpenedElementEnvironment())
+    return GenericEnvironment::mapTypeIntoContext(env, getInterfaceType());
+
   return getDeclContext()->mapTypeIntoContext(getInterfaceType());
 }
 

lib/AST/GenericSignature.cpp

@@ -105,8 +105,16 @@ void GenericSignatureImpl::forEachParam(
 
   for (auto req : getRequirements()) {
     GenericTypeParamType *gp;
+    bool isCanonical = false;
     switch (req.getKind()) {
     case RequirementKind::SameType: {
+      if (req.getSecondType()->isParameterPack() != 
+          req.getFirstType()->isParameterPack()) {
+        // This is a same-element requirement, which does not make
+        // type parameters non-canonical.
+        isCanonical = true;
+      }
+
       if (auto secondGP = req.getSecondType()->getAs<GenericTypeParamType>()) {
         // If two generic parameters are same-typed, then the right-hand one
         // is non-canonical.
@@ -136,7 +144,7 @@ void GenericSignatureImpl::forEachParam(
     }
 
     unsigned index = GenericParamKey(gp).findIndexIn(genericParams);
-    genericParamsAreCanonical[index] = false;
+    genericParamsAreCanonical[index] = isCanonical;
   }
 
   // Call the callback with each parameter and the result of the above analysis.

lib/AST/Stmt.cpp

@@ -446,6 +446,9 @@ void ForEachStmt::setPattern(Pattern *p) {
 }
 
 Expr *ForEachStmt::getTypeCheckedSequence() const {
+  if (auto *expansion = dyn_cast<PackExpansionExpr>(getParsedSequence()))
+    return expansion;
+
   return iteratorVar ? iteratorVar->getInit(/*index=*/0) : nullptr;
 }
 

lib/IDE/CodeCompletion.cpp

@@ -1062,6 +1062,8 @@ void CodeCompletionCallbacksImpl::addKeywords(CodeCompletionResultSink &Sink,
     addSuperKeyword(Sink, CurDeclContext);
     addExprKeywords(Sink, CurDeclContext);
     addAnyTypeKeyword(Sink, CurDeclContext->getASTContext().TheAnyType);
+    if (Kind == CompletionKind::ForEachSequence)
+      addKeyword(Sink, "repeat", CodeCompletionKeywordKind::kw_repeat);
     break;
 
   case CompletionKind::CallArg:

lib/SILGen/SILGenFunction.h

@@ -2724,26 +2724,24 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   ///
   ///   This function will be called within a cleanups scope and with
   ///   InnermostPackExpansion set up properly for the context.
-  void emitDynamicPackLoop(SILLocation loc,
-                           CanPackType formalPackType,
-                           unsigned componentIndex,
-                           SILValue startingAfterIndexWithinComponent,
-                           SILValue limitWithinComponent,
-                           GenericEnvironment *openedElementEnv,
-                           bool reverse,
-                        llvm::function_ref<void(SILValue indexWithinComponent,
-                                                SILValue packExpansionIndex,
-                                                SILValue packIndex)> emitBody);
+  void emitDynamicPackLoop(
+      SILLocation loc, CanPackType formalPackType, unsigned componentIndex,
+      SILValue startingAfterIndexWithinComponent, SILValue limitWithinComponent,
+      GenericEnvironment *openedElementEnv, bool reverse,
+      llvm::function_ref<void(SILValue indexWithinComponent,
+                              SILValue packExpansionIndex, SILValue packIndex)>
+          emitBody,
+      SILBasicBlock *loopLatch = nullptr);
 
   /// A convenience version of dynamic pack loop that visits an entire
   /// pack expansion component in forward order.
-  void emitDynamicPackLoop(SILLocation loc,
-                           CanPackType formalPackType,
-                           unsigned componentIndex,
-                           GenericEnvironment *openedElementEnv,
-                        llvm::function_ref<void(SILValue indexWithinComponent,
-                                                SILValue packExpansionIndex,
-                                                SILValue packIndex)> emitBody);
+  void emitDynamicPackLoop(
+      SILLocation loc, CanPackType formalPackType, unsigned componentIndex,
+      GenericEnvironment *openedElementEnv,
+      llvm::function_ref<void(SILValue indexWithinComponent,
+                              SILValue packExpansionIndex, SILValue packIndex)>
+          emitBody,
+      SILBasicBlock *loopLatch = nullptr);
 
   /// Emit a transform on each element of a pack-expansion component
   /// of a pack, write the result into a pack-expansion component of

lib/SILGen/SILGenPack.cpp

@@ -662,28 +662,27 @@ void SILGenFunction::projectTupleElementsToPack(SILLocation loc,
   });
 }
 
-void SILGenFunction::emitDynamicPackLoop(SILLocation loc,
-                                         CanPackType formalPackType,
-                                         unsigned componentIndex,
-                                         GenericEnvironment *openedElementEnv,
-                      llvm::function_ref<void(SILValue indexWithinComponent,
-                                              SILValue packExpansionIndex,
-                                              SILValue packIndex)> emitBody) {
+void SILGenFunction::emitDynamicPackLoop(
+    SILLocation loc, CanPackType formalPackType, unsigned componentIndex,
+    GenericEnvironment *openedElementEnv,
+    llvm::function_ref<void(SILValue indexWithinComponent,
+                            SILValue packExpansionIndex, SILValue packIndex)>
+        emitBody,
+    SILBasicBlock *loopLatch) {
   return emitDynamicPackLoop(loc, formalPackType, componentIndex,
                              /*startAfter*/ SILValue(), /*limit*/ SILValue(),
-                             openedElementEnv, /*reverse*/false, emitBody);
+                             openedElementEnv, /*reverse*/ false, emitBody,
+                             loopLatch);
 }
 
-void SILGenFunction::emitDynamicPackLoop(SILLocation loc,
-                                         CanPackType formalPackType,
-                                         unsigned componentIndex,
-                                         SILValue startingAfterIndexInComponent,
-                                         SILValue limitWithinComponent,
-                                         GenericEnvironment *openedElementEnv,
-                                         bool reverse,
-                      llvm::function_ref<void(SILValue indexWithinComponent,
-                                              SILValue packExpansionIndex,
-                                              SILValue packIndex)> emitBody) {
+void SILGenFunction::emitDynamicPackLoop(
+    SILLocation loc, CanPackType formalPackType, unsigned componentIndex,
+    SILValue startingAfterIndexInComponent, SILValue limitWithinComponent,
+    GenericEnvironment *openedElementEnv, bool reverse,
+    llvm::function_ref<void(SILValue indexWithinComponent,
+                            SILValue packExpansionIndex, SILValue packIndex)>
+        emitBody,
+    SILBasicBlock *loopLatch) {
   assert(isa<PackExpansionType>(formalPackType.getElementType(componentIndex)));
   assert((!startingAfterIndexInComponent || !reverse) &&
          "cannot reverse with a starting index");
@@ -764,6 +763,7 @@ void SILGenFunction::emitDynamicPackLoop(SILLocation loc,
   // the incoming index - 1 if reverse)
   SILValue curIndex = incomingIndex;
   if (reverse) {
+    assert(!loopLatch && "Only forward iteration supported with loop latch");
     curIndex = B.createBuiltinBinaryFunction(loc, "sub", wordTy, wordTy,
                                              { incomingIndex, one });
   }
@@ -791,6 +791,13 @@ void SILGenFunction::emitDynamicPackLoop(SILLocation loc,
   {
     FullExpr scope(Cleanups, CleanupLocation(loc));
     emitBody(curIndex, packExpansionIndex, packIndex);
+    if (loopLatch) {
+      B.createBranch(loc, loopLatch);
+    }
+  }
+
+  if (loopLatch) {
+    B.emitBlock(loopLatch);
   }
 
   // The index to pass to the loop condition block (the current index + 1