[Constraint system] Update type map code for recursion through salvage().

When we track which expressions we're already in the middle of type
checking, we need to ensure that we also track the constraint system
which has the types for that expression, and then transfer those types
into our current constraint system so that we do not fail to look them up.

Author: rudkx

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];
   }