Merge pull request #79782 from tshortli/semantic-availability-spec-request

AST: Request-ify retrieval of `SemanticAvailabilitySpec`

Author: tshortli

include/swift/AST/AvailabilityDomain.h

@@ -330,6 +330,12 @@ class AvailabilityDomainOrIdentifier {
   std::optional<AvailabilityDomain>
   lookUpInDeclContext(SourceLoc loc, const DeclContext *declContext) const;
 
+  void setResolved(std::optional<AvailabilityDomain> domain) {
+    if (domain)
+      storage.setPointer(*domain);
+    storage.setInt(true);
+  }
+
 public:
   AvailabilityDomainOrIdentifier(Identifier identifier)
       : storage(identifier) {};
@@ -359,18 +365,19 @@ class AvailabilityDomainOrIdentifier {
     return std::nullopt;
   }
 
+  /// Returns true if either a resolved domain is available or if the attempt
+  /// to look up the domain from the identifier was unsuccessful.
+  bool isResolved() const { return storage.getInt() || isDomain(); }
+
   std::optional<AvailabilityDomain>
   resolveInDeclContext(SourceLoc loc, const DeclContext *declContext) {
     // Return the domain directly if already resolved.
-    if (storage.getInt() || isDomain())
+    if (isResolved())
       return getAsDomain();
 
     // Look up the domain and cache the result.
     auto result = lookUpInDeclContext(loc, declContext);
-    if (result)
-      storage.setPointer(*result);
-    storage.setInt(true);
-
+    setResolved(result);
     return result;
   }
 

include/swift/AST/AvailabilitySpec.h

@@ -50,11 +50,13 @@ class AvailabilitySpec : public ASTAllocated<AvailabilitySpec> {
   // Location of the availability macro expanded to create this spec.
   SourceLoc MacroLoc;
 
+  unsigned IsInvalid : 1;
+
   AvailabilitySpec(AvailabilityDomainOrIdentifier DomainOrIdentifier,
                    SourceRange SrcRange, llvm::VersionTuple Version,
                    SourceLoc VersionStartLoc)
       : DomainOrIdentifier(DomainOrIdentifier), SrcRange(SrcRange),
-        Version(Version), VersionStartLoc(VersionStartLoc) {}
+        Version(Version), VersionStartLoc(VersionStartLoc), IsInvalid(false) {}
 
 public:
   /// Creates a wildcard availability specification that guards execution
@@ -91,6 +93,9 @@ class AvailabilitySpec : public ASTAllocated<AvailabilitySpec> {
   SourceRange getSourceRange() const { return SrcRange; }
   SourceLoc getStartLoc() const { return SrcRange.Start; }
 
+  /// Returns true if this spec did not type-check successfully.
+  bool isInvalid() const { return IsInvalid; }
+
   AvailabilityDomainOrIdentifier getDomainOrIdentifier() const {
     return DomainOrIdentifier;
   }
@@ -115,6 +120,14 @@ class AvailabilitySpec : public ASTAllocated<AvailabilitySpec> {
   void setMacroLoc(SourceLoc loc) { MacroLoc = loc; }
 
   void print(llvm::raw_ostream &os) const;
+
+private:
+  friend class SemanticAvailabilitySpecRequest;
+
+  std::optional<AvailabilityDomain>
+  resolveInDeclContext(const DeclContext *declContext);
+
+  void setInvalid() { IsInvalid = true; }
 };
 
 inline void simple_display(llvm::raw_ostream &os,

include/swift/AST/TypeCheckRequests.h

@@ -20,6 +20,7 @@
 #include "swift/AST/ASTTypeIDs.h"
 #include "swift/AST/ActorIsolation.h"
 #include "swift/AST/AnyFunctionRef.h"
+#include "swift/AST/AvailabilitySpec.h"
 #include "swift/AST/CatchNode.h"
 #include "swift/AST/Effects.h"
 #include "swift/AST/Evaluator.h"
@@ -5228,6 +5229,28 @@ class SemanticAvailableAttrRequest
   void cacheResult(std::optional<SemanticAvailableAttr> value) const;
 };
 
+class SemanticAvailabilitySpecRequest
+    : public SimpleRequest<SemanticAvailabilitySpecRequest,
+                           std::optional<SemanticAvailabilitySpec>(
+                               const AvailabilitySpec *, const DeclContext *),
+                           RequestFlags::SeparatelyCached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  std::optional<SemanticAvailabilitySpec>
+  evaluate(Evaluator &evaluator, const AvailabilitySpec *spec,
+           const DeclContext *declContext) const;
+
+public:
+  bool isCached() const { return true; }
+  std::optional<std::optional<SemanticAvailabilitySpec>>
+  getCachedResult() const;
+  void cacheResult(std::optional<SemanticAvailabilitySpec> value) const;
+};
+
 #define SWIFT_TYPEID_ZONE TypeChecker
 #define SWIFT_TYPEID_HEADER "swift/AST/TypeCheckerTypeIDZone.def"
 #include "swift/Basic/DefineTypeIDZone.h"

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -616,3 +616,7 @@ SWIFT_REQUEST(TypeChecker, SemanticAvailableAttrRequest,
               std::optional<SemanticAvailableAttr>
               (const AvailableAttr *, const Decl *),
               SeparatelyCached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, SemanticAvailabilitySpecRequest,
+              std::optional<SemanticAvailabilitySpec>
+              (const AvailabilitySpec *, const DeclContext *),
+              SeparatelyCached, NoLocationInfo)

lib/AST/AvailabilitySpec.cpp

@@ -18,6 +18,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/AvailabilityDomain.h"
 #include "swift/AST/DiagnosticsParse.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace swift;
@@ -70,16 +71,21 @@ void AvailabilitySpec::print(llvm::raw_ostream &os) const {
     os << " " << getRawVersion().getAsString();
 }
 
+std::optional<AvailabilityDomain>
+AvailabilitySpec::resolveInDeclContext(const DeclContext *declContext) {
+  auto domainOrIdentifier = getDomainOrIdentifier();
+  auto result =
+      domainOrIdentifier.resolveInDeclContext(getStartLoc(), declContext);
+  DomainOrIdentifier = domainOrIdentifier;
+  return result;
+}
+
 std::optional<SemanticAvailabilitySpec>
 AvailabilitySpec::getSemanticAvailabilitySpec(
     const DeclContext *declContext) const {
-  AvailabilitySpec *mutableThis = const_cast<AvailabilitySpec *>(this);
-  auto domain = mutableThis->DomainOrIdentifier.resolveInDeclContext(
-      getStartLoc(), declContext);
-
-  if (domain)
-    return SemanticAvailabilitySpec(this);
-  return std::nullopt;
+  return evaluateOrDefault(declContext->getASTContext().evaluator,
+                           SemanticAvailabilitySpecRequest{this, declContext},
+                           std::nullopt);
 }
 
 llvm::VersionTuple SemanticAvailabilitySpec::getVersion() const {
@@ -105,3 +111,33 @@ SemanticAvailabilitySpecs::Filter::operator()(
     return semanticSpec;
   return std::nullopt;
 }
+
+std::optional<SemanticAvailabilitySpec>
+SemanticAvailabilitySpecRequest::evaluate(
+    Evaluator &evaluator, const AvailabilitySpec *spec,
+    const DeclContext *declContext) const {
+  AvailabilitySpec *mutableSpec = const_cast<AvailabilitySpec *>(spec);
+  if (mutableSpec->resolveInDeclContext(declContext).has_value())
+    return SemanticAvailabilitySpec(spec);
+  return std::nullopt;
+}
+
+std::optional<std::optional<SemanticAvailabilitySpec>>
+SemanticAvailabilitySpecRequest::getCachedResult() const {
+  auto *spec = std::get<0>(getStorage());
+  auto domainOrIdentifier = spec->getDomainOrIdentifier();
+  if (!domainOrIdentifier.isResolved())
+    return {};
+
+  if (spec->isInvalid())
+    return {std::nullopt};
+
+  return SemanticAvailabilitySpec(spec);
+}
+
+void SemanticAvailabilitySpecRequest::cacheResult(
+    std::optional<SemanticAvailabilitySpec> value) const {
+  auto *mutableSpec = const_cast<AvailabilitySpec *>(std::get<0>(getStorage()));
+  if (!value)
+    mutableSpec->setInvalid();
+}