Merge pull request #61610 from slavapestov/fix-pack-expansion-matching

Fix pack expansion matching

Author: slavapestov

include/swift/Sema/ConstraintSystem.h

@@ -5197,6 +5197,12 @@ class ConstraintSystem {
                  ConstraintKind kind, TypeMatchOptions flags,
                  ConstraintLocatorBuilder locator);
 
+  TypeMatchResult
+  matchPackExpansionTypes(PackExpansionType *expansion1,
+                          PackExpansionType *expansion2,
+                          ConstraintKind kind, TypeMatchOptions flags,
+                          ConstraintLocatorBuilder locator);
+
   /// Subroutine of \c matchTypes(), which matches up two tuple types.
   ///
   /// \returns the result of performing the tuple-to-tuple conversion.

lib/AST/PackExpansionMatcher.cpp

@@ -62,15 +62,15 @@ bool TuplePackMatcher::match() {
     // A pack expansion type on the left hand side absorbs all elements
     // from the right hand side up to the next mismatched label.
     auto lhsElt = lhsElts.front();
-    if (lhsElt.getType()->is<PackExpansionType>()) {
+    if (auto *lhsExpansionType = lhsElt.getType()->getAs<PackExpansionType>()) {
       lhsElts = lhsElts.slice(1);
 
       assert(lhsElts.empty() || lhsElts.front().hasName() &&
              "Tuple element with pack expansion type cannot be followed "
              "by an unlabeled element");
 
       auto *rhs = gatherTupleElements(rhsElts, lhsElt.getName(), ctx);
-      pairs.emplace_back(lhsElt.getType(), rhs, idx++);
+      pairs.emplace_back(lhsExpansionType->getPatternType(), rhs, idx++);
       continue;
     }
 
@@ -82,15 +82,15 @@ bool TuplePackMatcher::match() {
     // A pack expansion type on the right hand side absorbs all elements
     // from the left hand side up to the next mismatched label.
     auto rhsElt = rhsElts.front();
-    if (rhsElt.getType()->is<PackExpansionType>()) {
+    if (auto *rhsExpansionType = rhsElt.getType()->getAs<PackExpansionType>()) {
       rhsElts = rhsElts.slice(1);
 
       assert(rhsElts.empty() || rhsElts.front().hasName() &&
              "Tuple element with pack expansion type cannot be followed "
              "by an unlabeled element");
 
       auto *lhs = gatherTupleElements(lhsElts, rhsElt.getName(), ctx);
-      pairs.emplace_back(lhs, rhsElt.getType(), idx++);
+      pairs.emplace_back(lhs, rhsExpansionType->getPatternType(), idx++);
       continue;
     }
 
@@ -169,34 +169,38 @@ bool ParamPackMatcher::match() {
 
   // If the left hand side is a single pack expansion type, bind it
   // to what remains of the right hand side.
-  if (lhsParams.size() == 1 &&
-      lhsParams[0].getPlainType()->is<PackExpansionType>()) {
-    SmallVector<Type, 2> rhsTypes;
-    for (auto rhsParam : rhsParams) {
-      // FIXME: Check rhs flags
-      rhsTypes.push_back(rhsParam.getPlainType());
-    }
-    auto rhs = PackType::get(ctx, rhsTypes);
+  if (lhsParams.size() == 1) {
+    auto lhsType = lhsParams[0].getPlainType();
+    if (auto *lhsExpansionType = lhsType->getAs<PackExpansionType>()) {
+      SmallVector<Type, 2> rhsTypes;
+      for (auto rhsParam : rhsParams) {
+        // FIXME: Check rhs flags
+        rhsTypes.push_back(rhsParam.getPlainType());
+      }
+      auto rhs = PackType::get(ctx, rhsTypes);
 
-    // FIXME: Check lhs flags
-    pairs.emplace_back(lhsParams[0].getPlainType(), rhs, prefixLength);
-    return false;
+      // FIXME: Check lhs flags
+      pairs.emplace_back(lhsExpansionType->getPatternType(), rhs, prefixLength);
+      return false;
+    }
   }
 
   // If the right hand side is a single pack expansion type, bind it
   // to what remains of the left hand side.
-  if (rhsParams.size() == 1 &&
-      rhsParams[0].getPlainType()->is<PackExpansionType>()) {
-    SmallVector<Type, 2> lhsTypes;
-    for (auto lhsParam : lhsParams) {
-      // FIXME: Check lhs flags
-      lhsTypes.push_back(lhsParam.getPlainType());
-    }
-    auto lhs = PackType::get(ctx, lhsTypes);
+  if (rhsParams.size() == 1) {
+    auto rhsType = rhsParams[0].getPlainType();
+    if (auto *rhsExpansionType = rhsType->getAs<PackExpansionType>()) {
+      SmallVector<Type, 2> lhsTypes;
+      for (auto lhsParam : lhsParams) {
+        // FIXME: Check lhs flags
+        lhsTypes.push_back(lhsParam.getPlainType());
+      }
+      auto lhs = PackType::get(ctx, lhsTypes);
 
-    // FIXME: Check rhs flags
-    pairs.emplace_back(lhs, rhsParams[0].getPlainType(), prefixLength);
-    return false;
+      // FIXME: Check rhs flags
+      pairs.emplace_back(lhs, rhsParams[0].getPlainType(), prefixLength);
+      return false;
+    }
   }
 
   // Otherwise, all remaining possibilities are invalid:

lib/Sema/CSSimplify.cpp

@@ -2425,6 +2425,23 @@ ConstraintSystem::matchPackTypes(PackType *pack1, PackType *pack2,
   return getTypeMatchSuccess();
 }
 
+ConstraintSystem::TypeMatchResult
+ConstraintSystem::matchPackExpansionTypes(PackExpansionType *expansion1,
+                                          PackExpansionType *expansion2,
+                                          ConstraintKind kind, TypeMatchOptions flags,
+                                          ConstraintLocatorBuilder locator) {
+  // FIXME: Should we downgrade kind to Bind or something here?
+  auto result = matchTypes(expansion1->getCountType(),
+                           expansion2->getCountType(),
+                           kind, flags, locator);
+  if (result.isFailure())
+    return result;
+
+  return matchTypes(expansion1->getPatternType(),
+                    expansion2->getPatternType(),
+                    kind, flags, locator);
+}
+
 /// Check where a representation is a subtype of another.
 ///
 /// The subtype relationship is defined as:
@@ -6636,10 +6653,13 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
                             kind, subflags, packLoc);
     }
     case TypeKind::PackExpansion: {
-      // FIXME: we need to match the count types as well
-      return matchTypes(cast<PackExpansionType>(desugar1)->getPatternType(),
-                        cast<PackExpansionType>(desugar2)->getPatternType(),
-                        kind, subflags, locator);
+      // FIXME: Need a new locator element
+
+      auto expansion1 = cast<PackExpansionType>(desugar1);
+      auto expansion2 = cast<PackExpansionType>(desugar2);
+
+      return matchPackExpansionTypes(expansion1, expansion2, kind, subflags,
+                                     locator);
     }
     }
   }
@@ -7034,20 +7054,6 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
     }
   }
 
-  if (isa<PackExpansionType>(desugar1) && isa<PackType>(desugar2)) {
-    auto *packExpansionType = cast<PackExpansionType>(desugar1);
-    auto *packType = cast<PackType>(desugar2);
-
-    if (packExpansionType->getPatternType()->is<TypeVariableType>())
-      return matchTypes(packExpansionType->getPatternType(), packType, kind, subflags, locator);
-  } else if (isa<PackType>(desugar1) && isa<PackExpansionType>(desugar2)) {
-    auto *packType = cast<PackType>(desugar1);
-    auto *packExpansionType = cast<PackExpansionType>(desugar2);
-
-    if (packExpansionType->getPatternType()->is<TypeVariableType>())
-      return matchTypes(packType, packExpansionType->getPatternType(), kind, subflags, locator);
-  }
-
   // Attempt fixes iff it's allowed, both types are concrete and
   // we are not in the middle of attempting one already.
   if (shouldAttemptFixes() && !flags.contains(TMF_ApplyingFix)) {