AST: Factor out duplication between PackTypeParameterCollector and PackTypeVariableCollector

Author: slavapestov

include/swift/AST/Types.h

@@ -724,8 +724,10 @@ class alignas(1 << TypeAlignInBits) TypeBase
       SmallVectorImpl<OpenedArchetypeType *> &rootOpenedArchetypes) const;
 
   /// Retrieve the set of type parameter packs that occur within this type.
-  void getTypeParameterPacks(
-      SmallVectorImpl<Type> &rootParameterPacks);
+  void getTypeParameterPacks(SmallVectorImpl<Type> &rootParameterPacks);
+
+  /// Retrieve the set of type parameter packs that occur within this type.
+  void walkPackReferences(llvm::function_ref<bool (Type)> fn);
 
   /// Replace opened archetypes with the given root with their most
   /// specific non-dependent upper bounds throughout this type.

lib/AST/ParameterPack.cpp

@@ -29,12 +29,13 @@ namespace {
 
 /// Collects all unique pack type parameters referenced from the pattern type,
 /// skipping those captured by nested pack expansion types.
-struct PackTypeParameterCollector: TypeWalker {
-  llvm::SetVector<Type> typeParams;
+struct PackReferenceCollector: TypeWalker {
+  llvm::function_ref<bool (Type)> fn;
   unsigned expansionLevel;
   SmallVector<unsigned, 2> elementLevel;
 
-  PackTypeParameterCollector() : expansionLevel(0) {
+  PackReferenceCollector(llvm::function_ref<bool (Type)> fn)
+    : fn(fn), expansionLevel(0) {
     elementLevel.push_back(0);
   }
 
@@ -72,12 +73,8 @@ struct PackTypeParameterCollector: TypeWalker {
     }
 
     if (elementLevel.back() == expansionLevel) {
-      if (auto *paramTy = t->getAs<GenericTypeParamType>()) {
-        if (paramTy->isParameterPack())
-          typeParams.insert(paramTy);
-      } else if (auto *archetypeTy = t->getAs<PackArchetypeType>()) {
-        typeParams.insert(archetypeTy->getRoot());
-      }
+      if (fn(t))
+        return Action::Stop;
     }
 
     return Action::Continue;
@@ -96,13 +93,28 @@ struct PackTypeParameterCollector: TypeWalker {
 
 }
 
+void TypeBase::walkPackReferences(
+    llvm::function_ref<bool (Type)> fn) {
+  Type(this).walk(PackReferenceCollector(fn));
+}
+
 void TypeBase::getTypeParameterPacks(
     SmallVectorImpl<Type> &rootParameterPacks) {
-  PackTypeParameterCollector collector;
-  Type(this).walk(collector);
+  llvm::SmallSetVector<Type, 2> rootParameterPackSet;
+
+  walkPackReferences([&](Type t) {
+    if (auto *paramTy = t->getAs<GenericTypeParamType>()) {
+      if (paramTy->isParameterPack())
+        rootParameterPackSet.insert(paramTy);
+    } else if (auto *archetypeTy = t->getAs<PackArchetypeType>()) {
+      rootParameterPackSet.insert(archetypeTy->getRoot());
+    }
+
+    return false;
+  });
 
-  rootParameterPacks.append(collector.typeParams.begin(),
-                            collector.typeParams.end());
+  rootParameterPacks.append(rootParameterPackSet.begin(),
+                            rootParameterPackSet.end());
 }
 
 bool GenericTypeParamType::isParameterPack() const {

lib/Sema/CSSimplify.cpp

@@ -2343,30 +2343,6 @@ ConstraintSystem::matchPackTypes(PackType *pack1, PackType *pack2,
   return getTypeMatchSuccess();
 }
 
-namespace {
-
-/// Collects all unique pack type variables referenced from the pattern type,
-/// skipping those captured by nested pack expansion types.
-///
-/// FIXME: This could probably be a common utility method somewhere.
-struct PackTypeVariableCollector: TypeWalker {
-  llvm::SetVector<TypeVariableType *> typeVars;
-
-  Action walkToTypePre(Type t) override {
-    if (t->is<PackExpansionType>())
-      return Action::SkipChildren;
-
-    if (auto *typeVar = t->getAs<TypeVariableType>()) {
-      if (typeVar->getImpl().canBindToPack())
-        typeVars.insert(typeVar);
-    }
-
-    return Action::Continue;
-  }
-};
-
-}
-
 /// Utility function used when matching a pack expansion type against a
 /// pack type.
 ///
@@ -2399,21 +2375,26 @@ static PackType *replaceTypeVariablesWithFreshPacks(ConstraintSystem &cs,
                                                     Type pattern,
                                                     PackType *pack,
                                                     ConstraintLocatorBuilder locator) {
-  SmallVector<Type, 2> elts;
-
+  llvm::SmallSetVector<TypeVariableType *, 2> typeVarSet;
   llvm::MapVector<TypeVariableType *, SmallVector<Type, 2>> typeVars;
 
-  PackTypeVariableCollector collector;
-  pattern.walk(collector);
+  pattern->walkPackReferences([&](Type t) {
+    if (auto *typeVar = t->getAs<TypeVariableType>()) {
+      if (typeVar->getImpl().canBindToPack())
+        typeVarSet.insert(typeVar);
+    }
+
+    return false;
+  });
 
-  if (collector.typeVars.empty())
+  if (typeVarSet.empty())
     return nullptr;
 
   auto *loc = cs.getConstraintLocator(locator);
 
   // For each pack type variable occurring in the pattern type, compute a
   // binding pack type comprised of fresh type variables.
-  for (auto *typeVar : collector.typeVars) {
+  for (auto *typeVar : typeVarSet) {
     auto &freshTypeVars = typeVars[typeVar];
     for (unsigned i = 0, e = pack->getNumElements(); i < e; ++i) {
       auto *packExpansionElt = pack->getElementType(i)->getAs<PackExpansionType>();
@@ -2441,6 +2422,8 @@ static PackType *replaceTypeVariablesWithFreshPacks(ConstraintSystem &cs,
     }
   }
 
+  SmallVector<Type, 2> elts;
+
   // For each element of the original pack type, instantiate the pattern type by
   // replacing each pack type variable with the corresponding element of the
   // pack type variable's binding pack.