Merge pull request #30322 from DougGregor/function-builders-extras

[Function builders] Additional entry points

Author: DougGregor

include/swift/AST/KnownIdentifiers.def

@@ -35,7 +35,9 @@ IDENTIFIER(buildBlock)
 IDENTIFIER(buildDo)
 IDENTIFIER(buildEither)
 IDENTIFIER(buildExpression)
+IDENTIFIER(buildFinalResult)
 IDENTIFIER(buildIf)
+IDENTIFIER(buildOptional)
 IDENTIFIER(callAsFunction)
 IDENTIFIER(Change)
 IDENTIFIER_WITH_NAME(code_, "_code")

lib/Sema/BuilderTransform.cpp

@@ -49,6 +49,7 @@ class BuilderClosureVisitor
   ASTContext &ctx;
   Type builderType;
   NominalTypeDecl *builder = nullptr;
+  Identifier buildOptionalId;
   llvm::SmallDenseMap<Identifier, bool> supportedOps;
 
   SkipUnhandledConstructInFunctionBuilder::UnhandledNode unhandledNode;
@@ -201,6 +202,14 @@ class BuilderClosureVisitor
     builder = builderType->getAnyNominal();
     applied.builderType = builderType;
     applied.bodyResultType = bodyResultType;
+
+    // Use buildOptional(_:) if available, otherwise fall back to buildIf
+    // when available.
+    if (builderSupports(ctx.Id_buildOptional) ||
+        !builderSupports(ctx.Id_buildIf))
+      buildOptionalId = ctx.Id_buildOptional;
+    else
+      buildOptionalId = ctx.Id_buildIf;
   }
 
   /// Apply the builder transform to the given statement.
@@ -210,6 +219,15 @@ class BuilderClosureVisitor
       return None;
 
     applied.returnExpr = buildVarRef(bodyVar, stmt->getEndLoc());
+
+    // If there is a buildFinalResult(_:), call it.
+    ASTContext &ctx = cs->getASTContext();
+    if (builderSupports(ctx.Id_buildFinalResult, { Identifier() })) {
+      applied.returnExpr = buildCallIfWanted(
+          applied.returnExpr->getLoc(), ctx.Id_buildFinalResult,
+          { applied.returnExpr }, { Identifier() });
+    }
+
     applied.returnExpr = cs->buildTypeErasedExpr(applied.returnExpr,
                                                  dc, applied.bodyResultType,
                                                  CTP_ReturnStmt);
@@ -407,8 +425,8 @@ class BuilderClosureVisitor
     if (!isBuildableIfChainRecursive(ifStmt, numPayloads, isOptional))
       return false;
 
-    // If there's a missing 'else', we need 'buildIf' to exist.
-    if (isOptional && !builderSupports(ctx.Id_buildIf))
+    // If there's a missing 'else', we need 'buildOptional' to exist.
+    if (isOptional && !builderSupports(buildOptionalId))
       return false;
 
     // If there are multiple clauses, we need 'buildEither(first:)' and
@@ -514,9 +532,9 @@ class BuilderClosureVisitor
     // 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 && isTopLevel) {
-      thenExpr = buildCallIfWanted(ifStmt->getEndLoc(), ctx.Id_buildIf,
+      thenExpr = buildCallIfWanted(ifStmt->getEndLoc(), buildOptionalId,
                                    thenExpr,  /*argLabels=*/{ });
-      elseExpr = buildCallIfWanted(ifStmt->getEndLoc(), ctx.Id_buildIf,
+      elseExpr = buildCallIfWanted(ifStmt->getEndLoc(), buildOptionalId,
                                    elseExpr,  /*argLabels=*/{ });
     }
 

test/Constraints/function_builder.swift

@@ -617,3 +617,50 @@ testSwitchCombined(getE(1))
 // CHECK: testSwitchCombined
 // CHECK-SAME: second("just 42")
 testSwitchCombined(getE(2))
+
+
+// Test buildOptional(_:) as an alternative to buildIf(_:).
+@_functionBuilder
+struct TupleBuilderWithOpt {
+  static func buildBlock<T1>(_ t1: T1) -> (T1) {
+    return (t1)
+  }
+
+  static func buildBlock<T1, T2>(_ t1: T1, _ t2: T2) -> (T1, T2) {
+    return (t1, t2)
+  }
+  
+  static func buildBlock<T1, T2, T3>(_ t1: T1, _ t2: T2, _ t3: T3)
+      -> (T1, T2, T3) {
+    return (t1, t2, t3)
+  }
+
+  static func buildBlock<T1, T2, T3, T4>(_ t1: T1, _ t2: T2, _ t3: T3, _ t4: T4)
+      -> (T1, T2, T3, T4) {
+    return (t1, t2, t3, t4)
+  }
+
+  static func buildBlock<T1, T2, T3, T4, T5>(
+    _ t1: T1, _ t2: T2, _ t3: T3, _ t4: T4, _ t5: T5
+  ) -> (T1, T2, T3, T4, T5) {
+    return (t1, t2, t3, t4, t5)
+  }
+
+  static func buildDo<T>(_ value: T) -> T { return value }
+  static func buildOptional<T>(_ value: T?) -> T? { return value }
+
+  static func buildEither<T,U>(first value: T) -> Either<T,U> {
+    return .first(value)
+  }
+  static func buildEither<T,U>(second value: U) -> Either<T,U> {
+    return .second(value)
+  }
+}
+
+func tuplifyWithOpt<T>(_ cond: Bool, @TupleBuilderWithOpt body: (Bool) -> T) {
+  print(body(cond))
+}
+tuplifyWithOpt(true) { c in
+  "1"
+  3.14159
+}

test/Constraints/function_builder_diags.swift

@@ -480,3 +480,63 @@ func testCaseVarTypes(e: E3) {
     }
   }
 }
+
+// Test for buildFinalResult.
+@_functionBuilder
+struct WrapperBuilder {
+  static func buildBlock() -> () { }
+  
+  static func buildBlock<T1>(_ t1: T1) -> T1 {
+    return t1
+  }
+  
+  static func buildBlock<T1, T2>(_ t1: T1, _ t2: T2) -> (T1, T2) {
+    return (t1, t2)
+  }
+  
+  static func buildBlock<T1, T2, T3>(_ t1: T1, _ t2: T2, _ t3: T3)
+      -> (T1, T2, T3) {
+    return (t1, t2, t3)
+  }
+
+  static func buildBlock<T1, T2, T3, T4>(_ t1: T1, _ t2: T2, _ t3: T3, _ t4: T4)
+      -> (T1, T2, T3, T4) {
+    return (t1, t2, t3, t4)
+  }
+
+  static func buildBlock<T1, T2, T3, T4, T5>(
+    _ t1: T1, _ t2: T2, _ t3: T3, _ t4: T4, _ t5: T5
+  ) -> (T1, T2, T3, T4, T5) {
+    return (t1, t2, t3, t4, t5)
+  }
+
+  static func buildDo<T>(_ value: T) -> T { return value }
+  static func buildIf<T>(_ value: T?) -> T? { return value }
+
+  static func buildEither<T,U>(first value: T) -> Either<T,U> {
+    return .first(value)
+  }
+  static func buildEither<T,U>(second value: U) -> Either<T,U> {
+    return .second(value)
+  }
+  static func buildFinalResult<T>(_ value: T) -> Wrapper<T> {
+    return Wrapper(value: value)
+  }
+}
+
+struct Wrapper<T> {
+  var value: T
+}
+
+func wrapperify<T>(_ cond: Bool, @WrapperBuilder body: (Bool) -> T) -> T{
+  return body(cond)
+}
+
+func testWrapperBuilder() {
+  let x = wrapperify(true) { c in
+    3.14159
+    "hello"
+  }
+
+  let _: Int = x // expected-error{{cannot convert value of type 'Wrapper<(Double, String)>' to specified type 'Int'}}
+}