[AST] Introduce PackReferenceTypeRepr for pack references spelled with

a contextual 'each' keyword.

Author: hborla

include/swift/AST/DiagnosticsSema.def

@@ -5150,6 +5150,12 @@ ERROR(expansion_not_allowed,none,
       "function result, tuple element or generic argument list", (Type))
 ERROR(expansion_not_variadic,none,
       "variadic expansion %0 must contain at least one variadic generic parameter", (Type))
+ERROR(pack_reference_outside_expansion,none,
+      "pack reference %0 can only appear in pack expansion or generic requirement",
+      (Type))
+ERROR(each_non_pack,none,
+      "'each' cannot be applied to non-pack type %0",
+      (Type))
 ERROR(tuple_duplicate_label,none,
       "cannot create a tuple with a duplicate element label", ())
 ERROR(multiple_ellipsis_in_tuple,none,

include/swift/AST/TypeRepr.h

@@ -743,6 +743,41 @@ class PackExpansionTypeRepr final : public TypeRepr {
   friend class TypeRepr;
 };
 
+/// A pack reference spelled with the \c each keyword.
+///
+/// Pack references can only appear inside pack expansions and in
+/// generic requirements.
+///
+/// \code
+/// struct Generic<T...> {
+///   func f(value: (each T)...) where each T: P {}
+/// }
+/// \endcode
+class PackReferenceTypeRepr: public TypeRepr {
+  TypeRepr *PackType;
+  SourceLoc EachLoc;
+
+public:
+  PackReferenceTypeRepr(SourceLoc eachLoc, TypeRepr *packType)
+    : TypeRepr(TypeReprKind::PackReference), PackType(packType),
+      EachLoc(eachLoc) {}
+
+  TypeRepr *getPackType() const { return PackType; }
+  SourceLoc getEachLoc() const { return EachLoc; }
+
+  static bool classof(const TypeRepr *T) {
+    return T->getKind() == TypeReprKind::PackReference;
+  }
+  static bool classof(const PackReferenceTypeRepr *T) { return true; }
+
+private:
+  SourceLoc getStartLocImpl() const { return EachLoc; }
+  SourceLoc getEndLocImpl() const { return PackType->getEndLoc(); }
+  SourceLoc getLocImpl() const { return EachLoc; }
+  void printImpl(ASTPrinter &Printer, const PrintOptions &Opts) const;
+  friend class TypeRepr;
+};
+
 /// A tuple type.
 /// \code
 ///   (Foo, Bar)
@@ -1336,6 +1371,7 @@ inline bool TypeRepr::isSimple() const {
   case TypeReprKind::OpaqueReturn:
   case TypeReprKind::NamedOpaqueReturn:
   case TypeReprKind::Existential:
+  case TypeReprKind::PackReference:
     return false;
   case TypeReprKind::SimpleIdent:
   case TypeReprKind::GenericIdent:

include/swift/AST/TypeReprNodes.def

@@ -64,6 +64,7 @@ TYPEREPR(Protocol, TypeRepr)
 TYPEREPR(OpaqueReturn, TypeRepr)
 TYPEREPR(NamedOpaqueReturn, TypeRepr)
 TYPEREPR(Existential, TypeRepr)
+TYPEREPR(PackReference, TypeRepr)
 TYPEREPR(Placeholder, TypeRepr)
 ABSTRACT_TYPEREPR(Specifier, TypeRepr)
   SPECIFIER_TYPEREPR(InOut, SpecifierTypeRepr)

lib/AST/ASTDumper.cpp

@@ -3141,6 +3141,12 @@ class PrintTypeRepr : public TypeReprVisitor<PrintTypeRepr> {
     PrintWithColorRAII(OS, ParenthesisColor) << ')';
   }
 
+  void visitPackReferenceTypeRepr(PackReferenceTypeRepr *T) {
+    printCommon("pack_reference");
+    printRec(T->getPackType());
+    PrintWithColorRAII(OS, ParenthesisColor) << ')';
+  }
+
   void visitTupleTypeRepr(TupleTypeRepr *T) {
     printCommon("type_tuple");
 

lib/AST/ASTWalker.cpp

@@ -1995,6 +1995,10 @@ bool Traversal::visitPackExpansionTypeRepr(PackExpansionTypeRepr *T) {
   return doIt(T->getPatternType());
 }
 
+bool Traversal::visitPackReferenceTypeRepr(PackReferenceTypeRepr *T) {
+  return doIt(T->getPackType());
+}
+
 bool Traversal::visitTupleTypeRepr(TupleTypeRepr *T) {
   for (auto &elem : T->getElements()) {
     if (doIt(elem.Type))

lib/AST/NameLookup.cpp

@@ -2516,6 +2516,13 @@ directReferencesForTypeRepr(Evaluator &evaluator,
                                        allowUsableFromInline);
   }
 
+  case TypeReprKind::PackReference: {
+    auto packReferenceRepr = cast<PackReferenceTypeRepr>(typeRepr);
+    return directReferencesForTypeRepr(evaluator, ctx,
+                                       packReferenceRepr->getPackType(), dc,
+                                       allowUsableFromInline);
+  }
+
   case TypeReprKind::Error:
   case TypeReprKind::Function:
   case TypeReprKind::InOut:

lib/AST/TypeRepr.cpp

@@ -413,6 +413,12 @@ void PackExpansionTypeRepr::printImpl(ASTPrinter &Printer,
   Printer << "...";
 }
 
+void PackReferenceTypeRepr::printImpl(ASTPrinter &Printer,
+                                      const PrintOptions &Opts) const {
+  Printer.printKeyword("each", Opts, /*Suffix=*/" ");
+  printTypeRepr(PackType, Printer, Opts);
+}
+
 void TupleTypeRepr::printImpl(ASTPrinter &Printer,
                               const PrintOptions &Opts) const {
   Printer.callPrintStructurePre(PrintStructureKind::TupleType);

lib/Sema/TypeCheckGeneric.cpp

@@ -966,6 +966,7 @@ RequirementRequest::evaluate(Evaluator &evaluator,
     context = TypeResolverContext::GenericRequirement;
   }
   auto options = TypeResolutionOptions(context);
+  options |= TypeResolutionFlags::AllowPackReferences;
   if (owner.dc->isInSpecializeExtensionContext())
     options |= TypeResolutionFlags::AllowUsableFromInline;
   Optional<TypeResolution> resolution;

lib/Sema/TypeCheckType.cpp

@@ -2088,6 +2088,8 @@ namespace {
                                   TypeResolutionOptions options);
     NeverNullType resolvePackExpansionType(PackExpansionTypeRepr *repr,
                                            TypeResolutionOptions options);
+    NeverNullType resolvePackReference(PackReferenceTypeRepr *repr,
+                                       TypeResolutionOptions options);
     NeverNullType resolveTupleType(TupleTypeRepr *repr,
                                    TypeResolutionOptions options);
     NeverNullType resolveCompositionType(CompositionTypeRepr *repr,
@@ -2277,6 +2279,9 @@ NeverNullType TypeResolver::resolveType(TypeRepr *repr,
   case TypeReprKind::PackExpansion:
     return resolvePackExpansionType(cast<PackExpansionTypeRepr>(repr), options);
 
+  case TypeReprKind::PackReference:
+    return resolvePackReference(cast<PackReferenceTypeRepr>(repr), options);
+
   case TypeReprKind::Tuple:
     return resolveTupleType(cast<TupleTypeRepr>(repr), options);
 
@@ -4182,6 +4187,7 @@ std::pair<Type, Type>
 TypeResolver::maybeResolvePackExpansionType(PackExpansionTypeRepr *repr,
                                             TypeResolutionOptions options) {
   auto elementOptions = options;
+  elementOptions |= TypeResolutionFlags::AllowPackReferences;
   auto patternTy = resolveType(repr->getPatternType(), elementOptions);
   if (patternTy->hasError())
     return std::make_pair(ErrorType::get(getASTContext()), Type());
@@ -4243,6 +4249,31 @@ NeverNullType TypeResolver::resolvePackExpansionType(PackExpansionTypeRepr *repr
   return PackExpansionType::get(pair.first, pair.second);
 }
 
+NeverNullType TypeResolver::resolvePackReference(PackReferenceTypeRepr *repr,
+                                                 TypeResolutionOptions options) {
+  auto &ctx = getASTContext();
+  auto packReference = resolveType(repr->getPackType(), options);
+
+  // If we already failed, don't diagnose again.
+  if (packReference->hasError())
+    return ErrorType::get(ctx);
+
+  if (!packReference->isParameterPack()) {
+    ctx.Diags.diagnose(repr->getLoc(), diag::each_non_pack,
+                       packReference);
+    return packReference;
+  }
+
+  if (!options.contains(TypeResolutionFlags::AllowPackReferences)) {
+    ctx.Diags.diagnose(repr->getLoc(),
+                       diag::pack_reference_outside_expansion,
+                       packReference);
+    return packReference;
+  }
+
+  return packReference;
+}
+
 NeverNullType TypeResolver::resolveTupleType(TupleTypeRepr *repr,
                                              TypeResolutionOptions options) {
   auto &ctx = getASTContext();
@@ -4739,6 +4770,7 @@ class ExistentialTypeVisitor
     case TypeReprKind::Placeholder:
     case TypeReprKind::CompileTimeConst:
     case TypeReprKind::PackExpansion:
+    case TypeReprKind::PackReference:
       return false;
     }
   }

lib/Sema/TypeCheckType.h

@@ -70,6 +70,12 @@ enum class TypeResolutionFlags : uint16_t {
 
   /// We are in a `@preconcurrency` declaration.
   Preconcurrency = 1 << 10,
+
+  /// Whether references to type parameter packs are allowed.
+  ///
+  /// Pack references are only allowed inside pack expansions
+  /// and in generic requirements.
+  AllowPackReferences = 1 << 11,
 };
 
 /// Type resolution contexts that require special handling.