[Variadic Generics] Add a contextual each keyword for pack element references in pack expansion patterns.

include/swift/AST/Expr.h

@@ -3515,6 +3515,50 @@ class VarargExpansionExpr : public Expr {
   }
 };
 
+/// A pack element expression spelled with the contextual \c each
+/// keyword applied to a pack reference expression.
+///
+/// \code
+///  func zip<T..., U...>(t: (each T)..., u: (each U)...) {
+///    let zipped = (each t, each u)...
+///  }
+/// \endcode
+///
+/// Pack elements can only appear in the pattern expression of a
+/// \c PackExpansionExpr.
+class PackElementExpr final : public Expr {
+  SourceLoc EachLoc;
+  Expr *PackRefExpr;
+
+  PackElementExpr(SourceLoc eachLoc, Expr *packRefExpr,
+                  bool implicit = false, Type type = Type())
+      : Expr(ExprKind::PackElement, implicit, type),
+        EachLoc(eachLoc), PackRefExpr(packRefExpr) {}
+
+public:
+  static PackElementExpr *create(ASTContext &ctx, SourceLoc eachLoc,
+                                 Expr *packRefExpr, bool implicit = false,
+                                 Type type = Type());
+
+  Expr *getPackRefExpr() const { return PackRefExpr; }
+
+  void setPackRefExpr(Expr *packRefExpr) {
+    PackRefExpr = packRefExpr;
+  }
+
+  SourceLoc getStartLoc() const {
+    return EachLoc;
+  }
+
+  SourceLoc getEndLoc() const {
+    return PackRefExpr->getEndLoc();
+  }
+
+  static bool classof(const Expr *E) {
+    return E->getKind() == ExprKind::PackElement;
+  }
+};
+
 /// A pack expansion expression is a pattern expression followed by
 /// the expansion operator '...'. The pattern expression contains
 /// references to parameter packs of length N, and the expansion
@@ -3525,51 +3569,41 @@ class VarargExpansionExpr : public Expr {
 /// call argument lists, the elements of a tuple value, and the source
 /// of a for-in loop.
 class PackExpansionExpr final : public Expr,
-    private llvm::TrailingObjects<PackExpansionExpr,
-                                  OpaqueValueExpr *, Expr *> {
+    private llvm::TrailingObjects<PackExpansionExpr, PackElementExpr *> {
   friend TrailingObjects;
 
   Expr *PatternExpr;
   SourceLoc DotsLoc;
   GenericEnvironment *Environment;
 
   PackExpansionExpr(Expr *patternExpr,
-                    ArrayRef<OpaqueValueExpr *> opaqueValues,
-                    ArrayRef<Expr *> bindings,
+                    ArrayRef<PackElementExpr *> packElements,
                     SourceLoc dotsLoc,
                     GenericEnvironment *environment,
                     bool implicit, Type type)
     : Expr(ExprKind::PackExpansion, implicit, type),
       PatternExpr(patternExpr), DotsLoc(dotsLoc), Environment(environment) {
-    assert(opaqueValues.size() == bindings.size());
-    Bits.PackExpansionExpr.NumBindings = opaqueValues.size();
+    Bits.PackExpansionExpr.NumBindings = packElements.size();
 
     assert(Bits.PackExpansionExpr.NumBindings > 0 &&
            "PackExpansionExpr must have pack references");
 
-    std::uninitialized_copy(opaqueValues.begin(), opaqueValues.end(),
-                            getTrailingObjects<OpaqueValueExpr *>());
-    std::uninitialized_copy(bindings.begin(), bindings.end(),
-                            getTrailingObjects<Expr *>());
-  }
-
-  size_t numTrailingObjects(OverloadToken<OpaqueValueExpr *>) const {
-    return getNumBindings();
+    std::uninitialized_copy(packElements.begin(), packElements.end(),
+                            getTrailingObjects<PackElementExpr *>());
   }
 
-  size_t numTrailingObjects(OverloadToken<Expr *>) const {
+  size_t numTrailingObjects(OverloadToken<PackElementExpr *>) const {
     return getNumBindings();
   }
 
-  MutableArrayRef<Expr *> getMutableBindings() {
-    return {getTrailingObjects<Expr *>(), getNumBindings()};
+  MutableArrayRef<PackElementExpr *> getMutableBindings() {
+    return {getTrailingObjects<PackElementExpr *>(), getNumBindings()};
   }
 
 public:
   static PackExpansionExpr *create(ASTContext &ctx,
                                    Expr *patternExpr,
-                                   ArrayRef<OpaqueValueExpr *> opaqueValues,
-                                   ArrayRef<Expr *> bindings,
+                                   ArrayRef<PackElementExpr *> packElements,
                                    SourceLoc dotsLoc,
                                    GenericEnvironment *environment,
                                    bool implicit = false,
@@ -3585,15 +3619,11 @@ class PackExpansionExpr final : public Expr,
     return Bits.PackExpansionExpr.NumBindings;
   }
 
-  ArrayRef<OpaqueValueExpr *> getOpaqueValues() {
-    return {getTrailingObjects<OpaqueValueExpr *>(), getNumBindings()};
-  }
-
-  ArrayRef<Expr *> getBindings() {
-    return {getTrailingObjects<Expr *>(), getNumBindings()};
+  ArrayRef<PackElementExpr *> getPackElements() {
+    return {getTrailingObjects<PackElementExpr *>(), getNumBindings()};
   }
 
-  void setBinding(unsigned i, Expr *e) {
+  void setBinding(unsigned i, PackElementExpr *e) {
     getMutableBindings()[i] = e;
   }
 

include/swift/AST/ExprNodes.def

@@ -130,6 +130,7 @@ ABSTRACT_EXPR(AbstractClosure, Expr)
 EXPR(InOut, Expr)
 EXPR(VarargExpansion, Expr)
 EXPR(PackExpansion, Expr)
+EXPR(PackElement, Expr)
 EXPR(DynamicType, Expr)
 EXPR(RebindSelfInConstructor, Expr)
 EXPR(OpaqueValue, Expr)

include/swift/Sema/ConstraintLocator.h

@@ -687,6 +687,21 @@ class LocatorPathElt::PackElement final : public StoredIntegerElement<1> {
   }
 };
 
+class LocatorPathElt::OpenedPackElement final
+    : public StoredPointerElement<GenericEnvironment> {
+public:
+  OpenedPackElement(GenericEnvironment *env)
+      : StoredPointerElement(PathElementKind::OpenedPackElement, env) {}
+
+  GenericEnvironment *getGenericEnvironment() const {
+    return getStoredPointer();
+  }
+
+  static bool classof(const LocatorPathElt *elt) {
+    return elt->getKind() == PathElementKind::OpenedPackElement;
+  }
+};
+
 class LocatorPathElt::KeyPathComponent final : public StoredIntegerElement<1> {
 public:
   KeyPathComponent(unsigned index)

include/swift/Sema/ConstraintLocatorPathElts.def

@@ -195,6 +195,9 @@ CUSTOM_LOCATOR_PATH_ELT(PackType)
 /// An element of a pack type - the T in <T, U, V, ...>
 CUSTOM_LOCATOR_PATH_ELT(PackElement)
 
+/// Stores the generic environment for an opened pack element.
+CUSTOM_LOCATOR_PATH_ELT(OpenedPackElement)
+
 /// The shape of a parameter pack.
 SIMPLE_LOCATOR_PATH_ELT(PackShape)
 

lib/AST/ASTDumper.cpp

@@ -2504,6 +2504,12 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
     PrintWithColorRAII(OS, ParenthesisColor) << ')';
   }
 
+  void visitPackElementExpr(PackElementExpr *E) {
+    printCommon(E, "pack_element_expr") << "\n";
+    printRec(E->getPackRefExpr());
+    PrintWithColorRAII(OS, ParenthesisColor) << ')';
+  }
+
   void visitForceTryExpr(ForceTryExpr *E) {
     printCommon(E, "force_try_expr");
     OS << '\n';

lib/AST/ASTPrinter.cpp

@@ -4858,6 +4858,10 @@ void PrintAST::visitPackExpansionExpr(PackExpansionExpr *expr) {
   visit(expr->getPatternExpr());
 }
 
+void PrintAST::visitPackElementExpr(PackElementExpr *expr) {
+  visit(expr->getPackRefExpr());
+}
+
 void PrintAST::visitArchetypeToSuperExpr(ArchetypeToSuperExpr *expr) {
 }
 

lib/AST/ASTVerifier.cpp

@@ -836,12 +836,6 @@ class Verifier : public ASTWalker {
         return false;
 
       Generics.push_back(expr->getGenericEnvironment());
-
-      for (auto *placeholder : expr->getOpaqueValues()) {
-        assert(!OpaqueValues.count(placeholder));
-        OpaqueValues[placeholder] = 0;
-      }
-
       return true;
     }
 
@@ -854,13 +848,6 @@ class Verifier : public ASTWalker {
       verifyCheckedAlwaysBase(E);
     }
 
-    void cleanup(PackExpansionExpr *expr) {
-      for (auto *placeholder : expr->getOpaqueValues()) {
-        assert(OpaqueValues.count(placeholder));
-        OpaqueValues.erase(placeholder);
-      }
-    }
-
     bool shouldVerify(MakeTemporarilyEscapableExpr *expr) {
       if (!shouldVerify(cast<Expr>(expr)))
         return false;

lib/AST/ASTWalker.cpp

@@ -859,6 +859,14 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     return nullptr;
   }
 
+  Expr *visitPackElementExpr(PackElementExpr *E) {
+    if (Expr *pattern = doIt(E->getPackRefExpr())) {
+      E->setPackRefExpr(pattern);
+      return E;
+    }
+    return nullptr;
+  }
+
   Expr *visitSequenceExpr(SequenceExpr *E) {
     for (unsigned i = 0, e = E->getNumElements(); i != e; ++i)
       if (Expr *Elt = doIt(E->getElement(i)))

lib/AST/Expr.cpp

@@ -394,6 +394,7 @@ ConcreteDeclRef Expr::getReferencedDecl(bool stopAtParenExpr) const {
 
   NO_REFERENCE(VarargExpansion);
   NO_REFERENCE(PackExpansion);
+  NO_REFERENCE(PackElement);
   NO_REFERENCE(DynamicType);
 
   PASS_THROUGH_REFERENCE(RebindSelfInConstructor, getSubExpr);
@@ -750,6 +751,7 @@ bool Expr::canAppendPostfixExpression(bool appendingPostfixOperator) const {
   case ExprKind::MakeTemporarilyEscapable:
   case ExprKind::VarargExpansion:
   case ExprKind::PackExpansion:
+  case ExprKind::PackElement:
     return false;
 
   case ExprKind::Call:
@@ -927,6 +929,7 @@ bool Expr::isValidParentOfTypeExpr(Expr *typeExpr) const {
   case ExprKind::InOut:
   case ExprKind::VarargExpansion:
   case ExprKind::PackExpansion:
+  case ExprKind::PackElement:
   case ExprKind::DynamicType:
   case ExprKind::RebindSelfInConstructor:
   case ExprKind::OpaqueValue:
@@ -1245,19 +1248,23 @@ VarargExpansionExpr *VarargExpansionExpr::createArrayExpansion(ASTContext &ctx,
 
 PackExpansionExpr *
 PackExpansionExpr::create(ASTContext &ctx, Expr *patternExpr,
-                          ArrayRef<OpaqueValueExpr *> opaqueValues,
-                          ArrayRef<Expr *> bindings, SourceLoc dotsLoc,
-                          GenericEnvironment *environment,
+                          ArrayRef<PackElementExpr *> packElements,
+                          SourceLoc dotsLoc, GenericEnvironment *environment,
                           bool implicit, Type type) {
   size_t size =
-      totalSizeToAlloc<OpaqueValueExpr *, Expr *>(opaqueValues.size(),
-                                                  bindings.size());
+      totalSizeToAlloc<PackElementExpr *>(packElements.size());
   void *mem = ctx.Allocate(size, alignof(PackExpansionExpr));
-  return ::new (mem) PackExpansionExpr(patternExpr, opaqueValues,
-                                       bindings, dotsLoc, environment,
+  return ::new (mem) PackExpansionExpr(patternExpr, packElements,
+                                       dotsLoc, environment,
                                        implicit, type);
 }
 
+PackElementExpr *
+PackElementExpr::create(ASTContext &ctx, SourceLoc eachLoc, Expr *packRefExpr,
+                        bool implicit, Type type) {
+  return new (ctx) PackElementExpr(eachLoc, packRefExpr, implicit, type);
+}
+
 SequenceExpr *SequenceExpr::create(ASTContext &ctx, ArrayRef<Expr*> elements) {
   assert(elements.size() & 1 && "even number of elements in sequence");
   size_t bytes = totalSizeToAlloc<Expr *>(elements.size());

lib/Parse/ParseExpr.cpp

@@ -1612,6 +1612,19 @@ ParserResult<Expr> Parser::parseExprPrimary(Diag<> ID, bool isExprBasic) {
       return makeParserResult(typeExpr);
     }
 
+    // 'each' followed by another identifier is a pack element expr.
+    if (Tok.isContextualKeyword("each") &&
+        peekToken().is(tok::identifier) &&
+        !peekToken().isAtStartOfLine()) {
+      SourceLoc loc = consumeToken();
+      ParserResult<Expr> ref = parseExpr(ID);
+      if (ref.isNull())
+        return ref;
+
+      auto *packRef = PackElementExpr::create(Context, loc, ref.get());
+      return makeParserResult(packRef);
+    }
+
     LLVM_FALLTHROUGH;
   case tok::kw_Self:     // Self
     return parseExprIdentifier();

lib/SILGen/SILGenExpr.cpp

@@ -444,6 +444,7 @@ namespace {
              SGFContext C);
     RValue visitBridgeToObjCExpr(BridgeToObjCExpr *E, SGFContext C);
     RValue visitPackExpansionExpr(PackExpansionExpr *E, SGFContext C);
+    RValue visitPackElementExpr(PackElementExpr *E, SGFContext C);
     RValue visitBridgeFromObjCExpr(BridgeFromObjCExpr *E, SGFContext C);
     RValue visitConditionalBridgeFromObjCExpr(ConditionalBridgeFromObjCExpr *E,
                                               SGFContext C);
@@ -1515,6 +1516,11 @@ RValueEmitter::visitPackExpansionExpr(PackExpansionExpr *E,
   llvm_unreachable("not implemented for PackExpansionExpr");
 }
 
+RValue
+RValueEmitter::visitPackElementExpr(PackElementExpr *E, SGFContext C) {
+  llvm_unreachable("not implemented for PackElementExpr");
+}
+
 RValue RValueEmitter::visitArchetypeToSuperExpr(ArchetypeToSuperExpr *E,
                                                 SGFContext C) {
   ManagedValue archetype = SGF.emitRValueAsSingleValue(E->getSubExpr());

lib/Sema/CSApply.cpp

@@ -3820,13 +3820,11 @@ namespace {
     }
 
     Expr *visitPackExpansionExpr(PackExpansionExpr *expr) {
-      for (unsigned i = 0; i < expr->getNumBindings(); ++i) {
-        auto *binding = expr->getBindings()[i];
-        expr->setBinding(i, visit(binding));
-      }
+      return simplifyExprType(expr);
+    }
 
-      simplifyExprType(expr);
-      return expr;
+    Expr *visitPackElementExpr(PackElementExpr *expr) {
+      return simplifyExprType(expr);
     }
 
     Expr *visitDynamicTypeExpr(DynamicTypeExpr *expr) {
@@ -7132,10 +7130,9 @@ Expr *ExprRewriter::coerceToType(Expr *expr, Type toType,
     auto shapeType = toExpansionType->getCountType();
     auto expansionTy = PackExpansionType::get(patternType, shapeType);
 
-    return cs.cacheType(PackExpansionExpr::create(ctx, pattern,
-        expansion->getOpaqueValues(), expansion->getBindings(),
-        expansion->getEndLoc(), expansion->getGenericEnvironment(),
-        expansion->isImplicit(), expansionTy));
+    expansion->setPatternExpr(pattern);
+    expansion->setType(expansionTy);
+    return cs.cacheType(expansion);
   }
 
   case TypeKind::BuiltinTuple: