[CS] Separate availability checking from SupportedOps cache

Previously we would cache the result of the first
query, with any further query of
`ResultBuilder::supports` ignoring the
`checkAvailability` parameter. Separate out the
availability checking such that we compute the
information up front, and adjust the result
depending on `checkAvailability`.

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;
 

lib/Sema/BuilderTransform.cpp

@@ -2770,11 +2770,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),
@@ -2793,17 +2804,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;
     }
@@ -2812,7 +2818,18 @@ bool TypeChecker::typeSupportsBuilderOp(
   if (allResults)
     allResults->append(foundDecls.begin(), foundDecls.end());
 
-  return foundMatch;
+  if (!foundMatch) {
+    return foundUnavailable ? ResultBuilderOpSupport::Unavailable
+                            : ResultBuilderOpSupport::Unsupported;
+  }
+  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) {
@@ -2976,13 +2993,13 @@ bool ResultBuilder::supports(Identifier fnBaseName,
                              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