Merge pull request swiftlang#29801 from DougGregor/constraint-solver-rewrite-target

[Constraint system] Factor out application to a SolutionApplicationTarget

Author: DougGregor

lib/Sema/BuilderTransform.cpp

@@ -627,8 +627,9 @@ class BuilderClosureRewriter
   const Solution &solution;
   DeclContext *dc;
   AppliedBuilderTransform builderTransform;
-  std::function<Expr *(Expr *)> rewriteExpr;
-  std::function<Expr *(Expr *, Type, ConstraintLocator *)> coerceToType;
+  std::function<
+      Optional<SolutionApplicationTarget> (SolutionApplicationTarget)>
+        rewriteTarget;
 
   /// Retrieve the temporary variable that will be used to capture the
   /// value of the given expression.
@@ -648,6 +649,17 @@ class BuilderClosureRewriter
     return recorded;
   }
 
+  /// Rewrite an expression without any particularly special context.
+  Expr *rewriteExpr(Expr *expr) {
+    auto result = rewriteTarget(
+      SolutionApplicationTarget(expr, dc, CTP_Unused, Type(),
+                                /*isDiscarded=*/false));
+    if (result)
+      return result->getAsExpr();
+
+    return nullptr;
+  }
+
 public:
   /// Retrieve information about a captured statement.
   std::pair<VarDecl *, llvm::TinyPtrVector<Expr *>>
@@ -673,19 +685,21 @@ class BuilderClosureRewriter
   ASTNode initializeTarget(FunctionBuilderTarget target) {
     assert(target.captured.second.size() == 1);
     auto capturedExpr = target.captured.second.front();
-    auto finalCapturedExpr = rewriteExpr(capturedExpr);
     SourceLoc implicitLoc = capturedExpr->getEndLoc();
     switch (target.kind) {
     case FunctionBuilderTarget::ReturnValue: {
       // Return the expression.
-      ConstraintSystem &cs = solution.getConstraintSystem();
       Type bodyResultType =
           solution.simplifyType(builderTransform.bodyResultType);
-      finalCapturedExpr = coerceToType(
-          finalCapturedExpr,
-          bodyResultType,
-          cs.getConstraintLocator(capturedExpr));
-      return new (ctx) ReturnStmt(implicitLoc, finalCapturedExpr);
+
+      SolutionApplicationTarget returnTarget(
+          capturedExpr, dc, CTP_ReturnStmt, bodyResultType,
+          /*isDiscarded=*/false);
+      Expr *resultExpr = nullptr;
+      if (auto resultTarget = rewriteTarget(returnTarget))
+        resultExpr = resultTarget->getAsExpr();
+
+      return new (ctx) ReturnStmt(implicitLoc, resultExpr);
     }
 
     case FunctionBuilderTarget::TemporaryVar: {
@@ -696,6 +710,7 @@ class BuilderClosureRewriter
       declRef->setType(LValueType::get(temporaryVar->getType()));
 
       // Load the right-hand side if needed.
+      auto finalCapturedExpr = rewriteExpr(capturedExpr);
       if (finalCapturedExpr->getType()->hasLValueType()) {
         finalCapturedExpr =
             TypeChecker::addImplicitLoadExpr(ctx, finalCapturedExpr);
@@ -734,12 +749,12 @@ class BuilderClosureRewriter
       const Solution &solution,
       DeclContext *dc,
       const AppliedBuilderTransform &builderTransform,
-      std::function<Expr *(Expr *)> rewriteExpr,
-      std::function<Expr *(Expr *, Type, ConstraintLocator *)> coerceToType
+      std::function<
+          Optional<SolutionApplicationTarget> (SolutionApplicationTarget)>
+            rewriteTarget
     ) : ctx(solution.getConstraintSystem().getASTContext()),
         solution(solution), dc(dc), builderTransform(builderTransform),
-        rewriteExpr(rewriteExpr),
-        coerceToType(coerceToType){ }
+        rewriteTarget(rewriteTarget) { }
 
   Stmt *visitBraceStmt(BraceStmt *braceStmt, FunctionBuilderTarget target,
                        Optional<FunctionBuilderTarget> innerTarget = None) {
@@ -954,9 +969,10 @@ BraceStmt *swift::applyFunctionBuilderTransform(
     AppliedBuilderTransform applied,
     BraceStmt *body,
     DeclContext *dc,
-    std::function<Expr *(Expr *)> rewriteExpr,
-    std::function<Expr *(Expr *, Type, ConstraintLocator *)> coerceToType) {
-  BuilderClosureRewriter rewriter(solution, dc, applied, rewriteExpr, coerceToType);
+    std::function<
+        Optional<SolutionApplicationTarget> (SolutionApplicationTarget)>
+          rewriteTarget) {
+  BuilderClosureRewriter rewriter(solution, dc, applied, rewriteTarget);
   auto captured = rewriter.takeCapturedStmt(body);
   return cast<BraceStmt>(
     rewriter.visitBraceStmt(

lib/Sema/CSApply.cpp

@@ -290,7 +290,7 @@ namespace {
   public:
     ConstraintSystem &cs;
     DeclContext *dc;
-    const Solution &solution;
+    Solution &solution;
     bool SuppressDiagnostics;
 
     /// Coerce the given tuple to another tuple type.
@@ -1792,7 +1792,7 @@ namespace {
     }
 
   public:
-    ExprRewriter(ConstraintSystem &cs, const Solution &solution,
+    ExprRewriter(ConstraintSystem &cs, Solution &solution,
                  bool suppressDiagnostics)
         : cs(cs), dc(cs.DC), solution(solution),
           SuppressDiagnostics(suppressDiagnostics) {}
@@ -7063,14 +7063,14 @@ namespace {
           llvm::SaveAndRestore<DeclContext *> savedDC(Rewriter.dc, closure);
           auto newBody = applyFunctionBuilderTransform(
               Rewriter.solution, *transform, closure->getBody(), closure,
-              [&](Expr *expr) {
-                Expr *result = expr->walk(*this);
-                if (result)
-                  Rewriter.solution.setExprTypes(result);
-                return result;
-              },
-              [&](Expr *expr, Type toType, ConstraintLocator *locator) {
-                return Rewriter.coerceToType(expr, toType, locator);
+              [&](SolutionApplicationTarget target) {
+                auto resultTarget = rewriteTarget(target);
+                if (resultTarget) {
+                  if (auto expr = resultTarget->getAsExpr())
+                    Rewriter.solution.setExprTypes(expr);
+                }
+
+                return resultTarget;
               });
           closure->setBody(newBody, /*isSingleExpression=*/false);
           closure->setAppliedFunctionBuilder();
@@ -7142,6 +7142,10 @@ namespace {
 
     /// Ignore declarations.
     bool walkToDeclPre(Decl *decl) override { return false; }
+
+    /// Rewrite the target, producing a new target.
+    Optional<SolutionApplicationTarget>
+    rewriteTarget(SolutionApplicationTarget target);
   };
 } // end anonymous namespace
 
@@ -7305,41 +7309,14 @@ static Optional<SolutionApplicationTarget> applySolutionToInitialization(
   return resultTarget;
 }
 
-/// Apply a given solution to the expression, producing a fully
-/// type-checked expression.
-Optional<SolutionApplicationTarget> ConstraintSystem::applySolution(
-    Solution &solution, SolutionApplicationTarget target,
-    bool performingDiagnostics) {
-  // If any fixes needed to be applied to arrive at this solution, resolve
-  // them to specific expressions.
-  if (!solution.Fixes.empty()) {
-    if (shouldSuppressDiagnostics())
-      return None;
-
-    bool diagnosedErrorsViaFixes = applySolutionFixes(solution);
-    // If all of the available fixes would result in a warning,
-    // we can go ahead and apply this solution to AST.
-    if (!llvm::all_of(solution.Fixes, [](const ConstraintFix *fix) {
-          return fix->isWarning();
-        })) {
-      // If we already diagnosed any errors via fixes, that's it.
-      if (diagnosedErrorsViaFixes)
-        return None;
-
-      // If we didn't manage to diagnose anything well, so fall back to
-      // diagnosing mining the system to construct a reasonable error message.
-      diagnoseFailureFor(target);
-      return None;
-    }
-  }
-
-  ExprRewriter rewriter(*this, solution, shouldSuppressDiagnostics());
-  ExprWalker walker(rewriter);
+Optional<SolutionApplicationTarget>
+ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
+  auto &solution = Rewriter.solution;
 
   // Apply the solution to the target.
   SolutionApplicationTarget result = target;
   if (auto expr = target.getAsExpr()) {
-    Expr *rewrittenExpr = expr->walk(walker);
+    Expr *rewrittenExpr = expr->walk(*this);
     if (!rewrittenExpr)
       return None;
 
@@ -7358,19 +7335,19 @@ Optional<SolutionApplicationTarget> ConstraintSystem::applySolution(
     auto fn = *target.getAsFunction();
 
     // Dig out the function builder transformation we applied.
-    auto transform = rewriter.getAppliedBuilderTransform(fn);
+    auto transform = Rewriter.getAppliedBuilderTransform(fn);
     assert(transform);
 
     auto newBody = applyFunctionBuilderTransform(
         solution, *transform, fn.getBody(), fn.getAsDeclContext(),
-        [&](Expr *expr) {
-          Expr *result = expr->walk(walker);
-          if (result)
-            solution.setExprTypes(result);
-          return result;
-        },
-        [&](Expr *expr, Type toType, ConstraintLocator *locator) {
-          return rewriter.coerceToType(expr, toType, locator);
+        [&](SolutionApplicationTarget target) {
+          auto resultTarget = rewriteTarget(target);
+          if (resultTarget) {
+            if (auto expr = resultTarget->getAsExpr())
+              Rewriter.solution.setExprTypes(expr);
+          }
+
+          return resultTarget;
         });
 
     if (!newBody)
@@ -7379,26 +7356,8 @@ Optional<SolutionApplicationTarget> ConstraintSystem::applySolution(
     result.setFunctionBody(newBody);
   }
 
-  // If we're re-typechecking an expression for diagnostics, don't
-  // visit closures that have non-single expression bodies.
-  if (!performingDiagnostics) {
-    bool hadError = false;
-    for (auto *closure : walker.getClosuresToTypeCheck())
-      hadError |= TypeChecker::typeCheckClosureBody(closure);
-
-    // Tap expressions too; they should or should not be
-    // type-checked under the same conditions as closure bodies.
-    for (auto tuple : walker.getTapsToTypeCheck()) {
-      auto tap = std::get<0>(tuple);
-      auto tapDC = std::get<1>(tuple);
-      hadError |= TypeChecker::typeCheckTapBody(tap, tapDC);
-    }
-
-    // If any of them failed to type check, bail.
-    if (hadError)
-      return None;
-  }
-
+  // Follow-up tasks.
+  auto &cs = solution.getConstraintSystem();
   if (auto resultExpr = result.getAsExpr()) {
     Expr *expr = target.getAsExpr();
     assert(expr && "Can't have expression result without expression target");
@@ -7410,23 +7369,24 @@ Optional<SolutionApplicationTarget> ConstraintSystem::applySolution(
     auto shouldCoerceToContextualType = [&]() {
       return convertType &&
           !target.isOptionalSomePatternInit() &&
-          !(getType(resultExpr)->isUninhabited() &&
-            getContextualTypePurpose(target.getAsExpr())
+          !(solution.getType(resultExpr)->isUninhabited() &&
+            cs.getContextualTypePurpose(target.getAsExpr())
               == CTP_ReturnSingleExpr);
     };
 
     // If we're supposed to convert the expression to some particular type,
     // do so now.
     if (shouldCoerceToContextualType()) {
-      resultExpr = rewriter.coerceToType(resultExpr,
-                                         simplifyType(convertType),
-                                         getConstraintLocator(expr));
-    } else if (getType(resultExpr)->hasLValueType() &&
+      resultExpr = Rewriter.coerceToType(resultExpr,
+                                         solution.simplifyType(convertType),
+                                         cs.getConstraintLocator(expr));
+    } else if (cs.getType(resultExpr)->hasLValueType() &&
                !target.isDiscardedExpr()) {
       // We referenced an lvalue. Load it.
-      resultExpr = rewriter.coerceToType(resultExpr,
-                                         getType(resultExpr)->getRValueType(),
-                                         getConstraintLocator(expr));
+      resultExpr = Rewriter.coerceToType(
+          resultExpr,
+          cs.getType(resultExpr)->getRValueType(),
+          cs.getConstraintLocator(expr));
     }
 
     if (!resultExpr)
@@ -7436,28 +7396,86 @@ Optional<SolutionApplicationTarget> ConstraintSystem::applySolution(
     // conversion.
     if (FunctionType *autoclosureParamType =
             target.getAsAutoclosureParamType()) {
-      resultExpr = buildAutoClosureExpr(resultExpr, autoclosureParamType);
+      resultExpr = cs.buildAutoClosureExpr(resultExpr, autoclosureParamType);
     }
 
     solution.setExprTypes(resultExpr);
     result.setExpr(resultExpr);
 
-    if (Context.TypeCheckerOpts.DebugConstraintSolver) {
-      auto &log = Context.TypeCheckerDebug->getStream();
+    auto &ctx = cs.getASTContext();
+    if (ctx.TypeCheckerOpts.DebugConstraintSolver) {
+      auto &log = ctx.TypeCheckerDebug->getStream();
       log << "---Type-checked expression---\n";
       resultExpr->dump(log);
       log << "\n";
     }
   }
 
+  return result;
+}
+
+/// Apply a given solution to the expression, producing a fully
+/// type-checked expression.
+Optional<SolutionApplicationTarget> ConstraintSystem::applySolution(
+    Solution &solution, SolutionApplicationTarget target,
+    bool performingDiagnostics) {
+  // If any fixes needed to be applied to arrive at this solution, resolve
+  // them to specific expressions.
+  if (!solution.Fixes.empty()) {
+    if (shouldSuppressDiagnostics())
+      return None;
+
+    bool diagnosedErrorsViaFixes = applySolutionFixes(solution);
+    // If all of the available fixes would result in a warning,
+    // we can go ahead and apply this solution to AST.
+    if (!llvm::all_of(solution.Fixes, [](const ConstraintFix *fix) {
+          return fix->isWarning();
+        })) {
+      // If we already diagnosed any errors via fixes, that's it.
+      if (diagnosedErrorsViaFixes)
+        return None;
+
+      // If we didn't manage to diagnose anything well, so fall back to
+      // diagnosing mining the system to construct a reasonable error message.
+      diagnoseFailureFor(target);
+      return None;
+    }
+  }
+
+  ExprRewriter rewriter(*this, solution, shouldSuppressDiagnostics());
+  ExprWalker walker(rewriter);
+  auto resultTarget = walker.rewriteTarget(target);
+  if (!resultTarget)
+    return None;
+
+  // If we're re-typechecking an expression for diagnostics, don't
+  // visit closures that have non-single expression bodies.
+  if (!performingDiagnostics) {
+    bool hadError = false;
+    for (auto *closure : walker.getClosuresToTypeCheck())
+      hadError |= TypeChecker::typeCheckClosureBody(closure);
+
+    // Tap expressions too; they should or should not be
+    // type-checked under the same conditions as closure bodies.
+    for (auto tuple : walker.getTapsToTypeCheck()) {
+      auto tap = std::get<0>(tuple);
+      auto tapDC = std::get<1>(tuple);
+      hadError |= TypeChecker::typeCheckTapBody(tap, tapDC);
+    }
+
+    // If any of them failed to type check, bail.
+    if (hadError)
+      return None;
+  }
+
   rewriter.finalize();
 
-  return result;
+  return resultTarget;
 }
 
 Expr *Solution::coerceToType(Expr *expr, Type toType,
                              ConstraintLocator *locator,
-                             Optional<Pattern*> typeFromPattern) const {
+                             Optional<Pattern*> typeFromPattern) {
   auto &cs = getConstraintSystem();
   ExprRewriter rewriter(cs, *this, /*suppressDiagnostics=*/false);
   Expr *result = rewriter.coerceToType(expr, toType, locator, typeFromPattern);