[BuilderTransform] AST Transform: Inject buildOptional to top-level…

lib/Sema/BuilderTransform.cpp

@@ -1245,10 +1245,26 @@ class ResultBuilderTransform
         auto *builderCall =
             buildWrappedChainPayload(branchVarRef, i, numPayloads, isOptional);
 
+        auto isTopLevel = [&](Stmt *anchor) {
+          if (ifStmt->getThenStmt() == anchor)
+            return true;
+
+          // The situation is this:
+          //
+          // if <cond> {
+          //   ...
+          // } else if <other-cond> {
+          //   ...
+          // }
+          if (auto *innerIf = getAsStmt<IfStmt>(ifStmt->getElseStmt()))
+            return innerIf->getThenStmt() == anchor;
+
+          return ifStmt->getElseStmt() == anchor;
+        };
+
         // 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.
-        if (isOptional && (ifStmt->getThenStmt() == anchor ||
-                           ifStmt->getElseStmt() == anchor)) {
+        if (isOptional && isTopLevel(anchor)) {
           builderCall = buildCallIfWanted(ifStmt->getEndLoc(),
                                           builder.getBuildOptionalId(),
                                           builderCall, /*argLabels=*/{});

test/Constraints/result_builder.swift

@@ -1216,3 +1216,39 @@ do {
   }
   // CHECK: (42, "", [true])
 }
+
+do {
+  @resultBuilder
+  struct MyBuilder {
+    static func buildBlock<T1: ExpressibleByStringLiteral>(_ t1: T1) -> (T1) {
+      return (t1)
+    }
+
+    static func buildBlock<T1, T2>(_ t1: T1, _ t2: T2) -> (T1, T2) {
+      return (t1, t2)
+    }
+
+    static func buildOptional<T>(_ value: T?) -> T { return value! }
+
+    static func buildEither<T>(first value: T) -> T {
+      return value
+    }
+
+    static func buildEither<U>(second value: U) -> U {
+      return value
+    }
+  }
+
+  func test<T>(@MyBuilder _ builder: (Int) -> T) {
+    print(builder(42))
+  }
+
+  test {
+    if $0 < 0 {
+      "\($0)"
+    } else if $0 == 42 {
+      "the answer"
+    }
+  }
+  // CHECK: the answer
+}