Merge pull request swiftlang#77728 from slavapestov/cg-cleanup

Sema: Small ConstraintGraph cleanup

Author: slavapestov

include/swift/Sema/CSBindings.h

@@ -218,11 +218,6 @@ struct LiteralRequirement {
 };
 
 struct PotentialBindings {
-  /// The constraint system this type variable and its bindings belong to.
-  ConstraintSystem &CS;
-
-  TypeVariableType *TypeVar;
-
   /// The set of all constraints that have been added via infer().
   llvm::SmallPtrSet<Constraint *, 2> Constraints;
 
@@ -265,18 +260,13 @@ struct PotentialBindings {
   llvm::SmallSetVector<std::pair<TypeVariableType *, Constraint *>, 4> SupertypeOf;
   llvm::SmallSetVector<std::pair<TypeVariableType *, Constraint *>, 4> EquivalentTo;
 
-  PotentialBindings(ConstraintSystem &cs, TypeVariableType *typeVar)
-      : CS(cs), TypeVar(typeVar) {}
-
   void addDefault(Constraint *constraint);
 
   void addLiteral(Constraint *constraint);
 
   /// Add a potential binding to the list of bindings,
   /// coalescing supertype bounds when we are able to compute the meet.
-  void addPotentialBinding(PotentialBinding binding);
-
-  bool isGenericParameter() const;
+  void addPotentialBinding(TypeVariableType *typeVar, PotentialBinding binding);
 
   bool isSubtypeOf(TypeVariableType *typeVar) const {
     return llvm::any_of(
@@ -291,17 +281,26 @@ struct PotentialBindings {
   /// Attempt to infer a new binding and other useful information
   /// (i.e. whether bindings should be delayed) from the given
   /// relational constraint.
-  std::optional<PotentialBinding> inferFromRelational(Constraint *constraint);
+  std::optional<PotentialBinding> inferFromRelational(
+      ConstraintSystem &CS,
+      TypeVariableType *TypeVar,
+      Constraint *constraint);
 
 public:
-  void infer(Constraint *constraint);
+  void infer(ConstraintSystem &CS,
+             TypeVariableType *TypeVar,
+             Constraint *constraint);
 
   /// Retract all bindings and other information related to a given
   /// constraint from this binding set.
   ///
   /// This would happen when constraint is simplified or solver backtracks
   /// (either from overload choice or (some) type variable binding).
-  void retract(Constraint *constraint);
+  void retract(ConstraintSystem &CS,
+               TypeVariableType *TypeVar,
+               Constraint *constraint);
+
+  void reset();
 };
 
 
@@ -388,8 +387,9 @@ class BindingSet {
   /// subtype/conversion/equivalence relations with other type variables.
   std::optional<llvm::SmallPtrSet<Constraint *, 4>> TransitiveProtocols;
 
-  BindingSet(const PotentialBindings &info)
-      : CS(info.CS), TypeVar(info.TypeVar), Info(info) {
+  BindingSet(ConstraintSystem &CS, TypeVariableType *TypeVar,
+             const PotentialBindings &info)
+      : CS(CS), TypeVar(TypeVar), Info(info) {
     for (const auto &binding : info.Bindings)
       addBinding(binding, /*isTransitive=*/false);
 
@@ -405,7 +405,7 @@ class BindingSet {
 
   ConstraintSystem &getConstraintSystem() const { return CS; }
 
-  TypeVariableType *getTypeVariable() const { return Info.TypeVar; }
+  TypeVariableType *getTypeVariable() const { return TypeVar; }
 
   /// Check whether this binding set belongs to a type variable
   /// that represents a result type of a closure.

include/swift/Sema/ConstraintGraph.h

@@ -84,7 +84,10 @@ class ConstraintGraphNode {
   /// as this type variable.
   ArrayRef<TypeVariableType *> getEquivalenceClass() const;
 
-  inference::PotentialBindings &getCurrentBindings();
+  inference::PotentialBindings &getCurrentBindings() {
+    assert(forRepresentativeVar());
+    return Bindings;
+  }
 
 private:
   /// Determines whether the type variable associated with this node
@@ -180,7 +183,7 @@ class ConstraintGraphNode {
   TypeVariableType *TypeVar;
 
   /// The set of bindings associated with this type variable.
-  std::optional<inference::PotentialBindings> Bindings;
+  inference::PotentialBindings Bindings;
 
   /// The vector of constraints that mention this type variable, in a stable
   /// order for iteration.
@@ -244,14 +247,11 @@ class ConstraintGraph {
   /// Retrieve the constraint system this graph describes.
   ConstraintSystem &getConstraintSystem() const { return CS; }
 
-  /// Access the node corresponding to the given type variable.
-  ConstraintGraphNode &operator[](TypeVariableType *typeVar) {
-    return lookupNode(typeVar).first;
-  }
+  /// Add a new vertex to the graph.
+  void addTypeVariable(TypeVariableType *typeVar);
 
-  /// Retrieve the node and index corresponding to the given type variable.
-  std::pair<ConstraintGraphNode &, unsigned> 
-  lookupNode(TypeVariableType *typeVar);
+  /// Look up the vertex associated with the given type variable.
+  ConstraintGraphNode &operator[](TypeVariableType *typeVar);
 
   /// Add a new constraint to the graph.
   void addConstraint(Constraint *constraint);

include/swift/Sema/ConstraintSystem.h

@@ -30,7 +30,6 @@
 #include "swift/Basic/Debug.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/Basic/OptionSet.h"
-#include "swift/Sema/CSBindings.h"
 #include "swift/Sema/CSFix.h"
 #include "swift/Sema/CSTrail.h"
 #include "swift/Sema/Constraint.h"
@@ -1494,7 +1493,7 @@ class Solution {
   DeclContext *getDC() const;
 
   /// The set of type bindings.
-  llvm::DenseMap<TypeVariableType *, Type> typeBindings;
+  llvm::MapVector<TypeVariableType *, Type> typeBindings;
 
   /// The set of overload choices along with their types.
   llvm::DenseMap<ConstraintLocator *, SelectedOverload> overloadChoices;
@@ -2689,10 +2688,6 @@ class ConstraintSystem {
   /// we're exploring.
   SolverState *solverState = nullptr;
 
-  bool isRecordingChanges() const {
-    return solverState && !solverState->Trail.isUndoActive();
-  }
-
   void recordChange(SolverTrail::Change change) {
     solverState->Trail.recordChange(change);
   }

lib/IDE/KeyPathCompletion.cpp

@@ -53,7 +53,7 @@ void KeyPathTypeCheckCompletionCallback::sawSolutionImpl(
             return Entry.first->getImpl().getLocator() == RootLocator;
           });
       if (BaseVariableTypeBinding != S.typeBindings.end()) {
-        BaseType = S.simplifyType(BaseVariableTypeBinding->getSecond());
+        BaseType = S.simplifyType(BaseVariableTypeBinding->second);
       }
     }
   } else {

lib/Sema/CSBindings.cpp

@@ -35,11 +35,11 @@ static std::optional<Type> checkTypeOfBinding(TypeVariableType *typeVar,
                                               Type type);
 
 bool BindingSet::forClosureResult() const {
-  return Info.TypeVar->getImpl().isClosureResultType();
+  return TypeVar->getImpl().isClosureResultType();
 }
 
 bool BindingSet::forGenericParameter() const {
-  return bool(Info.TypeVar->getImpl().getGenericParameter());
+  return bool(TypeVar->getImpl().getGenericParameter());
 }
 
 bool BindingSet::canBeNil() const {
@@ -54,10 +54,10 @@ bool BindingSet::isDirectHole() const {
     return false;
 
   return Bindings.empty() && getNumViableLiteralBindings() == 0 &&
-         Defaults.empty() && Info.TypeVar->getImpl().canBindToHole();
+         Defaults.empty() && TypeVar->getImpl().canBindToHole();
 }
 
-bool PotentialBindings::isGenericParameter() const {
+static bool isGenericParameter(TypeVariableType *TypeVar) {
   auto *locator = TypeVar->getImpl().getLocator();
   return locator && locator->isLastElement<LocatorPathElt::GenericParameter>();
 }
@@ -178,7 +178,7 @@ bool BindingSet::involvesTypeVariables() const {
 
 bool BindingSet::isPotentiallyIncomplete() const {
   // Generic parameters are always potentially incomplete.
-  if (Info.isGenericParameter())
+  if (isGenericParameter(TypeVar))
     return true;
 
   // Key path literal type is incomplete until there is a
@@ -1168,7 +1168,8 @@ LiteralRequirement::isCoveredBy(const PotentialBinding &binding, bool canBeNil,
   } while (true);
 }
 
-void PotentialBindings::addPotentialBinding(PotentialBinding binding) {
+void PotentialBindings::addPotentialBinding(TypeVariableType *TypeVar,
+                                            PotentialBinding binding) {
   assert(!binding.BindingType->is<ErrorType>());
 
   // If the type variable can't bind to an lvalue, make sure the
@@ -1418,7 +1419,7 @@ BindingSet ConstraintSystem::getBindingsFor(TypeVariableType *typeVar,
          "not a representative");
   assert(!typeVar->getImpl().getFixedType(nullptr) && "has a fixed type");
 
-  BindingSet bindings{CG[typeVar].getCurrentBindings()};
+  BindingSet bindings(*this, typeVar, CG[typeVar].getCurrentBindings());
 
   if (finalize) {
     llvm::SmallDenseMap<TypeVariableType *, BindingSet> cache;
@@ -1461,7 +1462,9 @@ static std::optional<Type> checkTypeOfBinding(TypeVariableType *typeVar,
 }
 
 std::optional<PotentialBinding>
-PotentialBindings::inferFromRelational(Constraint *constraint) {
+PotentialBindings::inferFromRelational(ConstraintSystem &CS,
+                                       TypeVariableType *TypeVar,
+                                       Constraint *constraint) {
   assert(constraint->getClassification() ==
              ConstraintClassification::Relational &&
          "only relational constraints handled here");
@@ -1657,7 +1660,7 @@ PotentialBindings::inferFromRelational(Constraint *constraint) {
     // Since inference now happens during constraint generation,
     // this hack should be allowed in both `Solving`
     // (during non-diagnostic mode) and `ConstraintGeneration` phases.
-    if (isGenericParameter() &&
+    if (isGenericParameter(TypeVar) &&
         (!CS.shouldAttemptFixes() ||
          CS.getPhase() == ConstraintSystemPhase::ConstraintGeneration)) {
       type = fnTy->withExtInfo(fnTy->getExtInfo().withNoEscape(false));
@@ -1766,12 +1769,14 @@ PotentialBindings::inferFromRelational(Constraint *constraint) {
 /// Retrieve the set of potential type bindings for the given
 /// representative type variable, along with flags indicating whether
 /// those types should be opened.
-void PotentialBindings::infer(Constraint *constraint) {
+void PotentialBindings::infer(ConstraintSystem &CS,
+                              TypeVariableType *TypeVar,
+                              Constraint *constraint) {
   if (!Constraints.insert(constraint).second)
     return;
 
   // Record the change, if there are active scopes.
-  if (CS.isRecordingChanges())
+  if (CS.solverState && !CS.solverState->Trail.isUndoActive())
     CS.recordChange(SolverTrail::Change::InferredBindings(TypeVar, constraint));
 
   switch (constraint->getKind()) {
@@ -1787,11 +1792,11 @@ void PotentialBindings::infer(Constraint *constraint) {
   case ConstraintKind::OptionalObject:
   case ConstraintKind::UnresolvedMemberChainBase:
   case ConstraintKind::LValueObject: {
-    auto binding = inferFromRelational(constraint);
+    auto binding = inferFromRelational(CS, TypeVar, constraint);
     if (!binding)
       break;
 
-    addPotentialBinding(*binding);
+    addPotentialBinding(TypeVar, *binding);
     break;
   }
   case ConstraintKind::KeyPathApplication: {
@@ -1940,12 +1945,14 @@ void PotentialBindings::infer(Constraint *constraint) {
   }
 }
 
-void PotentialBindings::retract(Constraint *constraint) {
+void PotentialBindings::retract(ConstraintSystem &CS,
+                                TypeVariableType *TypeVar,
+                                Constraint *constraint) {
   if (!Constraints.erase(constraint))
     return;
 
   // Record the change, if there are active scopes.
-  if (CS.isRecordingChanges())
+  if (CS.solverState && !CS.solverState->Trail.isUndoActive())
     CS.recordChange(SolverTrail::Change::RetractedBindings(TypeVar, constraint));
 
   LLVM_DEBUG(
@@ -2011,6 +2018,23 @@ void PotentialBindings::retract(Constraint *constraint) {
   EquivalentTo.remove_if(hasMatchingSource);
 }
 
+void PotentialBindings::reset() {
+  if (CONDITIONAL_ASSERT_enabled()) {
+    ASSERT(Constraints.empty());
+    ASSERT(Bindings.empty());
+    ASSERT(Protocols.empty());
+    ASSERT(Literals.empty());
+    ASSERT(Defaults.empty());
+    ASSERT(DelayedBy.empty());
+    ASSERT(AdjacentVars.empty());
+    ASSERT(SubtypeOf.empty());
+    ASSERT(SupertypeOf.empty());
+    ASSERT(EquivalentTo.empty());
+  }
+
+  AssociatedCodeCompletionToken = ASTNode();
+}
+
 void BindingSet::forEachLiteralRequirement(
     llvm::function_ref<void(KnownProtocolKind)> callback) const {
   for (const auto &literal : Literals) {

lib/Sema/CSGen.cpp

@@ -4930,7 +4930,7 @@ bool ConstraintSystem::generateConstraints(
     // Cache the outer generic environment, if it exists.
     if (target.getPackElementEnv()) {
       PackElementGenericEnvironments.push_back(target.getPackElementEnv());
-      ASSERT(!isRecordingChanges() && "Need to record a change");
+      ASSERT(!solverState && "Need to record a change");
     }
 
     // For a for-each statement, generate constraints for the pattern, where

lib/Sema/CSRanking.cpp

@@ -1350,6 +1350,9 @@ SolutionCompareResult ConstraintSystem::compareSolutions(
   const auto &bindings2 = solutions[idx2].typeBindings;
 
   for (const auto &binding1 : bindings1) {
+    if (!binding1.second)
+      continue;
+
     auto *typeVar = binding1.first;
     auto *loc = typeVar->getImpl().getLocator();
 
@@ -1375,6 +1378,9 @@ SolutionCompareResult ConstraintSystem::compareSolutions(
     if (binding2 == bindings2.end())
       continue;
 
+    if (!binding2->second)
+      continue;
+
     TypeBindingsToCompare typesToCompare(binding1.second, binding2->second);
 
     // For short-form and self-delegating init calls, we want to prefer

lib/Sema/CSSolver.cpp

@@ -103,8 +103,10 @@ Solution ConstraintSystem::finalize() {
     // This type variable has no binding. Allowed only
     // when `FreeTypeVariableBinding::Allow` is set,
     // which is checked above.
-    if (!getFixedType(tv))
+    if (!getFixedType(tv)) {
+      solution.typeBindings[tv] = Type();
       continue;
+    }
 
     solution.typeBindings[tv] = simplifyType(tv)->reconstituteSugar(false);
   }
@@ -274,10 +276,15 @@ void ConstraintSystem::replaySolution(const Solution &solution,
   if (shouldIncreaseScore)
     replayScore(solution.getFixedScore());
 
-  // Assign fixed types to the type variables solved by this solution.
   for (auto binding : solution.typeBindings) {
     // If we haven't seen this type variable before, record it now.
     addTypeVariable(binding.first);
+  }
+
+  // Assign fixed types to the type variables solved by this solution.
+  for (auto binding : solution.typeBindings) {
+    if (!binding.second)
+      continue;
 
     // If we don't already have a fixed type for this type variable,
     // assign the fixed type from the solution.