AST: Use TypeTransform::transformPackExpansion() to transform PackExpansionType

Author: slavapestov

include/swift/AST/TypeTransform.h

@@ -586,30 +586,8 @@ case TypeKind::Id:
     }
 
     case TypeKind::PackExpansion: {
-      auto expand = cast<PackExpansionType>(base);
-
-      // Substitution completely replaces this.
-
-      Type transformedPat = doIt(expand->getPatternType(), pos);
-      if (!transformedPat)
-        return Type();
-
-      Type transformedCount = doIt(expand->getCountType(), TypePosition::Shape);
-      if (!transformedCount)
-        return Type();
-
-      if (transformedPat.getPointer() == expand->getPatternType().getPointer() &&
-          transformedCount.getPointer() == expand->getCountType().getPointer())
-        return t;
-
-      // // If we transform the count to a pack type, expand the pattern.
-      // // This is necessary because of how we piece together types in
-      // // the constraint system.
-      // if (auto countPack = transformedCount->getAs<PackType>()) {
-      //   return PackExpansionType::expand(transformedPat, countPack);
-      // }
-
-      return PackExpansionType::get(transformedPat, transformedCount);
+      auto *expand = cast<PackExpansionType>(base);
+      return asDerived().transformPackExpansion(expand, pos);
     }
 
     case TypeKind::PackElement: {
@@ -1005,8 +983,26 @@ case TypeKind::Id:
   CanType transformSILField(CanType fieldTy, TypePosition pos) {
     return doIt(fieldTy, pos)->getCanonicalType();
   }
+
+  Type transformPackExpansion(PackExpansionType *expand, TypePosition pos) {
+    // Substitution completely replaces this.
+
+    Type transformedPat = doIt(expand->getPatternType(), pos);
+    if (!transformedPat)
+      return Type();
+
+    Type transformedCount = doIt(expand->getCountType(), TypePosition::Shape);
+    if (!transformedCount)
+      return Type();
+
+    if (transformedPat.getPointer() == expand->getPatternType().getPointer() &&
+        transformedCount.getPointer() == expand->getCountType().getPointer())
+      return expand;
+
+    return PackExpansionType::get(transformedPat, transformedCount);
+  }
 };
 
 }
 
-#endif
+#endif

lib/AST/TypeSubstitution.cpp

@@ -488,7 +488,9 @@ class TypeSubstituter : public TypeTransform<TypeSubstituter> {
   TypeSubstituter(unsigned level, InFlightSubstitution &IFS)
     : level(level), IFS(IFS) {}
 
-  std::optional<Type> transform(TypeBase *type, TypePosition position);
+  std::optional<Type> transform(TypeBase *type, TypePosition pos);
+
+  Type transformPackExpansion(PackExpansionType *expand, TypePosition pos);
 
   SubstitutionMap transformSubstitutionMap(SubstitutionMap subs);
 
@@ -506,13 +508,6 @@ TypeSubstituter::transform(TypeBase *type, TypePosition position) {
          "should not be doing AST type-substitution on a lowered SIL type;"
          "use SILType::subst");
 
-  if (auto packExpansionTy = dyn_cast<PackExpansionType>(type)) {
-    auto eltTys = IFS.expandPackExpansionType(packExpansionTy);
-    if (eltTys.size() == 1)
-      return eltTys[0];
-    return Type(PackType::get(packExpansionTy->getASTContext(), eltTys));
-  }
-
   auto oldLevel = level;
   SWIFT_DEFER { level = oldLevel; };
 
@@ -589,6 +584,14 @@ TypeSubstituter::transform(TypeBase *type, TypePosition position) {
                               level);
 }
 
+Type TypeSubstituter::transformPackExpansion(PackExpansionType *expand,
+                                             TypePosition pos) {
+  auto eltTys = IFS.expandPackExpansionType(expand);
+  if (eltTys.size() == 1)
+    return eltTys[0];
+  return Type(PackType::get(expand->getASTContext(), eltTys));
+}
+
 SubstitutionMap TypeSubstituter::transformSubstitutionMap(SubstitutionMap subs) {
   // FIXME: Take level into account? Move level down into IFS?
   return subs.subst(IFS);