Sema: Collect opened pack expansion types in PreparedOverload

Author: slavapestov

include/swift/Sema/ConstraintSystem.h

@@ -4360,7 +4360,9 @@ class ConstraintSystem {
 
   /// Update OpenedPackExpansionTypes and record a change in the trail.
   void recordOpenedPackExpansionType(PackExpansionType *expansion,
-                                     TypeVariableType *expansionVar);
+                                     TypeVariableType *expansionVar,
+                                     PreparedOverload *preparedOverload
+                                        = nullptr);
 
   /// Undo the above change.
   void removeOpenedPackExpansionType(PackExpansionType *expansion) {

include/swift/Sema/PreparedOverload.h

@@ -35,6 +35,8 @@ struct PreparedOverload {
   SmallVector<Constraint *, 2> Constraints;
   SmallVector<OpenedType, 2> Replacements;
   ExistentialArchetypeType *OpenedExistential = nullptr;
+  SmallVector<std::pair<PackExpansionType *, TypeVariableType *>>
+    OpenedPackExpansionTypes;
 
   void discharge(ConstraintSystem &cs, ConstraintLocatorBuilder locator) const;
 };

lib/Sema/TypeOfReference.cpp

@@ -314,31 +314,54 @@ Type ConstraintSystem::openPackExpansionType(PackExpansionType *expansion,
 
   auto openedPackExpansion = PackExpansionType::get(patternType, shapeType);
 
-  auto known = OpenedPackExpansionTypes.find(openedPackExpansion);
-  if (known != OpenedPackExpansionTypes.end())
-    return known->second;
+  // FIXME: It's silly that we need to do this. The whole concept of
+  // "opening" pack expansions is broken.
+  {
+    if (preparedOverload) {
+      for (auto pair : preparedOverload->OpenedPackExpansionTypes) {
+        if (pair.first == openedPackExpansion)
+          return pair.second;
+      }
+    }
+
+    auto known = OpenedPackExpansionTypes.find(openedPackExpansion);
+    if (known != OpenedPackExpansionTypes.end())
+      return known->second;
+  }
 
   auto *expansionLoc = getConstraintLocator(locator.withPathElement(
       LocatorPathElt::PackExpansionType(openedPackExpansion)));
 
   auto *expansionVar = createTypeVariable(expansionLoc, TVO_PackExpansion,
                                           preparedOverload);
 
-  ASSERT(!preparedOverload && "Implement below");
-
   // 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::FallbackType,
-                                           expansionVar, openedPackExpansion,
-                                           expansionLoc));
-
-  recordOpenedPackExpansionType(openedPackExpansion, expansionVar);
+  auto *c = Constraint::create(*this, ConstraintKind::FallbackType,
+                               expansionVar, openedPackExpansion,
+                               expansionLoc);
+  if (preparedOverload)
+    preparedOverload->Constraints.push_back(c);
+  else
+    addUnsolvedConstraint(c);
+
+  recordOpenedPackExpansionType(openedPackExpansion, expansionVar,
+                                preparedOverload);
   return expansionVar;
 }
 
 void ConstraintSystem::recordOpenedPackExpansionType(PackExpansionType *expansion,
-                                                     TypeVariableType *expansionVar) {
+                                                     TypeVariableType *expansionVar,
+                                                     PreparedOverload *preparedOverload) {
+  if (preparedOverload) {
+    ASSERT(PreparingOverload);
+    preparedOverload->OpenedPackExpansionTypes.push_back({expansion, expansionVar});
+    return;
+  }
+
+  ASSERT(!PreparingOverload);
+
   bool inserted = OpenedPackExpansionTypes.insert({expansion, expansionVar}).second;
   ASSERT(inserted);
 
@@ -2476,6 +2499,9 @@ void PreparedOverload::discharge(ConstraintSystem &cs,
     cs.recordOpenedExistentialType(cs.getConstraintLocator(locator),
                                    OpenedExistential);
   }
+  for (auto pair : OpenedPackExpansionTypes) {
+    cs.recordOpenedPackExpansionType(pair.first, pair.second);
+  }
 }
 
 void ConstraintSystem::resolveOverload(ConstraintLocator *locator,