Merge pull request #41576 from rxwei/buildpartialblock

[ResultBuilders] `buildPartialBlock` support

Author: rxwei

include/swift/AST/DiagnosticsSema.def

@@ -6112,6 +6112,15 @@ WARNING(result_builder_missing_limited_availability, none,
 ERROR(result_builder_static_buildblock, none,
       "result builder must provide at least one static 'buildBlock' "
       "method", ())
+ERROR(result_builder_static_buildblock_or_buildpartialblock, none,
+      "result builder must provide at least one static 'buildBlock' "
+      "method, or both 'buildPartialBlock(first:)' and "
+      "'buildPartialBlock(accumulated:next:)'", ())
+ERROR(result_builder_missing_available_buildpartialblock, none,
+      "result builder %0 does not implement any 'buildBlock' or a combination "
+      "of 'buildPartialBlock(first:)' and "
+      "'buildPartialBlock(accumulated:next:)' with sufficient availability for "
+      "this call site", (Type))
 NOTE(result_builder_non_static_buildblock, none,
      "did you mean to make instance method 'buildBlock' static?", ())
 NOTE(result_builder_buildblock_enum_case, none,

include/swift/AST/KnownIdentifiers.def

@@ -28,6 +28,7 @@ IDENTIFIER(allKeys)
 IDENTIFIER(alloc)
 IDENTIFIER(allocWithZone)
 IDENTIFIER(allZeros)
+IDENTIFIER(accumulated)
 IDENTIFIER(ActorType)
 IDENTIFIER(Any)
 IDENTIFIER(ArrayLiteralElement)
@@ -41,6 +42,7 @@ IDENTIFIER(buildFinalResult)
 IDENTIFIER(buildIf)
 IDENTIFIER(buildLimitedAvailability)
 IDENTIFIER(buildOptional)
+IDENTIFIER(buildPartialBlock)
 IDENTIFIER(callAsFunction)
 IDENTIFIER(Change)
 IDENTIFIER_WITH_NAME(code_, "_code")

lib/Sema/BuilderTransform.cpp

@@ -71,7 +71,7 @@ class BuilderClosureVisitor
   Type builderType;
   NominalTypeDecl *builder = nullptr;
   Identifier buildOptionalId;
-  llvm::SmallDenseMap<Identifier, bool> supportedOps;
+  llvm::SmallDenseMap<DeclName, bool> supportedOps;
 
   SkipUnhandledConstructInResultBuilder::UnhandledNode unhandledNode;
 
@@ -141,15 +141,18 @@ class BuilderClosureVisitor
   }
 
   /// Check whether the builder supports the given operation.
-  bool builderSupports(Identifier fnName,
-                       ArrayRef<Identifier> argLabels = {}) {
-    auto known = supportedOps.find(fnName);
+  bool builderSupports(Identifier fnBaseName,
+                       ArrayRef<Identifier> argLabels = {},
+                       bool checkAvailability = false) {
+    DeclName name(dc->getASTContext(), fnBaseName, argLabels);
+    auto known = supportedOps.find(name);
     if (known != supportedOps.end()) {
       return known->second;
     }
 
-    return supportedOps[fnName] = TypeChecker::typeSupportsBuilderOp(
-               builderType, dc, fnName, argLabels);
+    return supportedOps[name] = TypeChecker::typeSupportsBuilderOp(
+               builderType, dc, fnBaseName, argLabels, /*allResults*/ {},
+               checkAvailability);
   }
 
   /// Build an implicit variable in this context.
@@ -368,11 +371,39 @@ class BuilderClosureVisitor
       return nullptr;
 
     Expr *call = nullptr;
-    // If the builder supports `buildBlock(combining:into:)`, use this to
-    // combine subexpressions pairwise.
+    // If the builder supports `buildPartialBlock(first:)` and
+    // `buildPartialBlock(accumulated:next:)`, use this to combine
+    // subexpressions pairwise.
     if (ctx.LangOpts.EnableExperimentalPairwiseBuildBlock &&
         !expressions.empty() &&
-        builderSupports(ctx.Id_buildBlock, {ctx.Id_combining, ctx.Id_into})) {
+        builderSupports(ctx.Id_buildPartialBlock, {ctx.Id_first},
+                        /*checkAvailability*/ true) &&
+        builderSupports(ctx.Id_buildPartialBlock,
+                        {ctx.Id_accumulated, ctx.Id_next},
+                        /*checkAvailability*/ true)) {
+      // NOTE: The current implementation uses one-way constraints in between
+      // subexpressions. It's functionally equivalent to the following:
+      //   let v0 = Builder.buildPartialBlock(first: arg_0)
+      //   let v1 = Builder.buildPartialBlock(accumulated: arg_1, next: v0)
+      //   ...
+      //   return Builder.buildPartialBlock(accumulated: arg_n, next: ...)
+      call = buildCallIfWanted(braceStmt->getStartLoc(),
+                               ctx.Id_buildPartialBlock,
+                               {expressions.front()},
+                               /*argLabels=*/{ctx.Id_first});
+      for (auto *expr : llvm::drop_begin(expressions)) {
+        call = buildCallIfWanted(braceStmt->getStartLoc(),
+                                 ctx.Id_buildPartialBlock,
+                                 {new (ctx) OneWayExpr(call), expr},
+                                 {ctx.Id_accumulated, ctx.Id_next});
+      }
+    }
+    // TODO: Remove support for the old method name,
+    // `buildBlock(combining:into:)`.
+    else if (ctx.LangOpts.EnableExperimentalPairwiseBuildBlock &&
+             !expressions.empty() &&
+             builderSupports(ctx.Id_buildBlock,
+                             {ctx.Id_combining, ctx.Id_into})) {
       // NOTE: The current implementation uses one-way constraints in between
       // subexpressions. It's functionally equivalent to the following:
       //   let v0 = Builder.buildBlock(arg_0)
@@ -387,6 +418,17 @@ class BuilderClosureVisitor
                                  {ctx.Id_combining, ctx.Id_into});
       }
     }
+    // If `buildBlock` does not exist at this point, it could be the case that
+    // `buildPartialBlock` did not have the sufficient availability for this
+    // call site.  Diagnose it.
+    else if (ctx.LangOpts.EnableExperimentalPairwiseBuildBlock &&
+             !builderSupports(ctx.Id_buildBlock)) {
+      ctx.Diags.diagnose(
+          braceStmt->getStartLoc(),
+          diag::result_builder_missing_available_buildpartialblock,
+          builderType);
+      return nullptr;
+    }
     // Otherwise, call `buildBlock` on all subexpressions.
     else {
       // Call Builder.buildBlock(... args ...)
@@ -2016,7 +2058,8 @@ std::vector<ReturnStmt *> TypeChecker::findReturnStatements(AnyFunctionRef fn) {
 
 bool TypeChecker::typeSupportsBuilderOp(
     Type builderType, DeclContext *dc, Identifier fnName,
-    ArrayRef<Identifier> argLabels, SmallVectorImpl<ValueDecl *> *allResults) {
+    ArrayRef<Identifier> argLabels, SmallVectorImpl<ValueDecl *> *allResults,
+    bool checkAvailability) {
   bool foundMatch = false;
   SmallVector<ValueDecl *, 4> foundDecls;
   dc->lookupQualified(
@@ -2036,6 +2079,17 @@ bool TypeChecker::typeSupportsBuilderOp(
           continue;
       }
 
+      // If we are checking availability, the candidate must have enough
+      // availability in the calling context.
+      if (checkAvailability) {
+        if (AvailableAttr::isUnavailable(func))
+          continue;
+        if (TypeChecker::checkDeclarationAvailability(
+                func, ExportContext::forFunctionBody(
+                    dc, extractNearestSourceLoc(dc))))
+          continue;
+      }
+
       foundMatch = true;
       break;
     }

lib/Sema/TypeCheckAttr.cpp

@@ -3266,15 +3266,30 @@ void AttributeChecker::visitPropertyWrapperAttr(PropertyWrapperAttr *attr) {
 
 void AttributeChecker::visitResultBuilderAttr(ResultBuilderAttr *attr) {
   auto *nominal = dyn_cast<NominalTypeDecl>(D);
+  auto &ctx = D->getASTContext();
   SmallVector<ValueDecl *, 4> potentialMatches;
   bool supportsBuildBlock = TypeChecker::typeSupportsBuilderOp(
-      nominal->getDeclaredType(), nominal, D->getASTContext().Id_buildBlock,
+      nominal->getDeclaredType(), nominal, ctx.Id_buildBlock,
       /*argLabels=*/{}, &potentialMatches);
-
-  if (!supportsBuildBlock) {
+  bool isBuildPartialBlockFeatureEnabled =
+      ctx.LangOpts.EnableExperimentalPairwiseBuildBlock;
+  bool supportsBuildPartialBlock = isBuildPartialBlockFeatureEnabled &&
+      TypeChecker::typeSupportsBuilderOp(
+          nominal->getDeclaredType(), nominal,
+          ctx.Id_buildPartialBlock,
+          /*argLabels=*/{ctx.Id_first}, &potentialMatches) &&
+      TypeChecker::typeSupportsBuilderOp(
+          nominal->getDeclaredType(), nominal,
+          ctx.Id_buildPartialBlock,
+          /*argLabels=*/{ctx.Id_accumulated, ctx.Id_next}, &potentialMatches);
+
+  if (!supportsBuildBlock && !supportsBuildPartialBlock) {
     {
       auto diag = diagnose(
-          nominal->getLoc(), diag::result_builder_static_buildblock);
+          nominal->getLoc(),
+          isBuildPartialBlockFeatureEnabled
+              ? diag::result_builder_static_buildblock_or_buildpartialblock
+              : diag::result_builder_static_buildblock);
 
       // If there were no close matches, propose adding a stub.
       SourceLoc buildInsertionLoc;

lib/Sema/TypeChecker.h

@@ -1185,7 +1185,8 @@ UnresolvedMemberExpr *getUnresolvedMemberChainBase(Expr *expr);
 /// are verified against any candidates.
 bool typeSupportsBuilderOp(Type builderType, DeclContext *dc, Identifier fnName,
                            ArrayRef<Identifier> argLabels = {},
-                           SmallVectorImpl<ValueDecl *> *allResults = nullptr);
+                           SmallVectorImpl<ValueDecl *> *allResults = nullptr,
+                           bool checkAvailability = false);
 
 /// Forces all changes specified by the module's access notes file to be
 /// applied to this declaration. It is safe to call this function more than