Merge pull request #60803 from hamishknight/part-and-partial

Author: hamishknight

include/swift/Sema/ConstraintSystem.h

@@ -107,6 +107,33 @@ enum class FreeTypeVariableBinding {
   UnresolvedType
 };
 
+/// Describes whether or not a result builder method is supported.
+struct ResultBuilderOpSupport {
+  enum Classification {
+    Unsupported,
+    Unavailable,
+    Supported
+  };
+  Classification Kind;
+
+  ResultBuilderOpSupport(Classification Kind) : Kind(Kind) {}
+
+  /// Returns whether or not the builder method is supported. If
+  /// \p requireAvailable is true, an unavailable method will be considered
+  /// unsupported.
+  bool isSupported(bool requireAvailable) const {
+    switch (Kind) {
+    case Unsupported:
+      return false;
+    case Unavailable:
+      return !requireAvailable;
+    case Supported:
+      return true;
+    }
+    llvm_unreachable("Unhandled case in switch!");
+  }
+};
+
 namespace constraints {
 
 struct ResultBuilder {
@@ -115,7 +142,9 @@ struct ResultBuilder {
   /// An implicit variable that represents `Self` type of the result builder.
   VarDecl *BuilderSelf;
   Type BuilderType;
-  llvm::SmallDenseMap<DeclName, bool> SupportedOps;
+
+  /// Cache of supported result builder operations.
+  llvm::SmallDenseMap<DeclName, ResultBuilderOpSupport> SupportedOps;
 
   Identifier BuildOptionalId;
 
@@ -143,6 +172,13 @@ struct ResultBuilder {
 
   bool supportsOptional() { return supports(getBuildOptionalId()); }
 
+  /// Checks whether the `buildPartialBlock` method is supported.
+  bool supportsBuildPartialBlock(bool checkAvailability);
+
+  /// Checks whether the builder can use `buildPartialBlock` to combine
+  /// expressions, instead of `buildBlock`.
+  bool canUseBuildPartialBlock();
+
   Expr *buildCall(SourceLoc loc, Identifier fnName,
                   ArrayRef<Expr *> argExprs,
                   ArrayRef<Identifier> argLabels) const;

lib/Sema/BuilderTransform.cpp

@@ -432,12 +432,7 @@ class BuilderClosureVisitor
     // If the builder supports `buildPartialBlock(first:)` and
     // `buildPartialBlock(accumulated:next:)`, use this to combine
     // subexpressions pairwise.
-    if (!expressions.empty() &&
-        builder.supports(ctx.Id_buildPartialBlock, {ctx.Id_first},
-                         /*checkAvailability*/ true) &&
-        builder.supports(ctx.Id_buildPartialBlock,
-                         {ctx.Id_accumulated, ctx.Id_next},
-                         /*checkAvailability*/ true)) {
+    if (!expressions.empty() && builder.canUseBuildPartialBlock()) {
       // 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)
@@ -1087,12 +1082,7 @@ class ResultBuilderTransform
       // If the builder supports `buildPartialBlock(first:)` and
       // `buildPartialBlock(accumulated:next:)`, use this to combine
       // sub-expressions pairwise.
-      if (!buildBlockArguments.empty() &&
-          builder.supports(ctx.Id_buildPartialBlock, {ctx.Id_first},
-                           /*checkAvailability*/ true) &&
-          builder.supports(ctx.Id_buildPartialBlock,
-                           {ctx.Id_accumulated, ctx.Id_next},
-                           /*checkAvailability*/ true)) {
+      if (!buildBlockArguments.empty() && builder.canUseBuildPartialBlock()) {
         //   let v0 = Builder.buildPartialBlock(first: arg_0)
         //   let v1 = Builder.buildPartialBlock(accumulated: v0, next: arg_1)
         //   ...
@@ -2777,11 +2767,22 @@ std::vector<ReturnStmt *> TypeChecker::findReturnStatements(AnyFunctionRef fn) {
   return precheck.getReturnStmts();
 }
 
-bool TypeChecker::typeSupportsBuilderOp(
+ResultBuilderOpSupport TypeChecker::checkBuilderOpSupport(
     Type builderType, DeclContext *dc, Identifier fnName,
-    ArrayRef<Identifier> argLabels, SmallVectorImpl<ValueDecl *> *allResults,
-    bool checkAvailability) {
+    ArrayRef<Identifier> argLabels, SmallVectorImpl<ValueDecl *> *allResults) {
+
+  auto isUnavailable = [&](Decl *D) -> bool {
+    if (AvailableAttr::isUnavailable(D))
+      return true;
+
+    auto loc = extractNearestSourceLoc(dc);
+    auto context = ExportContext::forFunctionBody(dc, loc);
+    return TypeChecker::checkDeclarationAvailability(D, context).hasValue();
+  };
+
   bool foundMatch = false;
+  bool foundUnavailable = false;
+
   SmallVector<ValueDecl *, 4> foundDecls;
   dc->lookupQualified(
       builderType, DeclNameRef(fnName),
@@ -2800,17 +2801,12 @@ 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;
+      // Check if the the candidate has a suitable availability for the
+      // calling context.
+      if (isUnavailable(func)) {
+        foundUnavailable = true;
+        continue;
       }
-
       foundMatch = true;
       break;
     }
@@ -2819,7 +2815,24 @@ bool TypeChecker::typeSupportsBuilderOp(
   if (allResults)
     allResults->append(foundDecls.begin(), foundDecls.end());
 
-  return foundMatch;
+  if (!foundMatch) {
+    return foundUnavailable ? ResultBuilderOpSupport::Unavailable
+                            : ResultBuilderOpSupport::Unsupported;
+  }
+  // If the builder type itself isn't available, don't consider any builder
+  // method available.
+  if (auto *D = builderType->getAnyNominal()) {
+    if (isUnavailable(D))
+      return ResultBuilderOpSupport::Unavailable;
+  }
+  return ResultBuilderOpSupport::Supported;
+}
+
+bool TypeChecker::typeSupportsBuilderOp(
+    Type builderType, DeclContext *dc, Identifier fnName,
+    ArrayRef<Identifier> argLabels, SmallVectorImpl<ValueDecl *> *allResults) {
+  return checkBuilderOpSupport(builderType, dc, fnName, argLabels, allResults)
+      .isSupported(/*requireAvailable*/ false);
 }
 
 Type swift::inferResultBuilderComponentType(NominalTypeDecl *builder) {
@@ -2978,18 +2991,43 @@ ResultBuilder::ResultBuilder(ConstraintSystem *CS, DeclContext *DC,
   }
 }
 
+bool ResultBuilder::supportsBuildPartialBlock(bool checkAvailability) {
+  auto &ctx = DC->getASTContext();
+  return supports(ctx.Id_buildPartialBlock, {ctx.Id_first},
+                  checkAvailability) &&
+         supports(ctx.Id_buildPartialBlock, {ctx.Id_accumulated, ctx.Id_next},
+                  checkAvailability);
+}
+
+bool ResultBuilder::canUseBuildPartialBlock() {
+  // If buildPartialBlock doesn't exist at all, we can't use it.
+  if (!supportsBuildPartialBlock(/*checkAvailability*/ false))
+    return false;
+
+  // If buildPartialBlock exists and is available, use it.
+  if (supportsBuildPartialBlock(/*checkAvailability*/ true))
+    return true;
+
+  // We have buildPartialBlock, but it is unavailable. We can however still
+  // use it if buildBlock is also unavailable.
+  auto &ctx = DC->getASTContext();
+  return supports(ctx.Id_buildBlock) &&
+         !supports(ctx.Id_buildBlock, /*labels*/ {},
+                   /*checkAvailability*/ true);
+}
+
 bool ResultBuilder::supports(Identifier fnBaseName,
                              ArrayRef<Identifier> argLabels,
                              bool checkAvailability) {
   DeclName name(DC->getASTContext(), fnBaseName, argLabels);
   auto known = SupportedOps.find(name);
-  if (known != SupportedOps.end()) {
-    return known->second;
-  }
+  if (known != SupportedOps.end())
+    return known->second.isSupported(checkAvailability);
 
-  return SupportedOps[name] = TypeChecker::typeSupportsBuilderOp(
-             BuilderType, DC, fnBaseName, argLabels, /*allResults*/ {},
-             checkAvailability);
+  auto support = TypeChecker::checkBuilderOpSupport(
+      BuilderType, DC, fnBaseName, argLabels, /*allResults*/ {});
+  SupportedOps.insert({name, support});
+  return support.isSupported(checkAvailability);
 }
 
 Expr *ResultBuilder::buildCall(SourceLoc loc, Identifier fnName,

lib/Sema/TypeChecker.h

@@ -1234,10 +1234,20 @@ UnresolvedMemberExpr *getUnresolvedMemberChainBase(Expr *expr);
 /// Checks whether a result builder type has a well-formed result builder
 /// method with the given name. If provided and non-empty, the argument labels
 /// are verified against any candidates.
+ResultBuilderOpSupport
+checkBuilderOpSupport(Type builderType, DeclContext *dc, Identifier fnName,
+                      ArrayRef<Identifier> argLabels = {},
+                      SmallVectorImpl<ValueDecl *> *allResults = nullptr);
+
+/// Checks whether a result builder type has a well-formed result builder
+/// method with the given name. If provided and non-empty, the argument labels
+/// are verified against any candidates.
+///
+/// This will return \c true even if the builder method is unavailable. Use
+/// \c checkBuilderOpSupport if availability should be checked.
 bool typeSupportsBuilderOp(Type builderType, DeclContext *dc, Identifier fnName,
                            ArrayRef<Identifier> argLabels = {},
-                           SmallVectorImpl<ValueDecl *> *allResults = nullptr,
-                           bool checkAvailability = false);
+                           SmallVectorImpl<ValueDecl *> *allResults = nullptr);
 
 /// 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

test/Constraints/result_builder_availability.swift

@@ -167,3 +167,59 @@ tuplifyWithAvailabilityErasure(true) { cond in
     globalFuncAvailableOn10_52()
   }
 }
+
+// rdar://97533700 – Make sure we can prefer an unavailable buildPartialBlock if
+// buildBlock also isn't available.
+
+@resultBuilder
+struct UnavailableBuildPartialBlock {
+  static func buildPartialBlock(first: Int) -> Int { 0 }
+
+  @available(*, unavailable)
+  static func buildPartialBlock(accumulated: Int, next: Int) -> Int { 0 }
+
+  static func buildBlock(_ x: Int...) -> Int { 0 }
+}
+
+@UnavailableBuildPartialBlock
+func testUnavailableBuildPartialBlock() -> Int {
+  // We can use buildBlock here.
+  2
+  3
+}
+
+@resultBuilder
+struct UnavailableBuildPartialBlockAndBuildBlock {
+  @available(*, unavailable)
+  static func buildPartialBlock(first: Int) -> Int { 0 }
+  // expected-note@-1 {{'buildPartialBlock(first:)' has been explicitly marked unavailable here}}
+
+  static func buildPartialBlock(accumulated: Int, next: Int) -> Int { 0 }
+
+  @available(*, unavailable)
+  static func buildBlock(_ x: Int...) -> Int { 0 }
+}
+
+// We can still use buildPartialBlock here as both are unavailable.
+@UnavailableBuildPartialBlockAndBuildBlock
+func testUnavailableBuildPartialBlockAndBuildBlock() -> Int {
+  // expected-error@-1 {{'buildPartialBlock(first:)' is unavailable}}
+  2
+  3
+}
+
+@available(*, unavailable)
+@resultBuilder
+struct UnavailableBuilderWithPartialBlock { // expected-note {{'UnavailableBuilderWithPartialBlock' has been explicitly marked unavailable here}}
+  @available(*, unavailable)
+  static func buildPartialBlock(first: String) -> Int { 0 }
+  static func buildPartialBlock(accumulated: Int, next: Int) -> Int { 0 }
+  static func buildBlock(_ x: Int...) -> Int { 0 }
+}
+
+@UnavailableBuilderWithPartialBlock // expected-error {{'UnavailableBuilderWithPartialBlock' is unavailable}}
+func testUnavailableBuilderWithPartialBlock() -> Int {
+  // The builder itself is unavailable, so we can still opt for buildPartialBlock.
+  2 // expected-error {{cannot convert value of type 'Int' to expected argument type 'String'}}
+  3
+}

test/StringProcessing/Sema/regex_builder_unavailable.swift

@@ -0,0 +1,28 @@
+// RUN: %target-typecheck-verify-swift -enable-bare-slash-regex -target %target-cpu-apple-macosx12.0
+
+// REQUIRES: OS=macosx
+
+import RegexBuilder
+
+// rdar://97533700 – Make sure we can emit an availability diagnostic here.
+
+let _ = Regex { // expected-error {{'Regex' is only available in macOS 13.0 or newer}}
+  // expected-error@-1 {{'init(_:)' is only available in macOS 13.0 or newer}}
+  // expected-note@-2 2{{add 'if #available' version check}}
+
+  Capture {} // expected-error {{'Capture' is only available in macOS 13.0 or newer}}
+  // expected-error@-1 {{'init(_:)' is only available in macOS 13.0 or newer}}
+  // expected-note@-2 2{{add 'if #available' version check}}
+}
+
+let _ = Regex { // expected-error {{'Regex' is only available in macOS 13.0 or newer}}
+  // expected-error@-1 {{'init(_:)' is only available in macOS 13.0 or newer}}
+  // expected-error@-2 {{'buildPartialBlock(accumulated:next:)' is only available in macOS 13.0 or newer}}
+  // expected-note@-3 3{{add 'if #available' version check}}
+
+  /abc/ // expected-error {{'Regex' is only available in macOS 13.0 or newer}}
+  // expected-note@-1 {{add 'if #available' version check}}
+
+  /def/ // expected-error {{'Regex' is only available in macOS 13.0 or newer}}
+  // expected-note@-1 {{add 'if #available' version check}}
+}