Sema: Record conversion restrictions in the trail

Author: slavapestov

include/swift/Sema/CSTrail.h

@@ -36,7 +36,7 @@ class SolverTrail {
 public:
 
   /// The kind of change made to the graph.
-  enum class ChangeKind {
+  enum class ChangeKind: unsigned {
     /// Added a new vertex to the constraint graph.
     AddedTypeVariable,
     /// Added a new constraint to the constraint graph.
@@ -53,6 +53,8 @@ class SolverTrail {
     RetractedBindings,
     /// Set the fixed type or parent and flags for a type variable.
     UpdatedTypeVariable,
+    /// Recorded a conversion restriction kind.
+    AddedConversionRestriction,
   };
 
   /// A change made to the constraint system.
@@ -64,6 +66,9 @@ class SolverTrail {
     /// The kind of change.
     ChangeKind Kind;
 
+    /// Extra storage.
+    unsigned Options;
+
     union {
       TypeVariableType *TypeVar;
 
@@ -97,10 +102,15 @@ class SolverTrail {
 
         /// The representative of the equivalence class, or the fixed type.
         llvm::PointerUnion<TypeVariableType *, TypeBase *> ParentOrFixed;
-
-        /// The saved value of TypeVariableType::Implementation::getRawOptions().
-        unsigned Options;
       } Update;
+
+      struct {
+        /// The source type.
+        Type SrcType;
+
+        /// The destination type.
+        Type DstType;
+      } Restriction;
     };
 
     Change() : Kind(ChangeKind::AddedTypeVariable), TypeVar(nullptr) { }
@@ -137,6 +147,9 @@ class SolverTrail {
                llvm::PointerUnion<TypeVariableType *, TypeBase *> parentOrFixed,
                unsigned options);
 
+    /// Create a change that recorded a restriction.
+    static Change addedConversionRestriction(Type srcType, Type dstType);
+
     /// Undo this change, reverting the constraint graph to the state it
     /// had prior to this change.
     ///

include/swift/Sema/ConstraintSystem.h

@@ -2314,7 +2314,7 @@ class ConstraintSystem {
   /// there are multiple ways in which one type could convert to another, e.g.,
   /// given class types A and B, the solver might choose either a superclass
   /// conversion or a user-defined conversion.
-  llvm::MapVector<std::pair<TypeBase *, TypeBase *>, ConversionRestrictionKind>
+  llvm::DenseMap<std::pair<TypeBase *, TypeBase *>, ConversionRestrictionKind>
       ConstraintRestrictions;
 
   /// The set of fixes applied to make the solution work.
@@ -2858,9 +2858,6 @@ class ConstraintSystem {
     /// The length of \c Trail.
     unsigned numTrailChanges;
 
-    /// The length of \c ConstraintRestrictions.
-    unsigned numConstraintRestrictions;
-
     /// The length of \c Fixes.
     unsigned numFixes;
 
@@ -4218,6 +4215,13 @@ class ConstraintSystem {
                        bool updateState = true,
                        bool notifyBindingInference = true);
 
+  /// Update ConstraintRestrictions and record a change in the trail.
+  void addConversionRestriction(Type srcType, Type dstType,
+                                ConversionRestrictionKind restriction);
+
+  /// Called to undo the above change.
+  void removeConversionRestriction(Type srcType, Type dstType);
+
   /// Determine whether the given type is a dictionary and, if so, provide the
   /// key and value types for the dictionary.
   static std::optional<std::pair<Type, Type>> isDictionaryType(Type type);

lib/Sema/CSSimplify.cpp

@@ -14676,8 +14676,7 @@ ConstraintSystem::simplifyRestrictedConstraint(
       addFixConstraint(fix, matchKind, type1, type2, locator);
     }
 
-    ConstraintRestrictions.insert({
-        std::make_pair(type1.getPointer(), type2.getPointer()), restriction});
+    addConversionRestriction(type1, type2, restriction);
     return SolutionKind::Solved;
   }
   case SolutionKind::Unsolved:

lib/Sema/CSSolver.cpp

@@ -290,10 +290,10 @@ void ConstraintSystem::applySolution(const Solution &solution) {
   // Register constraint restrictions.
   // FIXME: Copy these directly into some kind of partial solution?
   for ( auto &restriction : solution.ConstraintRestrictions) {
-    auto *type1 = restriction.first.first.getPointer();
-    auto *type2 = restriction.first.second.getPointer();
+    auto type1 = restriction.first.first;
+    auto type2 = restriction.first.second;
 
-    ConstraintRestrictions.insert({{type1, type2}, restriction.second});
+    addConversionRestriction(type1, type2, restriction.second);
   }
 
   // Register the solution's disjunction choices.
@@ -652,7 +652,6 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numTrailChanges = cs.solverState->Trail.size();
 
   numTypeVariables = cs.TypeVariables.size();
-  numConstraintRestrictions = cs.ConstraintRestrictions.size();
   numFixes = cs.Fixes.size();
   numFixedRequirements = cs.FixedRequirements.size();
   numDisjunctionChoices = cs.DisjunctionChoices.size();
@@ -720,9 +719,6 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // constraints introduced by the current scope.
   cs.solverState->rollback(this);
 
-  // Remove any constraint restrictions.
-  truncate(cs.ConstraintRestrictions, numConstraintRestrictions);
-
   // Remove any fixes.
   truncate(cs.Fixes, numFixes);
 

lib/Sema/CSTrail.cpp

@@ -118,7 +118,16 @@ SolverTrail::Change::updatedTypeVariable(
   result.Kind = ChangeKind::UpdatedTypeVariable;
   result.Update.TypeVar = typeVar;
   result.Update.ParentOrFixed = parentOrFixed;
-  result.Update.Options = options;
+  result.Options = options;
+  return result;
+}
+
+SolverTrail::Change
+SolverTrail::Change::addedConversionRestriction(Type srcType, Type dstType) {
+  Change result;
+  result.Kind = ChangeKind::AddedConversionRestriction;
+  result.Restriction.SrcType = srcType;
+  result.Restriction.DstType = dstType;
   return result;
 }
 
@@ -157,9 +166,14 @@ void SolverTrail::Change::undo(ConstraintSystem &cs) const {
     break;
 
   case ChangeKind::UpdatedTypeVariable:
-    Update.TypeVar->getImpl().setRawOptions(Update.Options);
+    Update.TypeVar->getImpl().setRawOptions(Options);
     Update.TypeVar->getImpl().ParentOrFixed = Update.ParentOrFixed;
     break;
+
+  case ChangeKind::AddedConversionRestriction:
+    cs.removeConversionRestriction(Restriction.SrcType,
+                                   Restriction.DstType);
+    break;
   }
 }
 
@@ -229,7 +243,7 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
     out << ")\n";
     break;
 
-  case ChangeKind::UpdatedTypeVariable:
+  case ChangeKind::UpdatedTypeVariable: {
     out << "(updated type variable ";
     Update.TypeVar->print(out, PO);
 
@@ -243,7 +257,16 @@ void SolverTrail::Change::dump(llvm::raw_ostream &out,
       parentOrFixed.get<TypeBase *>()->print(out, PO);
     }
     out << " with options 0x";
-    out.write_hex(Update.Options);
+    out.write_hex(Options);
+    out << ")\n";
+    break;
+  }
+
+  case ChangeKind::AddedConversionRestriction:
+    out << "(added restriction with source ";
+    Restriction.SrcType->print(out, PO);
+    out << " and destination ";
+    Restriction.DstType->print(out, PO);
     out << ")\n";
     break;
   }

lib/Sema/ConstraintSystem.cpp

@@ -257,6 +257,28 @@ void ConstraintSystem::addTypeVariableConstraintsToWorkList(
       activateConstraint(constraint);
 }
 
+void ConstraintSystem::addConversionRestriction(
+    Type srcType, Type dstType,
+    ConversionRestrictionKind restriction) {
+  auto key = std::make_pair(srcType.getPointer(), dstType.getPointer());
+  bool inserted = ConstraintRestrictions.insert(
+      std::make_pair(key, restriction)).second;
+  if (!inserted)
+    return;
+
+  if (isRecordingChanges()) {
+    recordChange(SolverTrail::Change::addedConversionRestriction(
+      srcType, dstType));
+  }
+}
+
+void ConstraintSystem::removeConversionRestriction(
+    Type srcType, Type dstType) {
+  auto key = std::make_pair(srcType.getPointer(), dstType.getPointer());
+  bool erased = ConstraintRestrictions.erase(key);
+  ASSERT(erased);
+}
+
 /// Retrieve a dynamic result signature for the given declaration.
 static std::tuple<char, ObjCSelector, CanType>
 getDynamicResultSignature(ValueDecl *decl) {