Merge pull request #30065 from DougGregor/function-builders-if-case-fixes

[Constraint system] Cleanups for function builders support for if let/if case let

Author: DougGregor

lib/Sema/BuilderTransform.cpp

@@ -331,7 +331,7 @@ class BuilderClosureVisitor
       addChild(captureExpr(expr, /*oneWay=*/true, node.get<Expr *>()));
     }
 
-    if (!cs)
+    if (!cs || hadError)
       return nullptr;
 
     // Call Builder.buildBlock(... args ...)
@@ -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;
@@ -1383,6 +1354,11 @@ class PreCheckFunctionBuilderApplication : public ASTWalker {
     // Otherwise, recurse into the statement normally.
     return std::make_pair(true, S);
   }
+
+  /// Ignore patterns.
+  std::pair<bool, Pattern*> walkToPatternPre(Pattern *pat) override {
+    return { false, pat };
+  }
 };
 
 }

lib/Sema/CSApply.cpp

@@ -7093,6 +7093,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;
@@ -7395,6 +7437,45 @@ 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.
+        solution.setExprTypes(finalCondExpr);
+        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;
+
+        solution.setExprTypes(resolvedTarget->getAsExpr());
+        condElement.setInitializer(resolvedTarget->getAsExpr());
+        condElement.setPattern(resolvedTarget->getInitializationPattern());
+        continue;
+      }
+      }
+    }
+
+    return target;
   } else {
     auto fn = *target.getAsFunction();
 
@@ -7408,7 +7489,7 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
           auto resultTarget = rewriteTarget(target);
           if (resultTarget) {
             if (auto expr = resultTarget->getAsExpr())
-              Rewriter.solution.setExprTypes(expr);
+              solution.setExprTypes(expr);
           }
 
           return resultTarget;
@@ -7547,50 +7628,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) {
@@ -7705,5 +7742,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/CSGen.cpp

@@ -2197,6 +2197,18 @@ namespace {
     /// type information with fresh type variables.
     ///
     /// \param pattern The pattern.
+    ///
+    /// \param locator The locator to use for generated constraints and
+    /// type variables.
+    ///
+    /// \param externalPatternType The type imposed by the enclosing pattern,
+    /// if any. This will be non-null in cases where there is, e.g., a
+    /// pattern such as "is SubClass".
+    ///
+    /// \param bindPatternVarsOneWay When true, generate fresh type variables
+    /// for the types of each variable declared within the pattern, along
+    /// with a one-way constraint binding that to the type to which the
+    /// variable will be ascribed or inferred.
     Type getTypeForPattern(Pattern *pattern, ConstraintLocatorBuilder locator,
                            Type externalPatternType,
                            bool bindPatternVarsOneWay) {
@@ -2257,14 +2269,9 @@ namespace {
         if (bindPatternVarsOneWay) {
           oneWayVarType = CS.createTypeVariable(
               CS.getConstraintLocator(locator), TVO_CanBindToNoEscape);
-          if (externalPatternType) {
-            CS.addConstraint(
-                ConstraintKind::OneWayEqual, oneWayVarType,
-                externalPatternType, locator);
-          } else {
-            CS.addConstraint(
-                ConstraintKind::OneWayEqual, oneWayVarType, varType, locator);
-          }
+          CS.addConstraint(
+              ConstraintKind::OneWayEqual, oneWayVarType,
+              externalPatternType ? externalPatternType : varType, locator);
         }
 
         // If there is an externally-imposed type.
@@ -2357,16 +2364,13 @@ namespace {
       case PatternKind::OptionalSome: {
         // Remove an optional from the object type.
         if (externalPatternType) {
-          if (Type objType = externalPatternType->getOptionalObjectType()) {
-            externalPatternType = objType;
-          } else {
-            Type objVar = CS.createTypeVariable(
-                CS.getConstraintLocator(locator), TVO_CanBindToNoEscape);
-            CS.addConstraint(
-                ConstraintKind::Equal, OptionalType::get(objVar),
-                externalPatternType, locator);
-            externalPatternType = objVar;
-          }
+          Type objVar = CS.createTypeVariable(
+              CS.getConstraintLocator(locator), TVO_CanBindToNoEscape);
+          CS.addConstraint(
+              ConstraintKind::OptionalObject, externalPatternType,
+              objVar, locator);
+
+          externalPatternType = objVar;
         }
 
         // The subpattern must have optional type.