Sema: Don't reuse PreparedOverload from normal type checking in salvage()

We want to re-prepare our overloads, since now they will contain fixes.

Author: slavapestov

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

@@ -4876,7 +4876,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

@@ -160,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());
   }
@@ -183,6 +187,10 @@ class PreparedOverload final :
     return ThrownErrorType;
   }
 
+  bool wasForDiagnostics() const {
+    return ForDiagnostics;
+  }
+
   ArrayRef<Change> getChanges() const { return getTrailingObjects(Count); }
 };
 

lib/Sema/CSSimplify.cpp

@@ -16837,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/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/TypeOfReference.cpp

@@ -2887,7 +2887,8 @@ ConstraintSystem::prepareOverloadImpl(OverloadChoice choice,
 
 PreparedOverload *ConstraintSystem::prepareOverload(OverloadChoice choice,
                                                     DeclContext *useDC,
-                                                    ConstraintLocator *locator) {
+                                                    ConstraintLocator *locator,
+                                                    bool forDiagnostics) {
   ASSERT(!PreparingOverload);
   PreparingOverload = true;
 
@@ -2903,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,