[ConstraintSystem] Record opened pack expansion type in the locator element

Author: xedin

include/swift/Sema/ConstraintLocator.h

@@ -1177,6 +1177,22 @@ class LocatorPathElt::AnyPatternDecl final
   }
 };
 
+class LocatorPathElt::PackExpansionType final
+    : public StoredPointerElement<swift::PackExpansionType> {
+public:
+  PackExpansionType(swift::PackExpansionType *openedPackExpansionTy)
+      : StoredPointerElement(PathElementKind::PackExpansionType,
+                             openedPackExpansionTy) {
+    assert(openedPackExpansionTy);
+  }
+
+  swift::PackExpansionType *getOpenedType() const { return getStoredPointer(); }
+
+  static bool classof(const LocatorPathElt *elt) {
+    return elt->getKind() == PathElementKind::PackExpansionType;
+  }
+};
+
 namespace details {
   template <typename CustomPathElement>
   class PathElement {

include/swift/Sema/ConstraintLocatorPathElts.def

@@ -202,8 +202,8 @@ CUSTOM_LOCATOR_PATH_ELT(PackElement)
 /// The shape of a parameter pack.
 SIMPLE_LOCATOR_PATH_ELT(PackShape)
 
-/// The type of pack expansion
-SIMPLE_LOCATOR_PATH_ELT(PackExpansionType)
+/// The type of an "opened" pack expansion
+CUSTOM_LOCATOR_PATH_ELT(PackExpansionType)
 
 /// The pattern  of a pack expansion.
 SIMPLE_LOCATOR_PATH_ELT(PackExpansionPattern)

lib/Sema/CSSimplify.cpp

@@ -11313,24 +11313,16 @@ bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
 
 bool ConstraintSystem::resolvePackExpansion(TypeVariableType *typeVar,
                                             ConstraintLocatorBuilder locator) {
-  // TODO: Pack expansion type variable can carry opened type.
-  auto expansion =
-      llvm::find_if(OpenedPackExpansionTypes, [&](const auto &expansion) {
-        return expansion.second->isEqual(typeVar);
-      });
-
-  assert(expansion != OpenedPackExpansionTypes.end());
-
-  assignFixedType(typeVar, expansion->first, getConstraintLocator(locator));
-
-  if (isDebugMode()) {
-    PrintOptions PO;
-    PO.PrintTypesForDebugging = true;
-    llvm::errs().indent(solverState->getCurrentIndent())
-        << "(pack expansion variable " << typeVar->getString(PO)
-        << " is bound to '" << expansion->first->getString(PO) << "')\n";
+  Type openedExpansionType;
+  if (auto last = locator.last()) {
+    auto expansionElt = last->castTo<LocatorPathElt::PackExpansionType>();
+    openedExpansionType = expansionElt.getOpenedType();
   }
 
+  if (!openedExpansionType)
+    return false;
+
+  assignFixedType(typeVar, openedExpansionType, getConstraintLocator(locator));
   return true;
 }
 

lib/Sema/ConstraintLocator.cpp

@@ -504,7 +504,9 @@ void LocatorPathElt::dump(raw_ostream &out) const {
     break;
 
   case ConstraintLocator::PackExpansionType:
-    out << "pack expansion type";
+    auto expansionElt = elt.castTo<LocatorPathElt::PackExpansionType>();
+    out << "pack expansion type ("
+        << expansionElt.getOpenedType()->getString(PO) << ")";
     break;
   }
   }

lib/Sema/ConstraintSystem.cpp

@@ -1018,15 +1018,17 @@ Type ConstraintSystem::openPackExpansionType(PackExpansionType *expansion,
     return known->second;
 
   auto *expansionVar = createTypeVariable(
-      getConstraintLocator({}, ConstraintLocator::PackExpansionType),
+      getConstraintLocator(
+          {}, LocatorPathElt::PackExpansionType(openedPackExpansion)),
       TVO_PackExpansion);
 
   // This constraint is important to make sure that pack expansion always
   // has a binding and connect pack expansion var to any type variables
   // that appear in pattern and shape types.
   addUnsolvedConstraint(Constraint::create(
       *this, ConstraintKind::Defaultable, expansionVar, openedPackExpansion,
-      getConstraintLocator({}, ConstraintLocator::PackExpansionType)));
+      getConstraintLocator(
+          {}, LocatorPathElt::PackExpansionType(openedPackExpansion))));
 
   OpenedPackExpansionTypes[openedPackExpansion] = expansionVar;
   return expansionVar;