AST: Centralize AvailabilityDomain lookup.

Implement lookup of availability domains for identifiers on
`AvailabilityDomainOrIdentifier`. Add a bit to that type which represents
whether or not lookup has already been attempted. This allows both
`AvailableAttr` and `AvailabilitySpec` to share a common implementation of
domain lookup.

Author: tshortli

include/swift/AST/Attr.h

@@ -746,6 +746,7 @@ class AvailableAttr : public DeclAttribute {
 
 private:
   friend class SemanticAvailableAttr;
+  friend class SemanticAvailableAttrRequest;
 
   AvailabilityDomainOrIdentifier DomainOrIdentifier;
   const SourceLoc DomainLoc;
@@ -769,13 +770,6 @@ class AvailableAttr : public DeclAttribute {
     return DomainOrIdentifier;
   }
 
-  /// Returns the `AvailabilityDomain` associated with the attribute, or
-  /// `std::nullopt` if it has either not yet been resolved or could not be
-  /// resolved successfully.
-  std::optional<AvailabilityDomain> getCachedDomain() const {
-    return DomainOrIdentifier.getAsDomain();
-  }
-
   SourceLoc getDomainLoc() const { return DomainLoc; }
 
   /// Returns the parsed version for `introduced:`.
@@ -918,11 +912,6 @@ class AvailableAttr : public DeclAttribute {
 private:
   friend class SemanticAvailableAttrRequest;
 
-  void setCachedDomain(AvailabilityDomain domain) {
-    assert(!DomainOrIdentifier.isDomain());
-    DomainOrIdentifier.setDomain(domain);
-  }
-
   bool hasComputedSemanticAttr() const {
     return Bits.AvailableAttr.HasComputedSemanticAttr;
   }
@@ -3305,12 +3294,13 @@ class SemanticAvailableAttr final {
 public:
   SemanticAvailableAttr(const AvailableAttr *attr) : attr(attr) {
     assert(attr);
-    assert(attr->hasCachedDomain());
+    bool hasDomain = attr->getDomainOrIdentifier().isDomain();
+    assert(hasDomain);
   }
 
   const AvailableAttr *getParsedAttr() const { return attr; }
   const AvailabilityDomain getDomain() const {
-    return attr->getCachedDomain().value();
+    return attr->getDomainOrIdentifier().getAsDomain().value();
   }
 
   /// The version tuple for the `introduced:` component.

include/swift/AST/AvailabilityDomain.h

@@ -23,6 +23,7 @@
 #include "swift/AST/PlatformKind.h"
 #include "swift/Basic/Assertions.h"
 #include "swift/Basic/LLVM.h"
+#include "swift/Basic/SourceLoc.h"
 #include "llvm/ADT/FoldingSet.h"
 #include "llvm/ADT/PointerEmbeddedInt.h"
 #include "llvm/ADT/PointerUnion.h"
@@ -294,42 +295,65 @@ class CustomAvailabilityDomain : public ASTAllocated<CustomAvailabilityDomain> {
 class AvailabilityDomainOrIdentifier {
   friend struct llvm::PointerLikeTypeTraits<AvailabilityDomainOrIdentifier>;
 
-  using Storage = llvm::PointerUnion<AvailabilityDomain, Identifier>;
+  using DomainOrIdentifier = llvm::PointerUnion<AvailabilityDomain, Identifier>;
+
+  /// Stores an extra bit representing whether the domain has been resolved.
+  using Storage = llvm::PointerIntPair<DomainOrIdentifier, 1, bool>;
   Storage storage;
 
   AvailabilityDomainOrIdentifier(Storage storage) : storage(storage) {}
 
+  static AvailabilityDomainOrIdentifier fromOpaque(void *opaque) {
+    return AvailabilityDomainOrIdentifier(Storage::getFromOpaqueValue(opaque));
+  }
+
+  std::optional<AvailabilityDomain>
+  lookUpInDeclContext(SourceLoc loc, const DeclContext *declContext) const;
+
 public:
   AvailabilityDomainOrIdentifier(Identifier identifier)
       : storage(identifier) {};
   AvailabilityDomainOrIdentifier(AvailabilityDomain domain)
       : storage(domain) {};
 
-  static AvailabilityDomainOrIdentifier fromOpaque(void *opaque) {
-    return AvailabilityDomainOrIdentifier(Storage::getFromOpaqueValue(opaque));
+  bool isDomain() const {
+    return storage.getPointer().is<AvailabilityDomain>();
   }
-
-  bool isDomain() const { return storage.is<AvailabilityDomain>(); }
-  bool isIdentifier() const { return storage.is<Identifier>(); }
+  bool isIdentifier() const { return storage.getPointer().is<Identifier>(); }
 
   /// Overwrites the existing domain or identifier with the given domain.
   void setDomain(AvailabilityDomain domain) { storage = Storage(domain); }
 
   /// Returns the resolved domain, or `std::nullopt` if there isn't one.
   std::optional<AvailabilityDomain> getAsDomain() const {
     if (isDomain())
-      return storage.get<AvailabilityDomain>();
+      return storage.getPointer().get<AvailabilityDomain>();
     return std::nullopt;
   }
 
   /// Returns the unresolved identifier, or `std::nullopt` if the domain has
   /// been resolved.
   std::optional<Identifier> getAsIdentifier() const {
     if (isIdentifier())
-      return storage.get<Identifier>();
+      return storage.getPointer().get<Identifier>();
     return std::nullopt;
   }
 
+  std::optional<AvailabilityDomain>
+  resolveInDeclContext(SourceLoc loc, const DeclContext *declContext) {
+    // Return the domain directly if already resolved.
+    if (storage.getInt() || isDomain())
+      return getAsDomain();
+
+    // Look up the domain and cache the result.
+    auto result = lookUpInDeclContext(loc, declContext);
+    if (result)
+      storage.setPointer(*result);
+    storage.setInt(true);
+
+    return result;
+  }
+
   /// Creates a new `AvailabilityDomainOrIdentifier`, defensively copying
   /// members of the original into the given `ASTContext` in case it is
   /// different than the context that the original was created for.

include/swift/AST/AvailabilitySpec.h

@@ -35,8 +35,7 @@ class SemanticAvailabilitySpec;
 /// The root class for specifications of API availability in availability
 /// queries.
 class AvailabilitySpec : public ASTAllocated<AvailabilitySpec> {
-  using DomainStorage = llvm::PointerIntPair<AvailabilityDomainOrIdentifier, 1>;
-  DomainStorage Storage;
+  AvailabilityDomainOrIdentifier DomainOrIdentifier;
 
   /// The range of the entire spec, including the version if there is one.
   SourceRange SrcRange;
@@ -51,11 +50,11 @@ class AvailabilitySpec : public ASTAllocated<AvailabilitySpec> {
   // Location of the availability macro expanded to create this spec.
   SourceLoc MacroLoc;
 
-  AvailabilitySpec(DomainStorage Storage, SourceRange SrcRange,
-                   llvm::VersionTuple Version, SourceLoc VersionStartLoc)
-      : Storage(Storage), SrcRange(SrcRange), Version(Version),
-        VersionStartLoc(VersionStartLoc) {}
-
+  AvailabilitySpec(AvailabilityDomainOrIdentifier DomainOrIdentifier,
+                   SourceRange SrcRange, llvm::VersionTuple Version,
+                   SourceLoc VersionStartLoc)
+      : DomainOrIdentifier(DomainOrIdentifier), SrcRange(SrcRange),
+        Version(Version), VersionStartLoc(VersionStartLoc) {}
 
 public:
   /// Creates a wildcard availability specification that guards execution
@@ -99,7 +98,7 @@ class AvailabilitySpec : public ASTAllocated<AvailabilitySpec> {
   }
 
   AvailabilityDomainOrIdentifier getDomainOrIdentifier() const {
-    return Storage.getPointer();
+    return DomainOrIdentifier;
   }
 
   std::optional<AvailabilityDomain> getDomain() const {

include/swift/AST/DiagnosticsSema.def

@@ -6748,9 +6748,8 @@ WARNING(attr_availability_expected_version_spec, none,
        "expected 'introduced', 'deprecated', or 'obsoleted' in '%0' attribute "
        "for platform '%1'", (StringRef, StringRef))
 ERROR(attr_availability_requires_custom_availability, none,
-      "specifying '%0' in '%1' attribute requires "
-      "-enable-experimental-feature CustomAvailability",
-      (StringRef, const DeclAttribute))
+      "%0 requires -enable-experimental-feature CustomAvailability",
+      (Identifier))
 WARNING(attr_availability_unexpected_version,none,
        "unexpected version number in '%0' attribute for '%1'",
        (const DeclAttribute, StringRef))

lib/AST/Availability.cpp

@@ -797,69 +797,25 @@ bool AvailabilityInference::isAvailableAsSPI(const Decl *D) {
   return false;
 }
 
-static std::optional<AvailabilityDomain>
-getAvailabilityDomainForName(Identifier identifier,
-                             const DeclContext *declContext) {
-  if (auto builtinDomain = AvailabilityDomain::builtinDomainForString(
-          identifier.str(), declContext))
-    return builtinDomain;
-
-  auto &ctx = declContext->getASTContext();
-  if (auto customDomain =
-          ctx.MainModule->getAvailabilityDomainForIdentifier(identifier))
-    return customDomain;
-
-  return std::nullopt;
-}
-
 std::optional<SemanticAvailableAttr>
 SemanticAvailableAttrRequest::evaluate(swift::Evaluator &evaluator,
                                        const AvailableAttr *attr,
                                        const Decl *decl) const {
-  if (attr->hasCachedDomain())
+  if (attr->getDomainOrIdentifier().isDomain())
     return SemanticAvailableAttr(attr);
 
   auto &ctx = decl->getASTContext();
   auto &diags = ctx.Diags;
   auto attrLoc = attr->getLocation();
   auto attrName = attr->getAttrName();
   auto domainLoc = attr->getDomainLoc();
+  auto declContext = decl->getInnermostDeclContext();
   auto mutableAttr = const_cast<AvailableAttr *>(attr);
-  auto domain = attr->getCachedDomain();
-
-  if (!domain) {
-    auto domainIdentifier = attr->getDomainOrIdentifier().getAsIdentifier();
-    ASSERT(domainIdentifier);
-
-    // Attempt to resolve the domain specified for the attribute and diagnose
-    // if no domain is found.
-    auto declContext = decl->getInnermostDeclContext();
-    domain = getAvailabilityDomainForName(*domainIdentifier, declContext);
-    if (!domain) {
-      auto domainString = domainIdentifier->str();
-      if (auto suggestion = closestCorrectedPlatformString(domainString)) {
-        diags
-            .diagnose(domainLoc, diag::attr_availability_suggest_platform,
-                      domainString, attrName, *suggestion)
-            .fixItReplace(SourceRange(domainLoc), *suggestion);
-      } else {
-        diags.diagnose(attrLoc, diag::attr_availability_unknown_platform,
-                       domainString, attrName);
-      }
-      return std::nullopt;
-    }
-
-    if (domain->isCustom() &&
-        !ctx.LangOpts.hasFeature(Feature::CustomAvailability) &&
-        !declContext->isInSwiftinterface()) {
-      diags.diagnose(domainLoc,
-                     diag::attr_availability_requires_custom_availability,
-                     domain->getNameForAttributePrinting(), attr);
-      return std::nullopt;
-    }
+  auto domain = mutableAttr->DomainOrIdentifier.resolveInDeclContext(
+      domainLoc, declContext);
 
-    mutableAttr->setCachedDomain(*domain);
-  }
+  if (!domain)
+    return std::nullopt;
 
   auto domainName = domain->getNameForAttributePrinting();
   auto semanticAttr = SemanticAvailableAttr(attr);

lib/AST/AvailabilityDomain.cpp

@@ -13,6 +13,9 @@
 #include "swift/AST/AvailabilityDomain.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/Decl.h"
+#include "swift/AST/DiagnosticsParse.h"
+#include "swift/AST/DiagnosticsSema.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/Basic/Assertions.h"
 #include "llvm/ADT/StringSwitch.h"
 
@@ -201,6 +204,57 @@ CustomAvailabilityDomain::create(const ASTContext &ctx, StringRef name,
   return new (ctx) CustomAvailabilityDomain(ctx.getIdentifier(name), mod, kind);
 }
 
+static std::optional<AvailabilityDomain>
+getAvailabilityDomainForName(Identifier identifier,
+                             const DeclContext *declContext) {
+  if (auto builtinDomain = AvailabilityDomain::builtinDomainForString(
+          identifier.str(), declContext))
+    return builtinDomain;
+
+  auto &ctx = declContext->getASTContext();
+  if (auto customDomain =
+          ctx.MainModule->getAvailabilityDomainForIdentifier(identifier))
+    return customDomain;
+
+  return std::nullopt;
+}
+
+std::optional<AvailabilityDomain>
+AvailabilityDomainOrIdentifier::lookUpInDeclContext(
+    SourceLoc loc, const DeclContext *declContext) const {
+  DEBUG_ASSERT(isIdentifier());
+
+  auto &ctx = declContext->getASTContext();
+  auto &diags = ctx.Diags;
+  std::optional<AvailabilityDomain> domain;
+  auto identifier = getAsIdentifier().value();
+  domain = getAvailabilityDomainForName(identifier, declContext);
+
+  if (!domain) {
+    auto domainString = identifier.str();
+    if (auto suggestion = closestCorrectedPlatformString(domainString)) {
+      diags
+          .diagnose(loc, diag::avail_query_suggest_platform_name, identifier,
+                    *suggestion)
+          .fixItReplace(SourceRange(loc), *suggestion);
+    } else {
+      diags.diagnose(loc, diag::avail_query_unrecognized_platform_name,
+                     identifier);
+    }
+    return std::nullopt;
+  }
+
+  if (domain->isCustom() &&
+      !ctx.LangOpts.hasFeature(Feature::CustomAvailability) &&
+      !declContext->isInSwiftinterface()) {
+    diags.diagnose(loc, diag::attr_availability_requires_custom_availability,
+                   identifier);
+    return std::nullopt;
+  }
+
+  return domain;
+}
+
 AvailabilityDomainOrIdentifier
 AvailabilityDomainOrIdentifier::copy(ASTContext &ctx) const {
   if (auto identifier = getAsIdentifier())

lib/AST/AvailabilitySpec.cpp

@@ -23,10 +23,9 @@ using namespace swift;
 
 AvailabilitySpec *AvailabilitySpec::createWildcard(ASTContext &ctx,
                                                    SourceLoc starLoc) {
-  return new (ctx)
-      AvailabilitySpec({AvailabilityDomain::forUniversal(), true}, starLoc,
-                       /*Version=*/{},
-                       /*VersionStartLoc=*/{});
+  return new (ctx) AvailabilitySpec(AvailabilityDomain::forUniversal(), starLoc,
+                                    /*Version=*/{},
+                                    /*VersionStartLoc=*/{});
 }
 
 AvailabilitySpec *AvailabilitySpec::createForDomain(ASTContext &ctx,
@@ -35,23 +34,22 @@ AvailabilitySpec *AvailabilitySpec::createForDomain(ASTContext &ctx,
                                                     llvm::VersionTuple version,
                                                     SourceRange versionRange) {
   DEBUG_ASSERT(!version.empty());
-  return new (ctx)
-      AvailabilitySpec({domain, true}, SourceRange(loc, versionRange.End),
-                       version, versionRange.Start);
+  return new (ctx) AvailabilitySpec(domain, SourceRange(loc, versionRange.End),
+                                    version, versionRange.Start);
 }
 
 AvailabilitySpec *AvailabilitySpec::createForUnknownDomain(
     ASTContext &ctx, Identifier domainIdentifier, SourceLoc loc,
     llvm::VersionTuple version, SourceRange versionRange) {
-  return new (ctx) AvailabilitySpec({domainIdentifier, true},
-                                    SourceRange(loc, versionRange.End), version,
-                                    versionRange.Start);
+  DEBUG_ASSERT(!version.empty());
+  return new (ctx)
+      AvailabilitySpec(domainIdentifier, SourceRange(loc, versionRange.End),
+                       version, versionRange.Start);
 }
 
 AvailabilitySpec *AvailabilitySpec::clone(ASTContext &ctx) const {
-  return new (ctx)
-      AvailabilitySpec({getDomainOrIdentifier().copy(ctx), true},
-                       SrcRange, Version, VersionStartLoc);
+  return new (ctx) AvailabilitySpec(getDomainOrIdentifier().copy(ctx), SrcRange,
+                                    Version, VersionStartLoc);
 }
 
 void AvailabilitySpec::print(llvm::raw_ostream &os) const {

lib/AST/TypeCheckRequests.cpp

@@ -2725,7 +2725,7 @@ SemanticAvailableAttrRequest::getCachedResult() const {
   if (!attr->hasComputedSemanticAttr())
     return std::nullopt;
 
-  if (!attr->hasCachedDomain()) {
+  if (!attr->getDomainOrIdentifier().isDomain()) {
     return std::optional<SemanticAvailableAttr>{};
   }