Small fixes for diagnostics with prepared overloads enabled

include/swift/Sema/Constraint.h

@@ -883,11 +883,7 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     return Overload.Prepared;
   }
 
-  void setPreparedOverload(PreparedOverload *preparedOverload) {
-    ASSERT(Kind == ConstraintKind::BindOverload);
-    ASSERT(!Overload.Prepared);
-    Overload.Prepared = preparedOverload;
-  }
+  void setPreparedOverload(PreparedOverload *preparedOverload);
 
   FunctionType *getAppliedFunctionType() const {
     assert(Kind == ConstraintKind::ApplicableFunction);

include/swift/Sema/ConstraintSystem.h

@@ -2131,9 +2131,7 @@ class SyntacticElementTargetRewriter {
 
 enum class ConstraintSystemPhase {
   ConstraintGeneration,
-  Solving,
-  Diagnostics,
-  Finalization
+  Solving
 };
 
 /// Retrieve the closure type from the constraint system.
@@ -2693,18 +2691,7 @@ class ConstraintSystem {
     case ConstraintSystemPhase::Solving:
       // We can come back to constraint generation phase while
       // processing result builder body.
-      assert(newPhase == ConstraintSystemPhase::ConstraintGeneration ||
-             newPhase == ConstraintSystemPhase::Diagnostics ||
-             newPhase == ConstraintSystemPhase::Finalization);
-      break;
-
-    case ConstraintSystemPhase::Diagnostics:
-      assert(newPhase == ConstraintSystemPhase::Solving ||
-             newPhase == ConstraintSystemPhase::Finalization);
-      break;
-
-    case ConstraintSystemPhase::Finalization:
-      assert(newPhase == ConstraintSystemPhase::Diagnostics);
+      assert(newPhase == ConstraintSystemPhase::ConstraintGeneration);
       break;
     }
 #endif
@@ -4876,7 +4863,8 @@ class ConstraintSystem {
   /// Build and allocate a prepared overload in the solver arena.
   PreparedOverload *prepareOverload(OverloadChoice choice,
                                     DeclContext *useDC,
-                                    ConstraintLocator *locator);
+                                    ConstraintLocator *locator,
+                                    bool forDiagnostics);
 
   /// Populate the prepared overload with all type variables and constraints
   /// that are to be introduced into the constraint system when this choice

include/swift/Sema/PreparedOverload.h

@@ -104,9 +104,10 @@ struct PreparedOverloadChange {
     /// For ChangeKind::AddedConstraint.
     Constraint *TheConstraint;
 
+    /// For ChangeKind::AddedBindConstraint.
     struct {
       TypeBase *FirstType;
-      TypeBase * SecondType;
+      TypeBase *SecondType;
     } Bind;
 
     /// For ChangeKind::OpenedTypes.
@@ -159,17 +160,21 @@ class PreparedOverload final :
 private:
   Type OpenedType;
   Type ThrownErrorType;
-  size_t Count;
+  unsigned Count : 31;
+
+  /// A prepared overload for diagnostics is different than one without,
+  /// because of fixes and such.
+  unsigned ForDiagnostics : 1;
 
   size_t numTrailingObjects(OverloadToken<Change>) const {
     return Count;
   }
 
 public:
   PreparedOverload(Type openedType, Type thrownErrorType,
-                   ArrayRef<Change> changes)
+                   ArrayRef<Change> changes, bool forDiagnostics)
     : OpenedType(openedType), ThrownErrorType(thrownErrorType),
-      Count(changes.size()) {
+      Count(changes.size()), ForDiagnostics(forDiagnostics) {
     std::uninitialized_copy(changes.begin(), changes.end(),
                             getTrailingObjects());
   }
@@ -182,11 +187,19 @@ class PreparedOverload final :
     return ThrownErrorType;
   }
 
+  bool wasForDiagnostics() const {
+    return ForDiagnostics;
+  }
+
   ArrayRef<Change> getChanges() const { return getTrailingObjects(Count); }
 };
 
 struct PreparedOverloadBuilder {
   SmallVector<PreparedOverload::Change, 8> Changes;
+  ConstraintLocator *Locator;
+
+  PreparedOverloadBuilder(ConstraintLocator *locator)
+    : Locator(locator) {}
 
   void addedTypeVariable(TypeVariableType *typeVar) {
     PreparedOverload::Change change;

lib/Sema/CSSimplify.cpp

@@ -16529,14 +16529,22 @@ void ConstraintSystem::addConstraint(ConstraintKind kind, Type first,
                                      PreparedOverloadBuilder *preparedOverload) {
   if (preparedOverload) {
     ASSERT(PreparingOverload);
-    if (kind == ConstraintKind::Bind) {
+
+    auto *locatorPtr = getConstraintLocator(locator);
+
+    // Fast path to save on memory usage.
+    if (kind == ConstraintKind::Bind &&
+        locatorPtr == preparedOverload->Locator) {
       ASSERT(!isFavored);
       preparedOverload->addedBindConstraint(first, second);
       return;
     }
-    auto c = Constraint::create(*this, kind, first, second,
-                                getConstraintLocator(locator));
-    if (isFavored) c->setFavored();
+
+    auto c = Constraint::create(*this, kind, first, second, locatorPtr);
+
+    if (isFavored)
+      c->setFavored();
+
     preparedOverload->addedConstraint(c);
     return;
   }
@@ -16829,16 +16837,24 @@ ConstraintSystem::simplifyConstraint(const Constraint &constraint) {
 
     // FIXME: Transitional hack.
     bool enablePreparedOverloads = getASTContext().TypeCheckerOpts.SolverEnablePreparedOverloads;
+    bool forDiagnostics = inSalvageMode();
 
+    // Don't reuse prepared overloads from normal type checking in salvage(),
+    // since they will contain fixes and such.
     auto *preparedOverload = constraint.getPreparedOverload();
-    if (!preparedOverload) {
-      if (enablePreparedOverloads &&
-          constraint.getOverloadChoice().canBePrepared()) {
-        preparedOverload = prepareOverload(constraint.getOverloadChoice(),
-                                           constraint.getDeclContext(),
-                                           constraint.getLocator());
-        const_cast<Constraint &>(constraint).setPreparedOverload(preparedOverload);
-      }
+    if (preparedOverload &&
+        preparedOverload->wasForDiagnostics() != forDiagnostics) {
+      preparedOverload = nullptr;
+    }
+
+    if (!preparedOverload &&
+        enablePreparedOverloads &&
+        constraint.getOverloadChoice().canBePrepared()) {
+      preparedOverload = prepareOverload(constraint.getOverloadChoice(),
+                                         constraint.getDeclContext(),
+                                         constraint.getLocator(),
+                                         forDiagnostics);
+      const_cast<Constraint &>(constraint).setPreparedOverload(preparedOverload);
     }
 
     resolveOverload(constraint.getOverloadChoice(),

lib/Sema/CSSolver.cpp

@@ -1579,8 +1579,6 @@ void ConstraintSystem::solveImpl(SmallVectorImpl<Solution> &solutions) {
 
   setPhase(ConstraintSystemPhase::Solving);
 
-  SWIFT_DEFER { setPhase(ConstraintSystemPhase::Finalization); };
-
   // If constraint system failed while trying to
   // genenerate constraints, let's stop right here.
   if (failedConstraint)

lib/Sema/Constraint.cpp

@@ -20,6 +20,7 @@
 #include "swift/Basic/Compiler.h"
 #include "swift/Sema/Constraint.h"
 #include "swift/Sema/ConstraintSystem.h"
+#include "swift/Sema/PreparedOverload.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
@@ -1142,3 +1143,18 @@ void *Constraint::operator new(size_t bytes, ConstraintSystem& cs,
                                size_t alignment) {
   return ::operator new (bytes, cs, alignment);
 }
+
+// FIXME: Perhaps we should store the Constraint -> PreparedOverload mapping
+// in a SolverStep or something? Mutating Constraint feels wrong.
+
+void Constraint::setPreparedOverload(PreparedOverload *preparedOverload) {
+  ASSERT(Kind == ConstraintKind::BindOverload);
+
+  // We can only set a prepared overload at most once in normal
+  // type checking, and then once in salvage.
+  ASSERT(!Overload.Prepared ||
+         (!Overload.Prepared->wasForDiagnostics() &&
+          preparedOverload->wasForDiagnostics()));
+
+  Overload.Prepared = preparedOverload;
+}

lib/Sema/ConstraintSystem.cpp

@@ -2098,8 +2098,6 @@ SolutionResult ConstraintSystem::salvage() {
     llvm::errs() << "---Attempting to salvage and emit diagnostics---\n";
   }
 
-  setPhase(ConstraintSystemPhase::Diagnostics);
-
   // Attempt to solve again, capturing all states that come from our attempts to
   // select overloads or bind type variables.
   //
@@ -4797,10 +4795,6 @@ void SyntacticElementTargetKey::dump(raw_ostream &OS) const {
 /// This is guaranteed to always emit an error message.
 ///
 void ConstraintSystem::diagnoseFailureFor(SyntacticElementTarget target) {
-  setPhase(ConstraintSystemPhase::Diagnostics);
-
-  SWIFT_DEFER { setPhase(ConstraintSystemPhase::Finalization); };
-
   auto &DE = getASTContext().Diags;
 
   // If constraint system is in invalid state always produce

lib/Sema/TypeOfReference.cpp

@@ -2811,7 +2811,8 @@ void ConstraintSystem::replayChanges(
       addConstraint(ConstraintKind::Bind,
                     change.Bind.FirstType,
                     change.Bind.SecondType,
-                    locator, /*isFavored=*/false);
+                    locator,
+                    /*isFavored=*/false);
       break;
 
     case PreparedOverload::Change::OpenedTypes: {
@@ -2886,11 +2887,12 @@ ConstraintSystem::prepareOverloadImpl(OverloadChoice choice,
 
 PreparedOverload *ConstraintSystem::prepareOverload(OverloadChoice choice,
                                                     DeclContext *useDC,
-                                                    ConstraintLocator *locator) {
+                                                    ConstraintLocator *locator,
+                                                    bool forDiagnostics) {
   ASSERT(!PreparingOverload);
   PreparingOverload = true;
 
-  PreparedOverloadBuilder builder;
+  PreparedOverloadBuilder builder(locator);
   Type openedType;
   Type thrownErrorType;
   std::tie(openedType, thrownErrorType) = prepareOverloadImpl(
@@ -2902,7 +2904,8 @@ PreparedOverload *ConstraintSystem::prepareOverload(OverloadChoice choice,
   auto size = PreparedOverload::totalSizeToAlloc<PreparedOverload::Change>(count);
   auto mem = Allocator.Allocate(size, alignof(PreparedOverload));
 
-  return new (mem) PreparedOverload(openedType, thrownErrorType, builder.Changes);
+  return new (mem) PreparedOverload(openedType, thrownErrorType, builder.Changes,
+                                    forDiagnostics);
 }
 
 void ConstraintSystem::resolveOverload(OverloadChoice choice, DeclContext *useDC,