[Constraint system] Expand SolutionApplicationTarget to StmtConditions.

Handle StmtCondition as part of SolutionApplicationTarget, so we can
generate constraints from it and rewrite directly as part of a solution,
rather than open-coding the operation in the function builder transform.

Author: DougGregor

lib/Sema/BuilderTransform.cpp

@@ -920,38 +920,9 @@ class BuilderClosureRewriter
 
   Stmt *visitIfStmt(IfStmt *ifStmt, FunctionBuilderTarget target) {
     // Rewrite the condition.
-    auto condition = ifStmt->getCond();
-    for (auto &condElement : condition) {
-      switch (condElement.getKind()) {
-      case StmtConditionElement::CK_Availability:
-        continue;
-
-      case StmtConditionElement::CK_Boolean: {
-        auto condExpr = condElement.getBoolean();
-        auto finalCondExpr = rewriteExpr(condExpr);
-
-        // Load the condition if needed.
-        if (finalCondExpr->getType()->hasLValueType()) {
-          finalCondExpr = TypeChecker::addImplicitLoadExpr(ctx, finalCondExpr);
-        }
-
-        condElement.setBoolean(finalCondExpr);
-        continue;
-      }
-
-      case StmtConditionElement::CK_PatternBinding: {
-        ConstraintSystem &cs = solution.getConstraintSystem();
-        auto target = *cs.getStmtConditionTarget(&condElement);
-        auto resolvedTarget = rewriteTarget(target);
-        if (resolvedTarget) {
-          condElement.setInitializer(resolvedTarget->getAsExpr());
-          condElement.setPattern(resolvedTarget->getInitializationPattern());
-        }
-        continue;
-      }
-      }
-    }
-    ifStmt->setCond(condition);
+    if (auto condition = rewriteTarget(
+            SolutionApplicationTarget(ifStmt->getCond(), dc)))
+      ifStmt->setCond(*condition->getAsStmtCondition());
 
     assert(target.kind == FunctionBuilderTarget::TemporaryVar);
     auto temporaryVar = target.captured.first;

lib/Sema/CSApply.cpp

@@ -7092,6 +7092,48 @@ bool swift::exprNeedsParensAfterAddingNilCoalescing(DeclContext *DC,
 }
 
 namespace {
+  class SetExprTypes : public ASTWalker {
+    const Solution &solution;
+
+  public:
+    explicit SetExprTypes(const Solution &solution)
+        : solution(solution) {}
+
+    Expr *walkToExprPost(Expr *expr) override {
+      auto &cs = solution.getConstraintSystem();
+      auto exprType = cs.getType(expr);
+      exprType = solution.simplifyType(exprType);
+      // assert((!expr->getType() || expr->getType()->isEqual(exprType)) &&
+      //       "Mismatched types!");
+      assert(!exprType->hasTypeVariable() &&
+             "Should not write type variable into expression!");
+      expr->setType(exprType);
+
+      if (auto kp = dyn_cast<KeyPathExpr>(expr)) {
+        for (auto i : indices(kp->getComponents())) {
+          Type componentType;
+          if (cs.hasType(kp, i)) {
+            componentType = solution.simplifyType(cs.getType(kp, i));
+            assert(!componentType->hasTypeVariable() &&
+                   "Should not write type variable into key-path component");
+          }
+
+          kp->getMutableComponents()[i].setComponentType(componentType);
+        }
+      }
+
+      return expr;
+    }
+
+    /// Ignore statements.
+    std::pair<bool, Stmt *> walkToStmtPre(Stmt *stmt) override {
+      return { false, stmt };
+    }
+
+    /// Ignore declarations.
+    bool walkToDeclPre(Decl *decl) override { return false; }
+  };
+
   class ExprWalker : public ASTWalker {
     ExprRewriter &Rewriter;
     SmallVector<ClosureExpr *, 4> ClosuresToTypeCheck;
@@ -7129,8 +7171,13 @@ namespace {
               [&](SolutionApplicationTarget target) {
                 auto resultTarget = rewriteTarget(target);
                 if (resultTarget) {
-                  if (auto expr = resultTarget->getAsExpr())
+
+                  if (auto expr = resultTarget->getAsExpr()) {
                     Rewriter.solution.setExprTypes(expr);
+                  } else if (auto stmtCondition =
+                                 resultTarget->getAsStmtCondition()) {
+
+                  }
                 }
 
                 return resultTarget;
@@ -7394,6 +7441,43 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
     } else {
       result.setExpr(rewrittenExpr);
     }
+  } else if (auto stmtCondition = target.getAsStmtCondition()) {
+    for (auto &condElement : *stmtCondition) {
+      switch (condElement.getKind()) {
+      case StmtConditionElement::CK_Availability:
+        continue;
+
+      case StmtConditionElement::CK_Boolean: {
+        auto condExpr = condElement.getBoolean();
+        auto finalCondExpr = condExpr->walk(*this);
+        if (!finalCondExpr)
+          return None;
+
+        // Load the condition if needed.
+        if (finalCondExpr->getType()->hasLValueType()) {
+          ASTContext &ctx = solution.getConstraintSystem().getASTContext();
+          finalCondExpr = TypeChecker::addImplicitLoadExpr(ctx, finalCondExpr);
+        }
+
+        condElement.setBoolean(finalCondExpr);
+        continue;
+      }
+
+      case StmtConditionElement::CK_PatternBinding: {
+        ConstraintSystem &cs = solution.getConstraintSystem();
+        auto target = *cs.getStmtConditionTarget(&condElement);
+        auto resolvedTarget = rewriteTarget(target);
+        if (!resolvedTarget)
+          return None;
+
+        condElement.setInitializer(resolvedTarget->getAsExpr());
+        condElement.setPattern(resolvedTarget->getInitializationPattern());
+        continue;
+      }
+      }
+    }
+
+    return target;
   } else {
     auto fn = *target.getAsFunction();
 
@@ -7406,8 +7490,8 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
         [&](SolutionApplicationTarget target) {
           auto resultTarget = rewriteTarget(target);
           if (resultTarget) {
-            if (auto expr = resultTarget->getAsExpr())
-              Rewriter.solution.setExprTypes(expr);
+            SetExprTypes typeSetter(solution);
+            resultTarget->walk(typeSetter);
           }
 
           return resultTarget;
@@ -7546,50 +7630,6 @@ Expr *Solution::coerceToType(Expr *expr, Type toType,
   return result;
 }
 
-namespace {
-class SetExprTypes : public ASTWalker {
-  const Solution &solution;
-
-public:
-  explicit SetExprTypes(const Solution &solution)
-      : solution(solution) {}
-
-  Expr *walkToExprPost(Expr *expr) override {
-    auto &cs = solution.getConstraintSystem();
-    auto exprType = cs.getType(expr);
-    exprType = solution.simplifyType(exprType);
-    // assert((!expr->getType() || expr->getType()->isEqual(exprType)) &&
-    //       "Mismatched types!");
-    assert(!exprType->hasTypeVariable() &&
-           "Should not write type variable into expression!");
-    expr->setType(exprType);
-
-    if (auto kp = dyn_cast<KeyPathExpr>(expr)) {
-      for (auto i : indices(kp->getComponents())) {
-        Type componentType;
-        if (cs.hasType(kp, i)) {
-          componentType = solution.simplifyType(cs.getType(kp, i));
-          assert(!componentType->hasTypeVariable() &&
-                 "Should not write type variable into key-path component");
-        }
-
-        kp->getMutableComponents()[i].setComponentType(componentType);
-      }
-    }
-
-    return expr;
-  }
-
-  /// Ignore statements.
-  std::pair<bool, Stmt *> walkToStmtPre(Stmt *stmt) override {
-    return { false, stmt };
-  }
-
-  /// Ignore declarations.
-  bool walkToDeclPre(Decl *decl) override { return false; }
-};
-}
-
 ProtocolConformanceRef Solution::resolveConformance(
     ConstraintLocator *locator, ProtocolDecl *proto) {
   for (const auto &conformance : Conformances) {
@@ -7704,5 +7744,11 @@ SolutionApplicationTarget SolutionApplicationTarget::walk(ASTWalker &walker) {
     return SolutionApplicationTarget(
         *getAsFunction(),
         cast_or_null<BraceStmt>(getFunctionBody()->walk(walker)));
+
+  case Kind::stmtCondition:
+    for (auto &condElement : stmtCondition.stmtCondition) {
+      condElement = *condElement.walk(walker);
+    }
+    return *this;
   }
 }

lib/Sema/ConstraintSystem.h

@@ -1151,7 +1151,8 @@ struct DynamicCallableMethods {
 class SolutionApplicationTarget {
   enum class Kind {
     expression,
-    function
+    function,
+    stmtCondition
   } kind;
 
   union {
@@ -1189,6 +1190,11 @@ class SolutionApplicationTarget {
       AnyFunctionRef function;
       BraceStmt *body;
     } function;
+
+    struct {
+      StmtCondition stmtCondition;
+      DeclContext *dc;
+    } stmtCondition;
   };
 
   // If the pattern contains a single variable that has an attached
@@ -1211,6 +1217,12 @@ class SolutionApplicationTarget {
   SolutionApplicationTarget(AnyFunctionRef fn)
       : SolutionApplicationTarget(fn, fn.getBody()) { }
 
+  SolutionApplicationTarget(StmtCondition stmtCondition, DeclContext *dc) {
+    kind = Kind::stmtCondition;
+    this->stmtCondition.stmtCondition = stmtCondition;
+    this->stmtCondition.dc = dc;
+  }
+
   SolutionApplicationTarget(AnyFunctionRef fn, BraceStmt *body) {
     kind = Kind::function;
     function.function = fn;
@@ -1228,6 +1240,7 @@ class SolutionApplicationTarget {
       return expression.expression;
 
     case Kind::function:
+    case Kind::stmtCondition:
       return nullptr;
     }
   }
@@ -1239,6 +1252,9 @@ class SolutionApplicationTarget {
 
     case Kind::function:
       return function.function.getAsDeclContext();
+
+    case Kind::stmtCondition:
+      return stmtCondition.dc;
     }
   }
 
@@ -1346,13 +1362,25 @@ class SolutionApplicationTarget {
   Optional<AnyFunctionRef> getAsFunction() const {
     switch (kind) {
     case Kind::expression:
+    case Kind::stmtCondition:
       return None;
 
     case Kind::function:
       return function.function;
     }
   }
 
+  Optional<StmtCondition> getAsStmtCondition() const {
+    switch (kind) {
+    case Kind::expression:
+    case Kind::function:
+      return None;
+
+      case Kind::stmtCondition:
+      return stmtCondition.stmtCondition;
+    }
+  }
+
   BraceStmt *getFunctionBody() const {
     assert(kind == Kind::function);
     return function.body;
@@ -1371,6 +1399,10 @@ class SolutionApplicationTarget {
 
     case Kind::function:
       return function.body->getSourceRange();
+
+    case Kind::stmtCondition:
+      return SourceRange(stmtCondition.stmtCondition.front().getStartLoc(),
+                         stmtCondition.stmtCondition.back().getEndLoc());
     }
   }
 
@@ -1382,6 +1414,9 @@ class SolutionApplicationTarget {
 
     case Kind::function:
       return function.function.getLoc();
+
+    case Kind::stmtCondition:
+      return stmtCondition.stmtCondition.front().getStartLoc();
     }
   }
 
@@ -4339,6 +4374,10 @@ class ConstraintSystem {
   Optional<SolutionApplicationTarget> applySolution(
       Solution &solution, SolutionApplicationTarget target);
 
+  /// Apply the given solution to the given statement-condition.
+  Optional<StmtCondition> applySolution(
+      Solution &solution, StmtCondition condition, DeclContext *dc);
+
   /// Reorder the disjunctive clauses for a given expression to
   /// increase the likelihood that a favored constraint will be successfully
   /// resolved before any others.