AST: Introduce TypeBase::getTypeSequenceParameters()

Author: slavapestov

include/swift/AST/Types.h

@@ -654,6 +654,11 @@ class alignas(1 << TypeAlignInBits) TypeBase
   void getRootOpenedExistentials(
       SmallVectorImpl<OpenedArchetypeType *> &rootOpenedArchetypes) const;
 
+  /// Retrieve the set of type sequence generic parameters that occur
+  /// within this type.
+  void getTypeSequenceParameters(
+      SmallVectorImpl<Type> &rootTypeSequenceParams) const;
+
   /// Replace opened archetypes with the given root with their most
   /// specific non-dependent upper bounds throughout this type.
   Type typeEraseOpenedArchetypesWithRoot(const OpenedArchetypeType *root,

lib/AST/Type.cpp

@@ -580,6 +580,28 @@ Type TypeBase::typeEraseOpenedArchetypesWithRoot(
   return transformFn(type);
 }
 
+void TypeBase::getTypeSequenceParameters(
+    SmallVectorImpl<Type> &rootTypeSequenceParams) const {
+  llvm::SmallDenseSet<CanType, 2> visited;
+
+  auto recordType = [&](Type t) {
+    if (visited.insert(t->getCanonicalType()).second)
+      rootTypeSequenceParams.push_back(t);
+  };
+
+  Type(const_cast<TypeBase *>(this)).visit([&](Type t) {
+    if (auto *paramTy = t->getAs<GenericTypeParamType>()) {
+      if (paramTy->isTypeSequence()) {
+        recordType(paramTy);
+      }
+    } else if (auto *archetypeTy = t->getAs<SequenceArchetypeType>()) {
+      if (archetypeTy->isRoot()) {
+        recordType(t);
+      }
+    }
+  });
+}
+
 Type TypeBase::addCurriedSelfType(const DeclContext *dc) {
   if (!dc->isTypeContext())
     return this;

lib/Sema/TypeCheckDecl.cpp

@@ -2228,21 +2228,13 @@ 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;
-        }
-      }
+    SmallVector<Type, 2> rootTypeSequenceParams;
+    Ty->getTypeSequenceParameters(rootTypeSequenceParams);
 
-      return false;
-    });
     // Handle the monovariadic/polyvariadic interface type split.
-    if (countTy) {
+    if (!rootTypeSequenceParams.empty()) {
       // Polyvariadic types (T...) for <T...> resolve to pack expansions.
-      Ty = PackExpansionType::get(Ty, countTy);
+      Ty = PackExpansionType::get(Ty, rootTypeSequenceParams[0]);
     } else {
       // Monovariadic types (T...) for <T> resolve to [T].
       Ty = VariadicSequenceType::get(Ty);

lib/Sema/TypeCheckType.cpp

@@ -4031,27 +4031,18 @@ NeverNullType TypeResolver::resolveTupleType(TupleTypeRepr *repr,
         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;
-      });
+      SmallVector<Type, 1> rootTypeSequenceParams;
+      patternTy->getTypeSequenceParameters(rootTypeSequenceParams);
 
       // If there's no reference to a variadic generic parameter, complain
       // - the pack won't actually expand to anything meaningful.
-      if (!countTy) {
+      if (rootTypeSequenceParams.empty()) {
         diagnose(repr->getLoc(), diag::expansion_not_variadic, patternTy)
           .highlight(repr->getParens());
         return ErrorType::get(getASTContext());
       }
 
-      return PackExpansionType::get(patternTy, countTy);
+      return PackExpansionType::get(patternTy, rootTypeSequenceParams[0]);
     } else {
       // Variadic tuples are not permitted.
       //