AST: Add a count type field to PackExpansionType

Even if we can't spell them in source, we want to model expansions where
the pattern does not depend on any pack type parameters, eg

    func f<C...: Collection>(_ c: C...) {
      let x = (c.count...)
    }

Here, the type of 'x' is notionally 'Int * C.count'.

Author: slavapestov

Diff for: include/swift/AST/Types.h

@@ -6378,7 +6378,7 @@ class PackType final : public TypeBase, public llvm::FoldingSetNode,
   }
 
 public:
-  void Profile(llvm::FoldingSetNodeID &ID) {
+  void Profile(llvm::FoldingSetNodeID &ID) const {
     Profile(ID, getElementTypes());
   }
   static void Profile(llvm::FoldingSetNodeID &ID, ArrayRef<Type> Elements);
@@ -6427,41 +6427,52 @@ class PackExpansionType : public TypeBase, public llvm::FoldingSetNode {
   friend class ASTContext;
 
   Type patternType;
+  Type countType;
 
 public:
   /// Create a pack expansion type from the given pattern type.
   ///
-  /// It is not required that the pattern type actually contain a reference to
-  /// a variadic generic parameter.
-  static PackExpansionType *get(Type pattern);
+  /// It is not required that \p pattern actually contain a reference to
+  /// a variadic generic parameter, but any variadic generic parameters
+  /// appearing in the pattern type must have the same count as \p countType.
+  ///
+  /// As for \p countType itself, it must be a type sequence generic parameter
+  /// type, or a sequence archetype type.
+  static PackExpansionType *get(Type pattern, Type countType);
 
 public:
   /// Retrieves the pattern type of this pack expansion.
   Type getPatternType() const { return patternType; }
 
+  /// Retrieves the count type of this pack expansion.
+  Type getCountType() const { return countType; }
+
 public:
   void Profile(llvm::FoldingSetNodeID &ID) {
-    Profile(ID, getPatternType());
+    Profile(ID, getPatternType(), getCountType());
   }
 
-  static void Profile(llvm::FoldingSetNodeID &ID, Type patternType);
+  static void Profile(llvm::FoldingSetNodeID &ID,
+                      Type patternType, Type countType);
 
   // Implement isa/cast/dyncast/etc.
   static bool classof(const TypeBase *T) {
     return T->getKind() == TypeKind::PackExpansion;
   }
 
 private:
-  PackExpansionType(Type patternType, const ASTContext *CanCtx)
-    : TypeBase(TypeKind::PackExpansion, CanCtx,
-               patternType->getRecursiveProperties()), patternType(patternType) {
-      assert(patternType);
-    }
+  PackExpansionType(Type patternType, Type countType,
+                    RecursiveTypeProperties properties,
+                    const ASTContext *ctx);
 };
 BEGIN_CAN_TYPE_WRAPPER(PackExpansionType, Type)
   CanType getPatternType() const {
     return CanType(getPointer()->getPatternType());
   }
+
+  CanType getCountType() const {
+    return CanType(getPointer()->getCountType());
+  }
 END_CAN_TYPE_WRAPPER(PackExpansionType, Type)
 
 /// getASTContext - Return the ASTContext that this type belongs to.

Diff for: lib/AST/ASTContext.cpp

@@ -2897,33 +2897,48 @@ TupleTypeElt::TupleTypeElt(Type ty, Identifier name)
   assert(!ty->is<InOutType>() && "Cannot have InOutType in a tuple");
 }
 
-PackExpansionType *PackExpansionType::get(Type patternTy) {
-  assert(patternTy && "Missing pattern type in expansion");
+PackExpansionType::PackExpansionType(Type patternType, Type countType,
+                                     RecursiveTypeProperties properties,
+                                     const ASTContext *canCtx)
+  : TypeBase(TypeKind::PackExpansion, canCtx, properties),
+    patternType(patternType), countType(countType) {
+  assert(countType->is<TypeVariableType>() ||
+         countType->is<SequenceArchetypeType>() ||
+         countType->castTo<GenericTypeParamType>()->isTypeSequence());
+}
+
+PackExpansionType *PackExpansionType::get(Type patternType, Type countType) {
+  auto properties = patternType->getRecursiveProperties();
+  properties |= countType->getRecursiveProperties();
 
-  auto properties = patternTy->getRecursiveProperties();
   auto arena = getArena(properties);
 
-  auto &context = patternTy->getASTContext();
+  auto &context = patternType->getASTContext();
   llvm::FoldingSetNodeID id;
-  PackExpansionType::Profile(id, patternTy);
+  PackExpansionType::Profile(id, patternType, countType);
 
   void *insertPos;
   if (PackExpansionType *expType =
-          context.getImpl()
-              .getArena(arena)
-              .PackExpansionTypes.FindNodeOrInsertPos(id, insertPos))
+        context.getImpl().getArena(arena)
+          .PackExpansionTypes.FindNodeOrInsertPos(id, insertPos))
     return expType;
 
-  const ASTContext *canCtx = patternTy->isCanonical() ? &context : nullptr;
-  PackExpansionType *expansionTy = new (context, AllocationArena::Permanent)
-      PackExpansionType(patternTy, canCtx);
-  context.getImpl().getArena(arena).PackExpansionTypes.InsertNode(expansionTy,
+  const ASTContext *canCtx =
+      (patternType->isCanonical() && countType->isCanonical())
+      ? &context : nullptr;
+  PackExpansionType *expansionType =
+      new (context, arena) PackExpansionType(patternType, countType, properties,
+                                             canCtx);
+  context.getImpl().getArena(arena).PackExpansionTypes.InsertNode(expansionType,
                                                                   insertPos);
-  return expansionTy;
+  return expansionType;
 }
 
-void PackExpansionType::Profile(llvm::FoldingSetNodeID &ID, Type patternType) {
+void PackExpansionType::Profile(llvm::FoldingSetNodeID &ID,
+                                Type patternType,
+                                Type countType) {
   ID.AddPointer(patternType.getPointer());
+  ID.AddPointer(countType.getPointer());
 }
 
 PackType *PackType::getEmpty(const ASTContext &C) {

Diff for: lib/AST/ASTDumper.cpp

@@ -3686,6 +3686,7 @@ namespace {
     void visitPackExpansionType(PackExpansionType *T, StringRef label) {
       printCommon(label, "pack_expansion_type");
       printField("pattern", T->getPatternType());
+      printField("count", T->getCountType());
       PrintWithColorRAII(OS, ParenthesisColor) << ')';
     }
 

Diff for: lib/AST/Type.cpp

@@ -1593,8 +1593,9 @@ CanType TypeBase::computeCanonicalType() {
 
   case TypeKind::PackExpansion: {
     auto *expansion = cast<PackExpansionType>(this);
-    auto pattern = expansion->getPatternType()->getCanonicalType();
-    Result = PackExpansionType::get(pattern);
+    auto patternType = expansion->getPatternType()->getCanonicalType();
+    auto countType = expansion->getCountType()->getCanonicalType();
+    Result = PackExpansionType::get(patternType, countType);
     break;
   }
 
@@ -5535,12 +5536,9 @@ case TypeKind::Id:
           return remap;
         }
 
-        if (input->is<TypeVariableType>()) {
-          if (auto *PT = (*remap)->getAs<PackType>()) {
-            maxArity = std::max(maxArity, PT->getNumElements());
-            cache.insert({input, PT});
-          }
-        } else if (input->isTypeSequenceParameter()) {
+        if (input->is<TypeVariableType>() ||
+            input->isTypeSequenceParameter() ||
+            input->is<SequenceArchetypeType>()) {
           if (auto *PT = (*remap)->getAs<PackType>()) {
             maxArity = std::max(maxArity, PT->getNumElements());
             cache.insert({input, PT});
@@ -5563,7 +5561,13 @@ case TypeKind::Id:
     if (!transformedPat)
       return Type();
 
-    if (transformedPat.getPointer() == expand->getPatternType().getPointer())
+    Type transformedCount =
+        expand->getCountType().transformWithPosition(pos, gather);
+    if (!transformedCount)
+      return Type();
+
+    if (transformedPat.getPointer() == expand->getPatternType().getPointer() &&
+        transformedCount.getPointer() == expand->getCountType().getPointer())
       return *this;
 
     llvm::DenseMap<TypeBase *, PackType *> expansions;
@@ -5573,7 +5577,7 @@ case TypeKind::Id:
       // If we didn't find any expansions, either the caller wasn't interested
       // in expanding this pack, or something has gone wrong. Leave off the
       // expansion and return the transformed type.
-      return PackExpansionType::get(transformedPat);
+      return PackExpansionType::get(transformedPat, transformedCount);
     }
 
     SmallVector<Type, 8> elts;

Diff for: lib/Sema/TypeCheckDecl.cpp

@@ -2227,10 +2227,22 @@ static Type validateParameterType(ParamDecl *decl) {
   }
 
   if (decl->isVariadic()) {
+    // Find the first type sequence parameter and use that as the count type.
+    Type countTy;
+    (void) Ty.findIf([&](Type t) -> bool {
+      if (auto *paramTy = t->getAs<GenericTypeParamType>()) {
+        if (paramTy->isTypeSequence()) {
+          countTy = paramTy;
+          return true;
+        }
+      }
+
+      return false;
+    });
     // Handle the monovariadic/polyvariadic interface type split.
-    if (Ty->hasTypeSequence()) {
+    if (countTy) {
       // Polyvariadic types (T...) for <T...> resolve to pack expansions.
-      Ty = PackExpansionType::get(Ty);
+      Ty = PackExpansionType::get(Ty, countTy);
     } else {
       // Monovariadic types (T...) for <T> resolve to [T].
       Ty = VariadicSequenceType::get(Ty);

Diff for: lib/Sema/TypeCheckType.cpp

@@ -4018,13 +4018,28 @@ NeverNullType TypeResolver::resolveTupleType(TupleTypeRepr *repr,
       if (patternTy->hasError())
         complained = true;
 
+      // Find the first type sequence parameter and use that as the count type.
+      Type countTy;
+      (void) patternTy.get().findIf([&](Type t) -> bool {
+        if (auto *paramTy = t->getAs<GenericTypeParamType>()) {
+          if (paramTy->isTypeSequence()) {
+            countTy = paramTy;
+            return true;
+          }
+        }
+
+        return false;
+      });
+
       // If there's no reference to a variadic generic parameter, complain
       // - the pack won't actually expand to anything meaningful.
-      if (!patternTy->hasTypeSequence())
+      if (!countTy) {
         diagnose(repr->getLoc(), diag::expansion_not_variadic, patternTy)
           .highlight(repr->getParens());
+        return ErrorType::get(getASTContext());
+      }
 
-      return PackExpansionType::get(patternTy);
+      return PackExpansionType::get(patternTy, countTy);
     } else {
       // Variadic tuples are not permitted.
       //

Diff for: test/Constraints/type_sequence.swift

@@ -65,7 +65,7 @@ func bindPrefixAndSuffix() {
 
 func invalidPacks() {
   func monovariadic1() -> (String...) {} // expected-error {{cannot create expansion with non-variadic type 'String'}}
-  func monovariadic2<T>() -> (T...) {} // expected-error 2 {{cannot create expansion with non-variadic type 'T'}}
+  func monovariadic2<T>() -> (T...) {} // expected-error {{cannot create expansion with non-variadic type 'T'}}
   func monovariadic3<T, U>() -> (T, U...) {} // expected-error {{cannot create a variadic tuple}}
 }