Merge pull request #29409 from DougGregor/function-builder-multi-if

[Function builders] Support multiple Boolean conditions in 'if' statements

Author: DougGregor

lib/Sema/BuilderTransform.cpp

@@ -324,18 +324,27 @@ class BuilderClosureVisitor
   CONTROL_FLOW_STMT(Yield)
   CONTROL_FLOW_STMT(Defer)
 
-  static Expr *getTrivialBooleanCondition(StmtCondition condition) {
-    if (condition.size() != 1)
-      return nullptr;
+  /// Whether we can handle all of the conditions for this statement.
+  static bool canHandleStmtConditions(StmtCondition condition) {
+    for (const auto &element : condition) {
+      switch (element.getKind()) {
+      case StmtConditionElement::CK_Boolean:
+        continue;
 
-    return condition.front().getBooleanOrNull();
+      case StmtConditionElement::CK_PatternBinding:
+      case StmtConditionElement::CK_Availability:
+        return false;
+      }
+    }
+
+    return true;
   }
 
   static bool isBuildableIfChainRecursive(IfStmt *ifStmt,
                                           unsigned &numPayloads,
                                           bool &isOptional) {
-    // The conditional must be trivial.
-    if (!getTrivialBooleanCondition(ifStmt->getCond()))
+    // Check whether we can handle the conditional.
+    if (!canHandleStmtConditions(ifStmt->getCond()))
       return false;
 
     // The 'then' clause contributes a payload.
@@ -461,26 +470,37 @@ class BuilderClosureVisitor
           payloadIndex + 1, numPayloads, isOptional);
     }
 
-    // Generate constraints for the various subexpressions.
-    auto condExpr = getTrivialBooleanCondition(ifStmt->getCond());
-    assert(condExpr && "Cannot get here without a trivial Boolean condition");
-    condExpr = cs->generateConstraints(condExpr, dc);
-    if (!condExpr) {
-      hadError = true;
-      return nullptr;
-    }
-
     // Condition must convert to Bool.
     // FIXME: This should be folded into constraint generation for conditions.
     auto boolDecl = ctx.getBoolDecl();
     if (!boolDecl) {
       hadError = true;
       return nullptr;
     }
-    cs->addConstraint(ConstraintKind::Conversion,
-                      cs->getType(condExpr),
-                      boolDecl->getDeclaredType(),
-                      cs->getConstraintLocator(condExpr));
+
+    // Generate constraints for the conditions.
+    for (const auto &condElement : ifStmt->getCond()) {
+      switch (condElement.getKind()) {
+      case StmtConditionElement::CK_Boolean: {
+        Expr *condExpr = condElement.getBoolean();
+        condExpr = cs->generateConstraints(condExpr, dc);
+        if (!condExpr) {
+          hadError = true;
+          return nullptr;
+        }
+
+        cs->addConstraint(ConstraintKind::Conversion,
+                          cs->getType(condExpr),
+                          boolDecl->getDeclaredType(),
+                          cs->getConstraintLocator(condExpr));
+        continue;
+      }
+
+      case StmtConditionElement::CK_PatternBinding:
+      case StmtConditionElement::CK_Availability:
+        llvm_unreachable("unhandled statement condition");
+      }
+    }
 
     // The operand should have optional type if we had optional results,
     // so we just need to call `buildIf` now, since we're at the top level.
@@ -850,19 +870,29 @@ class BuilderClosureRewriter
 
   Stmt *visitIfStmt(IfStmt *ifStmt, FunctionBuilderTarget target) {
     // Rewrite the condition.
-    // FIXME: We should handle the whole condition within the type system.
-    auto cond = ifStmt->getCond();
-    auto condExpr = cond.front().getBoolean();
-    auto finalCondExpr = rewriteExpr(condExpr);
-
-    // Load the condition if needed.
-    if (finalCondExpr->getType()->is<LValueType>()) {
-      auto &cs = solution.getConstraintSystem();
-      finalCondExpr = cs.addImplicitLoadExpr(finalCondExpr);
-    }
+    auto condition = ifStmt->getCond();
+    for (auto &condElement : condition) {
+      switch (condElement.getKind()) {
+      case StmtConditionElement::CK_Boolean: {
+        auto condExpr = condElement.getBoolean();
+        auto finalCondExpr = rewriteExpr(condExpr);
+
+        // Load the condition if needed.
+        if (finalCondExpr->getType()->is<LValueType>()) {
+          auto &cs = solution.getConstraintSystem();
+          finalCondExpr = cs.addImplicitLoadExpr(finalCondExpr);
+        }
+
+        condElement.setBoolean(finalCondExpr);
+        continue;
+      }
 
-    cond.front().setBoolean(finalCondExpr);
-    ifStmt->setCond(cond);
+      case StmtConditionElement::CK_PatternBinding:
+      case StmtConditionElement::CK_Availability:
+        llvm_unreachable("unhandled statement condition");
+      }
+    }
+    ifStmt->setCond(condition);
 
     assert(target.kind == FunctionBuilderTarget::TemporaryVar);
     auto temporaryVar = target.captured.first;

stdlib/public/Darwin/CMakeLists.txt

@@ -44,7 +44,7 @@ set(all_overlays
   MetalKit
   ModelIO
   NaturalLanguage
-  Network
+#  Network
   ObjectiveC
   OpenCL
   os

test/Constraints/function_builder.swift

@@ -461,3 +461,18 @@ func testNestedClosuresWithDependencies(cond: Bool) {
     }
   }
 }
+
+// Check that we can handle multiple conditions in an 'if' statement.
+func testIfConditions(cond: Bool, c1: Bool, i1: Int, i2: Int) {
+  tuplify(cond) { x in
+    "testIfConditions"
+    if i1 == i2, c1, x {
+      1
+      "hello"
+    }
+    3.14159
+  }
+}
+testIfConditions(cond: true, c1: true, i1: 1, i2: 1)
+// CHECK: testIfConditions
+// CHECK-SAME: hello