[Variadic Generics][WIP] Enable Pack Iteration

Author: simanerush

include/swift/Sema/SyntacticElementTarget.h

@@ -27,9 +27,9 @@
 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;
 
@@ -46,6 +46,23 @@ struct ForEachStmtInfo {
   Expr *nextCall;
 };
 
+/// Describes information specific to a pack expansion expression
+/// in a for-each loop.
+struct PackIterationInfo {
+  /// The pack expansion expression.
+  PackExpansionExpr *expansion;
+
+  /// The mapping from pack types in the outer context to element types inside
+  /// the for-each loop.
+  GenericEnvironment *openedElementEnvironment;
+};
+
+/// 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/Sema/CSApply.cpp

@@ -9139,20 +9139,22 @@ static llvm::Optional<SyntacticElementTarget> applySolutionToForEachStmt(
   auto *parsedSequence = stmt->getParsedSequence();
   bool isAsync = stmt->getAwaitLoc().isValid();
 
-  // Simplify the various types.
-  forEachStmtInfo.sequenceType =
-      solution.simplifyType(forEachStmtInfo.sequenceType);
-  forEachStmtInfo.elementType =
-      solution.simplifyType(forEachStmtInfo.elementType);
-  forEachStmtInfo.initType =
-      solution.simplifyType(forEachStmtInfo.initType);
-
   auto &cs = solution.getConstraintSystem();
   auto *dc = target.getDeclContext();
 
-  // First, let's apply the solution to the sequence expression.
-  {
-    auto *makeIteratorVar = forEachStmtInfo.makeIteratorVar;
+  if (forEachStmtInfo.isa<SequenceIterationInfo>()) {
+    auto sequenceIterationInfo =
+        *forEachStmtInfo.dyn_cast<SequenceIterationInfo>();
+    // Simplify the various types.
+    sequenceIterationInfo.sequenceType =
+        solution.simplifyType(sequenceIterationInfo.sequenceType);
+    sequenceIterationInfo.elementType =
+        solution.simplifyType(sequenceIterationInfo.elementType);
+    sequenceIterationInfo.initType =
+        solution.simplifyType(sequenceIterationInfo.initType);
+
+    // First, let's apply the solution to the expression.
+    auto *makeIteratorVar = sequenceIterationInfo.makeIteratorVar;
 
     auto makeIteratorTarget = *cs.getTargetFor({makeIteratorVar, /*index=*/0});
 
@@ -9167,127 +9169,126 @@ static llvm::Optional<SyntacticElementTarget> applySolutionToForEachStmt(
     }
 
     stmt->setIteratorVar(makeIteratorVar);
-  }
 
-  // Now, `$iterator.next()` call.
-  {
-    auto nextTarget = *cs.getTargetFor(forEachStmtInfo.nextCall);
+    // Now, `$iterator.next()` call.
+    {
+      auto nextTarget = *cs.getTargetFor(sequenceIterationInfo.nextCall);
 
-    auto rewrittenTarget = rewriteTarget(nextTarget);
-    if (!rewrittenTarget)
-      return llvm::None;
+      auto rewrittenTarget = rewriteTarget(nextTarget);
+      if (!rewrittenTarget)
+        return llvm::None;
 
-    Expr *nextCall = rewrittenTarget->getAsExpr();
-    // Wrap a call to `next()` into `try await` since `AsyncIteratorProtocol`
-    // witness could be `async throws`.
-    if (isAsync) {
-      // Cannot use `forEachChildExpr` here because we need to
-      // to wrap a call in `try` and then stop immediately after.
-      struct TryInjector : ASTWalker {
-        ASTContext &C;
-        const Solution &S;
+      Expr *nextCall = rewrittenTarget->getAsExpr();
+      // Wrap a call to `next()` into `try await` since `AsyncIteratorProtocol`
+      // witness could be `async throws`.
+      if (isAsync) {
+        // Cannot use `forEachChildExpr` here because we need to
+        // to wrap a call in `try` and then stop immediately after.
+        struct TryInjector : ASTWalker {
+          ASTContext &C;
+          const Solution &S;
 
-        bool ShouldStop = false;
+          bool ShouldStop = false;
 
-        TryInjector(ASTContext &ctx, const Solution &solution)
-            : C(ctx), S(solution) {}
+          TryInjector(ASTContext &ctx, const Solution &solution)
+              : C(ctx), S(solution) {}
 
-        MacroWalking getMacroWalkingBehavior() const override {
-          return MacroWalking::Expansion;
-        }
+          MacroWalking getMacroWalkingBehavior() const override {
+            return MacroWalking::Expansion;
+          }
 
-        PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
-          if (ShouldStop)
-            return Action::Stop();
-
-          if (auto *call = dyn_cast<CallExpr>(E)) {
-            // There is a single call expression in `nextCall`.
-            ShouldStop = true;
-
-            auto nextRefType =
-              S.getResolvedType(call->getFn())->castTo<FunctionType>();
-
-            // If the inferred witness is throwing, we need to wrap the call
-            // into `try` expression.
-            if (nextRefType->isThrowing()) {
-              auto *tryExpr = TryExpr::createImplicit(
-                  C, /*tryLoc=*/call->getStartLoc(), call, call->getType());
-              // Cannot stop here because we need to make sure that
-              // the new expression gets injected into AST.
-              return Action::SkipChildren(tryExpr);
+          PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
+            if (ShouldStop)
+              return Action::Stop();
+
+            if (auto *call = dyn_cast<CallExpr>(E)) {
+              // There is a single call expression in `nextCall`.
+              ShouldStop = true;
+
+              auto nextRefType =
+                  S.getResolvedType(call->getFn())->castTo<FunctionType>();
+
+              // If the inferred witness is throwing, we need to wrap the call
+              // into `try` expression.
+              if (nextRefType->isThrowing()) {
+                auto *tryExpr = TryExpr::createImplicit(
+                    C, /*tryLoc=*/call->getStartLoc(), call, call->getType());
+                // Cannot stop here because we need to make sure that
+                // the new expression gets injected into AST.
+                return Action::SkipChildren(tryExpr);
+              }
             }
+
+            return Action::Continue(E);
           }
+        };
 
-          return Action::Continue(E);
-        }
-      };
+        nextCall->walk(TryInjector(cs.getASTContext(), solution));
+      }
 
-      nextCall->walk(TryInjector(cs.getASTContext(), solution));
+      stmt->setNextCall(nextCall);
     }
 
-    stmt->setNextCall(nextCall);
-  }
-
-  // Coerce the pattern to the element type.
-  {
-    TypeResolutionOptions options(TypeResolverContext::ForEachStmt);
-    options |= TypeResolutionFlags::OverrideType;
+    // Coerce the pattern to the element type.
+    {
+      TypeResolutionOptions options(TypeResolverContext::ForEachStmt);
+      options |= TypeResolutionFlags::OverrideType;
 
-    auto tryRewritePattern = [&](Pattern *EP, Type ty) {
-      return ::tryRewriteExprPattern(EP, solution, ty, rewriteTarget);
-    };
+      auto tryRewritePattern = [&](Pattern *EP, Type ty) {
+        return ::tryRewriteExprPattern(EP, solution, ty, rewriteTarget);
+      };
 
-    // Apply the solution to the pattern as well.
-    auto contextualPattern = target.getContextualPattern();
-    auto coercedPattern = TypeChecker::coercePatternToType(
-        contextualPattern, forEachStmtInfo.initType, options,
-        tryRewritePattern);
-    if (!coercedPattern)
-      return llvm::None;
+      // Apply the solution to the pattern as well.
+      auto contextualPattern = target.getContextualPattern();
+      auto coercedPattern = TypeChecker::coercePatternToType(
+          contextualPattern, sequenceIterationInfo.initType, options,
+          tryRewritePattern);
+      if (!coercedPattern)
+        return llvm::None;
 
-    stmt->setPattern(coercedPattern);
-    resultTarget.setPattern(coercedPattern);
-  }
+      stmt->setPattern(coercedPattern);
+      resultTarget.setPattern(coercedPattern);
+    }
 
-  // Apply the solution to the filtering condition, if there is one.
-  if (auto *whereExpr = stmt->getWhere()) {
-    auto whereTarget = *cs.getTargetFor(whereExpr);
+    // Apply the solution to the filtering condition, if there is one.
+    if (auto *whereExpr = stmt->getWhere()) {
+      auto whereTarget = *cs.getTargetFor(whereExpr);
 
-    auto rewrittenTarget = rewriteTarget(whereTarget);
-    if (!rewrittenTarget)
-      return llvm::None;
+      auto rewrittenTarget = rewriteTarget(whereTarget);
+      if (!rewrittenTarget)
+        return llvm::None;
 
-    stmt->setWhere(rewrittenTarget->getAsExpr());
-  }
+      stmt->setWhere(rewrittenTarget->getAsExpr());
+    }
 
-  // Convert that llvm::Optional<Element> value to the type of the pattern.
-  auto optPatternType = OptionalType::get(forEachStmtInfo.initType);
-  Type nextResultType = OptionalType::get(forEachStmtInfo.elementType);
-  if (!optPatternType->isEqual(nextResultType)) {
-    ASTContext &ctx = cs.getASTContext();
-    OpaqueValueExpr *elementExpr = new (ctx) OpaqueValueExpr(
-        stmt->getInLoc(), nextResultType->getOptionalObjectType(),
-        /*isPlaceholder=*/true);
-    Expr *convertElementExpr = elementExpr;
-    if (TypeChecker::typeCheckExpression(
-            convertElementExpr, dc,
-            /*contextualInfo=*/{forEachStmtInfo.initType, CTP_CoerceOperand})
-            .isNull()) {
-      return llvm::None;
+    // Convert that llvm::Optional<Element> value to the type of the pattern.
+    auto optPatternType = OptionalType::get(sequenceIterationInfo.initType);
+    Type nextResultType = OptionalType::get(sequenceIterationInfo.elementType);
+    if (!optPatternType->isEqual(nextResultType)) {
+      ASTContext &ctx = cs.getASTContext();
+      OpaqueValueExpr *elementExpr = new (ctx) OpaqueValueExpr(
+          stmt->getInLoc(), nextResultType->getOptionalObjectType(),
+          /*isPlaceholder=*/true);
+      Expr *convertElementExpr = elementExpr;
+      if (TypeChecker::typeCheckExpression(
+              convertElementExpr, dc,
+              /*contextualInfo=*/
+              {sequenceIterationInfo.initType, CTP_CoerceOperand})
+              .isNull()) {
+        return llvm::None;
+      }
+      elementExpr->setIsPlaceholder(false);
+      stmt->setElementExpr(elementExpr);
+      stmt->setConvertElementExpr(convertElementExpr);
     }
-    elementExpr->setIsPlaceholder(false);
-    stmt->setElementExpr(elementExpr);
-    stmt->setConvertElementExpr(convertElementExpr);
-  }
 
-  // Get the conformance of the sequence type to the Sequence protocol.
-  {
+    // Get the conformance of the sequence type to the Sequence protocol.
     auto sequenceProto = TypeChecker::getProtocol(
         cs.getASTContext(), stmt->getForLoc(),
         stmt->getAwaitLoc().isValid() ? KnownProtocolKind::AsyncSequence
                                       : KnownProtocolKind::Sequence);
 
-    auto type = forEachStmtInfo.sequenceType->getRValueType();
+    auto type = sequenceIterationInfo.sequenceType->getRValueType();
     if (type->isExistentialType()) {
       auto *contextualLoc = solution.getConstraintLocator(
           parsedSequence, LocatorPathElt::ContextualType(CTP_ForEachSequence));
@@ -9300,6 +9301,48 @@ static llvm::Optional<SyntacticElementTarget> applySolutionToForEachStmt(
     stmt->setSequenceConformance(sequenceConformance);
   }
 
+  if (forEachStmtInfo.isa<PackIterationInfo>()) {
+    auto packIterationInfo = *forEachStmtInfo.dyn_cast<PackIterationInfo>();
+
+    // First, let's apply the solution to the expression.
+    auto makeSequenceTarget = *cs.getTargetFor(parsedSequence);
+    auto rewrittenTarget = rewriteTarget(makeSequenceTarget);
+    if (!rewrittenTarget)
+      return llvm::None;
+
+    // Coerce the pattern to the element type.
+    {
+      TypeResolutionOptions options(TypeResolverContext::ForEachStmt);
+      options |= TypeResolutionFlags::OverrideType;
+
+      auto tryRewritePattern = [&](Pattern *EP, Type ty) {
+        return ::tryRewriteExprPattern(EP, solution, ty, rewriteTarget);
+      };
+
+      // Apply the solution to the pattern as well.
+      auto contextualPattern = target.getContextualPattern();
+      auto coercedPattern = TypeChecker::coercePatternToType(
+          contextualPattern, packIterationInfo.expansion->getType(), options,
+          tryRewritePattern);
+      if (!coercedPattern)
+        return llvm::None;
+
+      stmt->setPattern(coercedPattern);
+      resultTarget.setPattern(coercedPattern);
+    }
+
+    // Apply the solution to the filtering condition, if there is one.
+    if (auto *whereExpr = stmt->getWhere()) {
+      auto whereTarget = *cs.getTargetFor(whereExpr);
+
+      auto rewrittenTarget = rewriteTarget(whereTarget);
+      if (!rewrittenTarget)
+        return llvm::None;
+
+      stmt->setWhere(rewrittenTarget->getAsExpr());
+    }
+  }
+
   return resultTarget;
 }