AST: Centralize AvailabilityDomain lookup

include/swift/AST/Attr.h

@@ -165,14 +165,20 @@ class DeclAttribute : public AttributeBase {
       /// Whether this attribute was spelled `@_spi_available`.
       IsSPI : 1,
 
-      /// Whether this attribute belongs to a chain of adjacent `@available` attributes that were generated from a single attribute written in source using short form syntax e.g. (`@available(macOS 15, iOS 18, *)`).
+      /// Whether this attribute belongs to a chain of adjacent `@available`
+      /// attributes that were generated from a single attribute written in
+      /// source using short form syntax, e.g.
+      ///
+      ///     @available(macOS 15, iOS 18, *)
+      ///
       IsGroupMember : 1,
 
       /// Whether this attribute is the final one in its group.
       IsGroupTerminator : 1,
 
-      /// Whether this attribute's specification was followed by `, *` in source.
-      IsAdjacentToWildcard : 1
+      /// Whether any members of the group were written as a wildcard
+      /// specification (`*`) in source.
+      IsGroupedWithWildcard : 1
     );
 
     SWIFT_INLINE_BITFIELD(ClangImporterSynthesizedTypeAttr, DeclAttribute, 1,
@@ -740,6 +746,7 @@ class AvailableAttr : public DeclAttribute {
 
 private:
   friend class SemanticAvailableAttr;
+  friend class SemanticAvailableAttrRequest;
 
   AvailabilityDomainOrIdentifier DomainOrIdentifier;
   const SourceLoc DomainLoc;
@@ -759,18 +766,8 @@ class AvailableAttr : public DeclAttribute {
   /// has been resolved successfully.
   bool hasCachedDomain() const { return DomainOrIdentifier.isDomain(); }
 
-  /// 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();
-  }
-
-  /// If the attribute does not already have a cached `AvailabilityDomain`, this
-  /// returns the domain identifier that was written in source, from which an
-  /// `AvailabilityDomain` can be resolved.
-  std::optional<Identifier> getDomainIdentifier() const {
-    return DomainOrIdentifier.getAsIdentifier();
+  AvailabilityDomainOrIdentifier getDomainOrIdentifier() const {
+    return DomainOrIdentifier;
   }
 
   SourceLoc getDomainLoc() const { return DomainLoc; }
@@ -848,12 +845,13 @@ class AvailableAttr : public DeclAttribute {
   }
   void setIsGroupTerminator() { Bits.AvailableAttr.IsGroupTerminator = true; }
 
-  /// Whether this attribute's specification was followed by `, *` in source.
-  bool isAdjacentToWildcard() const {
-    return Bits.AvailableAttr.IsAdjacentToWildcard;
+  /// Whether any members of the group were written as a wildcard specification
+  /// (`*`) in source.
+  bool isGroupedWithWildcard() const {
+    return Bits.AvailableAttr.IsGroupedWithWildcard;
   }
-  void setIsAdjacentToWildcard() {
-    Bits.AvailableAttr.IsAdjacentToWildcard = true;
+  void setIsGroupedWithWildcard() {
+    Bits.AvailableAttr.IsGroupedWithWildcard = true;
   }
 
   /// Returns the kind of availability the attribute specifies.
@@ -914,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;
   }
@@ -3301,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"
@@ -248,12 +249,19 @@ class AvailabilityDomain final {
     ID.AddPointer(getOpaqueValue());
   }
 
+  void print(llvm::raw_ostream &os) const;
+
 private:
   friend class AvailabilityDomainOrIdentifier;
 
   AvailabilityDomain copy(ASTContext &ctx) const;
 };
 
+inline void simple_display(llvm::raw_ostream &os,
+                           const AvailabilityDomain &domain) {
+  domain.print(os);
+}
+
 /// Represents an availability domain that has been defined in a module.
 class CustomAvailabilityDomain : public ASTAllocated<CustomAvailabilityDomain> {
 public:
@@ -287,46 +295,71 @@ 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.
   AvailabilityDomainOrIdentifier copy(ASTContext &ctx) const;
+
+  void print(llvm::raw_ostream &os) const;
 };
 
 } // end namespace swift

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,14 +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) {}
-
-  AvailabilityDomainOrIdentifier getDomainOrIdentifier() const {
-    return Storage.getPointer();
-  }
+  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
@@ -101,6 +97,10 @@ class AvailabilitySpec : public ASTAllocated<AvailabilitySpec> {
     return false;
   }
 
+  AvailabilityDomainOrIdentifier getDomainOrIdentifier() const {
+    return DomainOrIdentifier;
+  }
+
   std::optional<AvailabilityDomain> getDomain() const {
     return getDomainOrIdentifier().getAsDomain();
   }
@@ -123,8 +123,15 @@ class AvailabilitySpec : public ASTAllocated<AvailabilitySpec> {
   // Location of the macro expanded to create this spec.
   SourceLoc getMacroLoc() const { return MacroLoc; }
   void setMacroLoc(SourceLoc loc) { MacroLoc = loc; }
+
+  void print(llvm::raw_ostream &os) const;
 };
 
+inline void simple_display(llvm::raw_ostream &os,
+                           const AvailabilitySpec *spec) {
+  spec->print(os);
+}
+
 /// The type-checked representation of `AvailabilitySpec` which guaranatees that
 /// the spec has a valid `AvailabilityDomain`.
 class SemanticAvailabilitySpec {
@@ -153,6 +160,8 @@ class SemanticAvailabilitySpec {
   // This is required to support beta versions of macOS Big Sur that
   // report 10.16 at run time.
   llvm::VersionTuple getRuntimeVersion() const { return spec->getRawVersion(); }
+
+  void print(llvm::raw_ostream &os) const { spec->print(os); }
 };
 
 /// Wraps an array of availability specs and provides an iterator for their

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/ASTDumper.cpp

@@ -1103,11 +1103,11 @@ namespace {
       printRecArbitrary(
           [&](Label label) {
             printHead("availability_spec", PatternColor, label);
-            StringRef domainName =
-                Spec->isWildcard()
-                    ? "*"
-                    : Spec->getDomain()->getNameForAttributePrinting();
-            printField(domainName, Label::always("domain"));
+            printFieldRaw(
+                [&](llvm::raw_ostream &OS) {
+                  Spec->getDomainOrIdentifier().print(OS);
+                },
+                Label::always("domain"));
             if (!Spec->getRawVersion().empty())
               printFieldRaw(
                   [&](llvm::raw_ostream &OS) { OS << Spec->getRawVersion(); },
@@ -4944,13 +4944,8 @@ class PrintAttribute : public AttributeVisitor<PrintAttribute, void, Label>,
   void visitAvailableAttr(AvailableAttr *Attr, Label label) {
     printCommon(Attr, "available_attr", label);
 
-    if (auto domain = Attr->getCachedDomain()) {
-      printField(domain->getNameForAttributePrinting(),
-                 Label::always("domain"));
-    } else {
-      printField(*Attr->getDomainIdentifier(),
-                 Label::always("domainIdentifier"));
-    }
+    printFieldRaw([&](auto &out) { Attr->getDomainOrIdentifier().print(out); },
+                  Label::always("domain"));
 
     switch (Attr->getKind()) {
     case swift::AvailableAttr::Kind::Default:

lib/AST/Attr.cpp

@@ -2146,7 +2146,7 @@ AvailableAttr::AvailableAttr(
   Bits.AvailableAttr.IsSPI = IsSPI;
   Bits.AvailableAttr.IsGroupMember = false;
   Bits.AvailableAttr.IsGroupTerminator = false;
-  Bits.AvailableAttr.IsAdjacentToWildcard = false;
+  Bits.AvailableAttr.IsGroupedWithWildcard = false;
 }
 
 AvailableAttr *AvailableAttr::createUniversallyUnavailable(ASTContext &C,