[Constraint solver] Update type map code for recursion through salvag…

lib/Sema/CSDiag.cpp

@@ -3282,14 +3282,23 @@ Expr *FailureDiagnosis::typeCheckChildIndependently(
   // expression (which may lead to infinite recursion).  If the client is
   // telling us that it knows what it is doing, then believe it.
   if (!options.contains(TCC_ForceRecheck)) {
-    if (Expr *res = CS->TC.isExprBeingDiagnosed(subExpr)) {
+    if (CS->TC.isExprBeingDiagnosed(subExpr)) {
+      auto exprAndCS = CS->TC.getExprBeingDiagnosed(subExpr);
+      auto *savedExpr = exprAndCS.first;
+      if (subExpr == savedExpr)
+        return subExpr;
+
+      auto *oldCS = exprAndCS.second;
+
       // The types on the result might have already been cached into
       // another CS, but likely not this one.
-      CS->cacheExprTypes(res);
-      return res;
+      if (oldCS != CS)
+        CS->transferExprTypes(oldCS, savedExpr);
+
+      return savedExpr;
     }
 
-    CS->TC.addExprForDiagnosis(subExpr, subExpr);
+    CS->TC.addExprForDiagnosis(subExpr, std::make_pair(subExpr, CS));
   }
 
   // Validate contextual type before trying to use it.
@@ -3377,8 +3386,8 @@ Expr *FailureDiagnosis::typeCheckChildIndependently(
     SavedTypeData.restore();
   }
 
-  CS->TC.addExprForDiagnosis(preCheckedExpr, subExpr);
   CS->cacheExprTypes(subExpr);
+  CS->TC.addExprForDiagnosis(preCheckedExpr, std::make_pair(subExpr, CS));
 
   return subExpr;
 }

lib/Sema/ConstraintSystem.h

@@ -1287,6 +1287,30 @@ class ConstraintSystem {
     bool walkToDeclPre(Decl *decl) override { return false; }
   };
 
+  class TransferExprTypes : public ASTWalker {
+    ConstraintSystem &toCS;
+    ConstraintSystem &fromCS;
+
+  public:
+    TransferExprTypes(ConstraintSystem &toCS, ConstraintSystem &fromCS)
+        : toCS(toCS), fromCS(fromCS) {}
+
+    Expr *walkToExprPost(Expr *expr) override {
+      if (fromCS.hasType(expr))
+        toCS.setType(expr, fromCS.getType(expr));
+
+      return expr;
+    }
+
+    /// \brief Ignore statements.
+    std::pair<bool, Stmt *> walkToStmtPre(Stmt *stmt) override {
+      return { false, stmt };
+    }
+
+    /// \brief Ignore declarations.
+    bool walkToDeclPre(Decl *decl) override { return false; }
+  };
+
 public:
 
   void setExprTypes(Expr *expr) {
@@ -1312,6 +1336,10 @@ class ConstraintSystem {
     expr->walk(CacheExprTypes(expr, *this, excludeRoot));
   }
 
+  void transferExprTypes(ConstraintSystem *oldCS, Expr *expr) {
+    expr->walk(TransferExprTypes(*this, *oldCS));
+  }
+
   /// \brief The current solver state.
   ///
   /// This will be non-null when we're actively solving the constraint

lib/Sema/TypeChecker.h

@@ -55,6 +55,12 @@ namespace constraints {
 typedef llvm::DenseMap<SubstitutableType *,
                        SmallVector<ProtocolConformance *, 2>> ConformanceMap;
 
+/// \brief Used for recursive lookups into an expr that is already
+/// being type-checked and the constraint system in which its type is
+/// stored.
+typedef std::pair<Expr *,
+                  constraints::ConstraintSystem *> ExprAndConstraintSystem;
+
 /// Special-case type checking semantics for certain declarations.
 enum class DeclTypeCheckingSemantics {
   /// A normal declaration.
@@ -816,7 +822,7 @@ class TypeChecker final : public LazyResolver {
   llvm::DenseSet<CanType> CIntegerTypes;
 
   /// The set of expressions currently being analyzed for failures.
-  llvm::DenseMap<Expr*, Expr*> DiagnosedExprs;
+  llvm::DenseMap<Expr*, ExprAndConstraintSystem> DiagnosedExprs;
 
   ModuleDecl *StdlibModule = nullptr;
 
@@ -2307,10 +2313,13 @@ class TypeChecker final : public LazyResolver {
   void checkInitializerErrorHandling(Initializer *I, Expr *E);
   void checkEnumElementErrorHandling(EnumElementDecl *D);
 
-  void addExprForDiagnosis(Expr *E1, Expr *Result) {
+  void addExprForDiagnosis(Expr *E1, ExprAndConstraintSystem Result) {
     DiagnosedExprs[E1] = Result;
   }
-  Expr *isExprBeingDiagnosed(Expr *E) {
+  bool isExprBeingDiagnosed(Expr *E) {
+    return DiagnosedExprs.count(E);
+  }
+  ExprAndConstraintSystem getExprBeingDiagnosed(Expr *E) {
     return DiagnosedExprs[E];
   }