[Constraint solver] Capture/roll back type mappings generated while solving.

When we perform constraint generation while solving, capture the
type mappings we generate as part of the solver scope and solutions,
rolling them back and replaying them as necessary. Otherwise, we’ll
end up with uses of stale type variables, expression/parameter/type-loc
types set twice, etc.

Fixes rdar://problem/50390449 and rdar://problem/50150314.

Author: DougGregor

lib/Sema/CSApply.cpp

@@ -7487,6 +7487,11 @@ Expr *ConstraintSystem::applySolution(Solution &solution, Expr *expr,
                                       Type convertType,
                                       bool discardedExpr,
                                       bool skipClosures) {
+  // Add the node types back.
+  for (auto &nodeType : solution.addedNodeTypes) {
+    setType(nodeType.first, nodeType.second);
+  }
+
   // If any fixes needed to be applied to arrive at this solution, resolve
   // them to specific expressions.
   if (!solution.Fixes.empty()) {

lib/Sema/CSSolver.cpp

@@ -169,10 +169,19 @@ Solution ConstraintSystem::finalize() {
   solution.DefaultedConstraints.insert(DefaultedConstraints.begin(),
                                        DefaultedConstraints.end());
 
+  for (auto &nodeType : addedNodeTypes) {
+    solution.addedNodeTypes.push_back(nodeType);
+  }
+
   for (auto &e : CheckedConformances)
     solution.Conformances.push_back({e.first, e.second});
 
   for (const auto &transformed : builderTransformedClosures) {
+    auto known =
+        solution.builderTransformedClosures.find(std::get<0>(transformed));
+    if (known != solution.builderTransformedClosures.end()) {
+      assert(known->second.second == std::get<2>(transformed));
+    }
     solution.builderTransformedClosures.insert(
       std::make_pair(std::get<0>(transformed),
                      std::make_pair(std::get<1>(transformed),
@@ -240,6 +249,11 @@ void ConstraintSystem::applySolution(const Solution &solution) {
                               solution.DefaultedConstraints.begin(),
                               solution.DefaultedConstraints.end());
 
+  // Add the node types back.
+  for (auto &nodeType : solution.addedNodeTypes) {
+    setType(nodeType.first, nodeType.second);
+  }
+
   // Register the conformances checked along the way to arrive to solution.
   for (auto &conformance : solution.Conformances)
     CheckedConformances.push_back(conformance);
@@ -442,6 +456,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numOpenedTypes = cs.OpenedTypes.size();
   numOpenedExistentialTypes = cs.OpenedExistentialTypes.size();
   numDefaultedConstraints = cs.DefaultedConstraints.size();
+  numAddedNodeTypes = cs.addedNodeTypes.size();
   numCheckedConformances = cs.CheckedConformances.size();
   numMissingMembers = cs.MissingMembers.size();
   numDisabledConstraints = cs.solverState->getNumDisabledConstraints();
@@ -495,6 +510,12 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Remove any defaulted type variables.
   truncate(cs.DefaultedConstraints, numDefaultedConstraints);
 
+  // Remove any node types we registered.
+  for (unsigned i : range(numAddedNodeTypes, cs.addedNodeTypes.size())) {
+    cs.eraseType(cs.addedNodeTypes[i].first);
+  }
+  truncate(cs.addedNodeTypes, numAddedNodeTypes);
+
   // Remove any conformances checked along the current path.
   truncate(cs.CheckedConformances, numCheckedConformances);
 

lib/Sema/ConstraintSystem.cpp

@@ -3261,6 +3261,12 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::applyFunctionBuilder(
   assert(transformedType && "Missing type");
 
   // Record the transformation.
+  assert(std::find_if(builderTransformedClosures.begin(),
+                      builderTransformedClosures.end(),
+                      [&](const std::tuple<ClosureExpr *, Type, Expr *> &elt) {
+                        return std::get<0>(elt) == closure;
+                      }) == builderTransformedClosures.end() &&
+         "already transformed this closure along this path!?!");
   builderTransformedClosures.push_back(
     std::make_tuple(closure, builderType, singleExpr));
 

lib/Sema/ConstraintSystem.h

@@ -570,6 +570,9 @@ struct Score {
 
 };
 
+/// An AST node that can gain type information while solving.
+using TypedNode = llvm::PointerUnion3<Expr *, TypeLoc *, ParamDecl *>;
+
 /// Display a score.
 llvm::raw_ostream &operator<<(llvm::raw_ostream &out, const Score &score);
 
@@ -642,6 +645,9 @@ class Solution {
   /// The locators of \c Defaultable constraints whose defaults were used.
   llvm::SmallPtrSet<ConstraintLocator *, 2> DefaultedConstraints;
 
+  /// The node -> type mappings introduced by this solution.
+  llvm::SmallVector<std::pair<TypedNode, Type>, 8> addedNodeTypes;
+
   std::vector<std::pair<ConstraintLocator *, ProtocolConformanceRef>>
       Conformances;
 
@@ -1106,6 +1112,9 @@ class ConstraintSystem {
   SmallVector<std::pair<ConstraintLocator *, OpenedArchetypeType *>, 4>
     OpenedExistentialTypes;
 
+  /// The node -> type mappings introduced by generating constraints.
+  llvm::SmallVector<std::pair<TypedNode, Type>, 8> addedNodeTypes;
+
   std::vector<std::pair<ConstraintLocator *, ProtocolConformanceRef>>
       CheckedConformances;
 
@@ -1586,6 +1595,8 @@ class ConstraintSystem {
     /// The length of \c DefaultedConstraints.
     unsigned numDefaultedConstraints;
 
+    unsigned numAddedNodeTypes;
+
     unsigned numCheckedConformances;
 
     unsigned numMissingMembers;
@@ -1736,33 +1747,56 @@ class ConstraintSystem {
     this->FavoredTypes[E] = T;
   }
 
+  /// Set the type in our type map for the given node.
+  ///
+  /// The side tables are used through the expression type checker to avoid mutating nodes until
+  /// we know we have successfully type-checked them.
+  void setType(TypedNode node, Type type) {
+    assert(!node.isNull() && "Cannot set type information on null node");
+    assert(type && "Expected non-null type");
+
+    // Record the type.
+    if (auto expr = node.dyn_cast<Expr *>()) {
+      ExprTypes[expr] = type.getPointer();
+    } else if (auto typeLoc = node.dyn_cast<TypeLoc *>()) {
+      TypeLocTypes[typeLoc] = type.getPointer();
+    } else {
+      auto param = node.get<ParamDecl *>();
+      ParamTypes[param] = type.getPointer();
+    }
+
+    // Record the fact that we ascribed a type to this node.
+    if (solverState && solverState->depth > 0) {
+      addedNodeTypes.push_back({node, type});
+    }
+  }
+
   /// Set the type in our type map for a given expression. The side
   /// map is used throughout the expression type checker in order to
   /// avoid mutating expressions until we know we have successfully
   /// type-checked them.
   void setType(Expr *E, Type T) {
-    assert(E != nullptr && "Expected non-null expression!");
-    assert(T && "Expected non-null type!");
-
-    // FIXME: We sometimes set the type and then later set it to a
-    //        value that is slightly different, e.g. not an lvalue.
-    // assert((ExprTypes.find(E) == ExprTypes.end() ||
-    //         ExprTypes.find(E)->second->isEqual(T) ||
-    //         ExprTypes.find(E)->second->hasTypeVariable()) &&
-    //        "Expected type to be invariant!");
-
-    ExprTypes[E] = T.getPointer();
+    setType(TypedNode(E), T);
   }
 
   void setType(TypeLoc &L, Type T) {
-    assert(T && "Expected non-null type!");
-    TypeLocTypes[&L] = T.getPointer();
+    setType(TypedNode(&L), T);
   }
 
   void setType(ParamDecl *P, Type T) {
-    assert(P && "Expected non-null parameter!");
-    assert(T && "Expected non-null type!");
-    ParamTypes[P] = T.getPointer();
+    setType(TypedNode(P), T);
+  }
+
+  /// Erase the type for the given node.
+  void eraseType(TypedNode node) {
+    if (auto expr = node.dyn_cast<Expr *>()) {
+      ExprTypes.erase(expr);
+    } else if (auto typeLoc = node.dyn_cast<TypeLoc *>()) {
+      TypeLocTypes.erase(typeLoc);
+    } else {
+      auto param = node.get<ParamDecl *>();
+      ParamTypes.erase(param);
+    }
   }
 
   void setType(KeyPathExpr *KP, unsigned I, Type T) {

test/Constraints/function_builder_diags.swift

@@ -114,13 +114,8 @@ func testOverloading(name: String) {
 
   let _: A = a1
 
-#if false
-  // FIXME: Using the overloads causes a crash, which is likely related
-  // to us generating the constraints from the child expressions multiple
-  // times.
-  let b1 = overloadedTuplify(true) { b in
+  _ = overloadedTuplify(true) { b in
     b ? "Hello, \(name)" : "Goodbye"
     42
   }
-#endif
 }