Sema: Rework PreparedOverload to store a list of changes

Author: slavapestov

include/swift/Sema/ConstraintSystem.h

@@ -4438,7 +4438,8 @@ class ConstraintSystem {
 
   /// Update OpenedTypes and record a change in the trail.
   void recordOpenedType(
-      ConstraintLocator *locator, ArrayRef<OpenedType> openedTypes);
+      ConstraintLocator *locator, ArrayRef<OpenedType> openedTypes,
+      PreparedOverload *preparedOverload = nullptr);
 
   /// Record the set of opened types for the given locator.
   void recordOpenedTypes(

include/swift/Sema/PreparedOverload.h

@@ -31,15 +31,130 @@ class ConstraintSystem;
 /// variable.
 using OpenedType = std::pair<GenericTypeParamType *, TypeVariableType *>;
 
+/// A "pre-cooked" representation of all type variables and constraints
+/// that are generated as part of an overload choice.
 struct PreparedOverload {
-  SmallVector<TypeVariableType *, 2> TypeVariables;
-  SmallVector<Constraint *, 2> Constraints;
-  SmallVector<OpenedType, 2> Replacements;
-  ExistentialArchetypeType *OpenedExistential = nullptr;
-  SmallVector<std::pair<PackExpansionType *, TypeVariableType *>>
-    OpenedPackExpansionTypes;
-  SmallVector<AppliedPropertyWrapper, 2> PropertyWrappers;
-  SmallVector<std::pair<ConstraintFix *, unsigned>, 2> Fixes;
+  /// A change to be introduced into the constraint system when this
+  /// overload choice is chosen.
+  struct Change {
+    enum ChangeKind : unsigned {
+      /// A generic parameter was opened to a type variable.
+      AddedTypeVariable,
+
+      /// A generic requirement was opened to a constraint.
+      AddedConstraint,
+
+      /// A mapping of generic parameter types to type variables
+      /// was recorded.
+      OpenedTypes,
+
+      /// An existential type was opened.
+      OpenedExistentialType,
+
+      /// A pack expansion type was opened.
+      OpenedPackExpansionType,
+
+      /// A property wrapper was applied to a parameter.
+      AppliedPropertyWrapper,
+
+      /// A fix was recorded because a property wrapper application failed.
+      AddedFix
+    };
+
+    /// The kind of change.
+    ChangeKind Kind;
+
+    union {
+      /// For ChangeKind::AddedTypeVariable.
+      TypeVariableType *TypeVar;
+
+      /// For ChangeKind::AddedConstraint.
+      Constraint *TheConstraint;
+
+      /// For ChangeKind::OpenedTypes.
+      struct {
+        const OpenedType *Data;
+        size_t Count;
+      } Replacements;
+
+      /// For ChangeKind::OpenedExistentialType.
+      ExistentialArchetypeType *TheExistential;
+
+      /// For ChangeKind::OpenedPackExpansionType.
+      struct {
+        PackExpansionType *TheExpansion;
+        TypeVariableType *TypeVar;
+      } PackExpansion;
+
+      /// For ChangeKind::AppliedPropertyWrapper.
+      struct {
+        TypeBase *WrapperType;
+        PropertyWrapperInitKind InitKind;
+      } PropertyWrapper;
+
+      /// For ChangeKind::Fix.
+      struct {
+        ConstraintFix *TheFix;
+        unsigned Impact;
+      } Fix;
+    };
+  };
+
+  SmallVector<Change, 8> Changes;
+
+  void addedTypeVariable(TypeVariableType *typeVar) {
+    Change change;
+    change.Kind = Change::AddedTypeVariable;
+    change.TypeVar = typeVar;
+    Changes.push_back(change);
+  }
+
+  void addedConstraint(Constraint *constraint) {
+    Change change;
+    change.Kind = Change::AddedConstraint;
+    change.TheConstraint = constraint;
+    Changes.push_back(change);
+  }
+
+  void openedTypes(ArrayRef<OpenedType> replacements) {
+    Change change;
+    change.Kind = Change::OpenedTypes;
+    change.Replacements.Data = replacements.data();
+    change.Replacements.Count = replacements.size();
+    Changes.push_back(change);
+  }
+
+  void openedExistentialType(ExistentialArchetypeType *openedExistential) {
+    Change change;
+    change.Kind = Change::OpenedExistentialType;
+    change.TheExistential = openedExistential;
+    Changes.push_back(change);
+  }
+
+  void openedPackExpansionType(PackExpansionType *packExpansion,
+                               TypeVariableType *typeVar) {
+    Change change;
+    change.Kind = Change::OpenedPackExpansionType;
+    change.PackExpansion.TheExpansion = packExpansion;
+    change.PackExpansion.TypeVar = typeVar;
+    Changes.push_back(change);
+  }
+
+  void appliedPropertyWrapper(AppliedPropertyWrapper wrapper) {
+    Change change;
+    change.Kind = Change::AppliedPropertyWrapper;
+    change.PropertyWrapper.WrapperType = wrapper.wrapperType.getPointer();
+    change.PropertyWrapper.InitKind = wrapper.initKind;
+    Changes.push_back(change);
+  }
+
+  void addedFix(ConstraintFix *fix, unsigned impact) {
+    Change change;
+    change.Kind = Change::AddedFix;
+    change.Fix.TheFix = fix;
+    change.Fix.Impact = impact;
+    Changes.push_back(change);
+  }
 
   void discharge(ConstraintSystem &cs, ConstraintLocatorBuilder locator) const;
 };

lib/Sema/CSGen.cpp

@@ -5122,7 +5122,7 @@ void ConstraintSystem::applyPropertyWrapper(
     PreparedOverload *preparedOverload) {
   if (preparedOverload) {
     ASSERT(PreparingOverload);
-    preparedOverload->PropertyWrappers.push_back(applied);
+    preparedOverload->appliedPropertyWrapper(applied);
     return;
   }
 

lib/Sema/CSSimplify.cpp

@@ -15407,7 +15407,7 @@ bool ConstraintSystem::recordFix(ConstraintFix *fix, unsigned impact,
                                  PreparedOverload *preparedOverload) {
   if (preparedOverload) {
     ASSERT(PreparingOverload);
-    preparedOverload->Fixes.push_back({fix, impact});
+    preparedOverload->addedFix(fix, impact);
     return true;
   }
 
@@ -16350,7 +16350,7 @@ void ConstraintSystem::addConstraint(ConstraintKind kind, Type first,
     auto c = Constraint::create(*this, kind, first, second,
                                 getConstraintLocator(locator));
     if (isFavored) c->setFavored();
-    preparedOverload->Constraints.push_back(c);
+    preparedOverload->addedConstraint(c);
     return;
   }
 

lib/Sema/CSSolver.cpp

@@ -66,7 +66,7 @@ TypeVariableType *ConstraintSystem::createTypeVariable(
                                      locator, options);
   if (preparedOverload) {
     ASSERT(PreparingOverload);
-    preparedOverload->TypeVariables.push_back(tv);
+    preparedOverload->addedTypeVariable(tv);
   } else {
     addTypeVariable(tv);
   }

lib/Sema/ConstraintSystem.cpp

@@ -906,8 +906,7 @@ void ConstraintSystem::recordOpenedExistentialType(
     ExistentialArchetypeType *opened,
     PreparedOverload *preparedOverload) {
   if (preparedOverload) {
-    ASSERT(!preparedOverload->OpenedExistential);
-    preparedOverload->OpenedExistential = opened;
+    preparedOverload->openedExistentialType(opened);
     return;
   }
 

lib/Sema/TypeOfReference.cpp

@@ -321,9 +321,10 @@ Type ConstraintSystem::openPackExpansionType(PackExpansionType *expansion,
   // "opening" pack expansions is broken.
   {
     if (preparedOverload) {
-      for (auto pair : preparedOverload->OpenedPackExpansionTypes) {
-        if (pair.first == openedPackExpansion)
-          return pair.second;
+      for (auto change : preparedOverload->Changes) {
+        if (change.Kind == PreparedOverload::Change::OpenedPackExpansionType &&
+            change.PackExpansion.TheExpansion == openedPackExpansion)
+          return change.PackExpansion.TypeVar;
       }
     }
 
@@ -345,7 +346,7 @@ Type ConstraintSystem::openPackExpansionType(PackExpansionType *expansion,
                                expansionVar, openedPackExpansion,
                                expansionLoc);
   if (preparedOverload)
-    preparedOverload->Constraints.push_back(c);
+    preparedOverload->addedConstraint(c);
   else
     addUnsolvedConstraint(c);
 
@@ -359,7 +360,7 @@ void ConstraintSystem::recordOpenedPackExpansionType(PackExpansionType *expansio
                                                      PreparedOverload *preparedOverload) {
   if (preparedOverload) {
     ASSERT(PreparingOverload);
-    preparedOverload->OpenedPackExpansionTypes.push_back({expansion, expansionVar});
+    preparedOverload->openedPackExpansionType(expansion, expansionVar);
     return;
   }
 
@@ -699,8 +700,17 @@ Type ConstraintSystem::getUnopenedTypeOfReference(
 }
 
 void ConstraintSystem::recordOpenedType(
-    ConstraintLocator *locator, ArrayRef<OpenedType> openedTypes) {
-  bool inserted = OpenedTypes.insert({locator, openedTypes}).second;
+    ConstraintLocator *locator, ArrayRef<OpenedType> replacements,
+    PreparedOverload *preparedOverload) {
+  if (preparedOverload) {
+    ASSERT(PreparingOverload);
+    preparedOverload->openedTypes(replacements);
+    return;
+  }
+
+  ASSERT(!PreparingOverload);
+
+  bool inserted = OpenedTypes.insert({locator, replacements}).second;
   ASSERT(inserted);
 
   if (solverState)
@@ -715,16 +725,6 @@ void ConstraintSystem::recordOpenedTypes(
   if (replacements.empty())
     return;
 
-  if (preparedOverload) {
-    ASSERT(PreparingOverload);
-    ASSERT(preparedOverload->Replacements.empty());
-    preparedOverload->Replacements.append(
-      replacements.begin(), replacements.end());
-    return;
-  }
-
-  ASSERT(!PreparingOverload);
-
   // If the last path element is an archetype or associated type, ignore it.
   SmallVector<LocatorPathElt, 2> pathElts;
   auto anchor = locator.getLocatorParts(pathElts);
@@ -746,7 +746,8 @@ void ConstraintSystem::recordOpenedTypes(
   // FIXME: Get rid of fixmeAllowDuplicates.
   if (!fixmeAllowDuplicates || OpenedTypes.count(locatorPtr) == 0)
     recordOpenedType(
-      locatorPtr, llvm::ArrayRef(openedTypes, replacements.size()));
+      locatorPtr, llvm::ArrayRef(openedTypes, replacements.size()),
+      preparedOverload);
 }
 
 /// Determine how many levels of argument labels should be removed from the
@@ -2496,35 +2497,55 @@ void ConstraintSystem::recordResolvedOverload(ConstraintLocator *locator,
 
 void PreparedOverload::discharge(ConstraintSystem &cs,
                                  ConstraintLocatorBuilder locator) const {
-  for (auto *tv : TypeVariables) {
-    cs.addTypeVariable(tv);
-  }
+  for (auto change : Changes) {
+    switch (change.Kind) {
+    case PreparedOverload::Change::AddedTypeVariable:
+      cs.addTypeVariable(change.TypeVar);
+      break;
 
-  for (auto *c : Constraints) {
-    cs.addUnsolvedConstraint(c);
-    cs.activateConstraint(c);
-  }
+    case PreparedOverload::Change::AddedConstraint:
+      cs.addUnsolvedConstraint(change.TheConstraint);
+      cs.activateConstraint(change.TheConstraint);
+      break;
 
-  cs.recordOpenedTypes(locator, Replacements);
+    case PreparedOverload::Change::OpenedTypes: {
+      auto *locatorPtr = cs.getConstraintLocator(locator);
+      ArrayRef<OpenedType> replacements(
+          change.Replacements.Data,
+          change.Replacements.Count);
 
-  if (OpenedExistential) {
-    cs.recordOpenedExistentialType(cs.getConstraintLocator(locator),
-                                   OpenedExistential);
-  }
+      // FIXME: Get rid of this conditional.
+      if (cs.getOpenedTypes(locatorPtr).empty())
+        cs.recordOpenedType(locatorPtr, replacements);
+      break;
+    }
 
-  for (auto pair : OpenedPackExpansionTypes) {
-    cs.recordOpenedPackExpansionType(pair.first, pair.second);
-  }
+    case PreparedOverload::Change::OpenedExistentialType: {
+      auto *locatorPtr = cs.getConstraintLocator(locator);
+      cs.recordOpenedExistentialType(locatorPtr,
+                                     change.TheExistential);
+      break;
+    }
+
+    case PreparedOverload::Change::OpenedPackExpansionType:
+      cs.recordOpenedPackExpansionType(
+          change.PackExpansion.TheExpansion,
+          change.PackExpansion.TypeVar);
+      break;
 
-  if (!PropertyWrappers.empty()) {
-    Expr *anchor = getAsExpr(cs.getConstraintLocator(locator)->getAnchor());
-    for (auto applied : PropertyWrappers) {
-      cs.applyPropertyWrapper(anchor, applied);
+    case PreparedOverload::Change::AppliedPropertyWrapper: {
+      auto *locatorPtr = cs.getConstraintLocator(locator);
+      Expr *anchor = getAsExpr(locatorPtr->getAnchor());
+      cs.applyPropertyWrapper(anchor,
+          { Type(change.PropertyWrapper.WrapperType),
+            change.PropertyWrapper.InitKind });
+      break;
     }
-  }
 
-  for (auto pair : Fixes) {
-    cs.recordFix(pair.first, pair.second);
+    case PreparedOverload::Change::AddedFix:
+      cs.recordFix(change.Fix.TheFix, change.Fix.Impact);
+      break;
+    }
   }
 }