[Constraint solver] More type map updates.

lib/Sema/CSApply.cpp

@@ -1303,7 +1303,7 @@ namespace {
       // Apply a key path if we have one.
       if (choice.getKind() == OverloadChoiceKind::KeyPathApplication) {
         // The index argument should be (keyPath: KeyPath<Root, Value>).
-        auto keyPathTTy = index->getType()->castTo<TupleType>()
+        auto keyPathTTy = cs.getType(index)->castTo<TupleType>()
           ->getElementType(0);
 
         Type valueTy;
@@ -1722,7 +1722,7 @@ namespace {
 
       // Form the arguments.
       Expr *args[2] = {
-        cs.cacheType(object),
+        object,
         cs.cacheType(
             new (tc.Context) DotSelfExpr(
                 TypeExpr::createImplicitHack(object->getLoc(),
@@ -2816,8 +2816,10 @@ namespace {
                             /*Implicit=*/true,
                             argType);
 
-      cs.cacheExprTypes(typeRef);
+      cs.cacheExprTypes(arg);
+
       cs.setExprTypes(typeRef);
+      cs.setExprTypes(arg);
 
       Expr *result = tc.callWitness(typeRef, dc, arrayProto, *conformance,
                                     name, arg, diag::array_protocol_broken);
@@ -2896,8 +2898,10 @@ namespace {
                               /*Implicit=*/false,
                               argType);
 
-      cs.cacheExprTypes(typeRef);
+      cs.cacheExprTypes(arg);
+
       cs.setExprTypes(typeRef);
+      cs.setExprTypes(arg);
 
       Expr *result = tc.callWitness(typeRef, dc, dictionaryProto,
                                     *conformance, name, arg,
@@ -3611,7 +3615,7 @@ namespace {
       // already.
       Expr *simplified = simplifyExprType(expr);
       if (!SuppressDiagnostics
-          && !simplified->getType()->is<UnresolvedType>()) {
+          && !cs.getType(simplified)->is<UnresolvedType>()) {
         cs.TC.diagnose(simplified->getLoc(), diag::pattern_in_expr,
                        expr->getSubPattern()->getKind());
       }

lib/Sema/CSDiag.cpp

@@ -3353,8 +3353,6 @@ Expr *FailureDiagnosis::typeCheckChildIndependently(
                                              convertTypePurpose, TCEOptions,
                                              listener, CS);
 
-  CS->cacheExprTypes(subExpr);
-
   // This is a terrible hack to get around the fact that typeCheckExpression()
   // might change subExpr to point to a new OpenExistentialExpr. In that case,
   // since the caller passed subExpr by value here, they would be left
@@ -3368,9 +3366,11 @@ Expr *FailureDiagnosis::typeCheckChildIndependently(
   if (hadError)
     return nullptr;
 
+  CS->cacheExprTypes(subExpr);
+
   // If we type checked the result but failed to get a usable output from it,
   // just pretend as though nothing happened.
-  if (subExpr->getType()->is<ErrorType>()) {
+  if (CS->getType(subExpr)->is<ErrorType>()) {
     subExpr = preCheckedExpr;
     SavedTypeData.restore();
   }

lib/Sema/CSGen.cpp

@@ -1710,7 +1710,7 @@ namespace {
         unsigned index = 0;
         for (auto element : expr->getElements()) {
           CS.addConstraint(ConstraintKind::Conversion,
-                           element->getType(),
+                           CS.getType(element),
                            contextualArrayElementType,
                            CS.getConstraintLocator(expr,
                                                    LocatorPathElt::
@@ -1738,7 +1738,7 @@ namespace {
       unsigned index = 0;
       for (auto element : expr->getElements()) {
         CS.addConstraint(ConstraintKind::Conversion,
-                         element->getType(),
+                         CS.getType(element),
                          arrayElementTy,
                          CS.getConstraintLocator(
                            expr,
@@ -2002,7 +2002,7 @@ namespace {
         // This will avoid exponential typecheck behavior in the case of nested
         // array and dictionary literals.
         Type ty = haveBoundCollectionLiteral ?
-                    boundExpr->getType() :
+                    CS.getType(boundExpr) :
                     CS.createTypeVariable(CS.getConstraintLocator(locator),
                                           TVO_CanBindToInOut);
 
@@ -2060,7 +2060,7 @@ namespace {
 
     Type visitCaptureListExpr(CaptureListExpr *expr) {
       // The type of the capture list is just the type of its closure.
-      return expr->getClosureBody()->getType();
+      return CS.getType(expr->getClosureBody());
     }
 
     /// \brief Walk a closure body to determine if it's possible for
@@ -2773,7 +2773,7 @@ namespace {
 
     Type visitKeyPathExpr(KeyPathExpr *E) {
       if (E->isObjC())
-        return E->getObjCStringLiteralExpr()->getType();
+        return CS.getType(E->getObjCStringLiteralExpr());
 
       auto kpDecl = CS.getASTContext().getKeyPathDecl();
 

lib/Sema/ConstraintSystem.h

@@ -1196,14 +1196,6 @@ class ConstraintSystem {
     SetExprTypes(Expr *expr, ConstraintSystem &cs, bool excludeRoot)
         : RootExpr(expr), CS(cs), ExcludeRoot(excludeRoot) {}
 
-    std::pair<bool, Expr *> walkToExprPre(Expr *expr) override {
-      if (auto *closure = dyn_cast<ClosureExpr>(expr))
-        if (!closure->hasSingleExpressionBody())
-          return { false, closure };
-
-      return { true, expr };
-    }
-
     Expr *walkToExprPost(Expr *expr) override {
       if (ExcludeRoot && expr == RootExpr)
         return expr;

lib/Sema/TypeCheckConstraints.cpp

@@ -2181,7 +2181,7 @@ bool TypeChecker::typeCheckBinding(Pattern *&pattern, Expr *&initializer,
         return nullptr;
       }
 
-      assert(expr->getType()->isEqual(InitType));
+      assert(solution.getConstraintSystem().getType(expr)->isEqual(InitType));
 
       initializer = expr;
       return expr;