AST: Add printing utilities for AvailabilityDomain and AvailabilitySpec.

tshortli · tshortli · commit 8ebf2aa7e6bc · 2025-02-25T19:45:39.000-08:00
diff --git a/include/swift/AST/Attr.h b/include/swift/AST/Attr.h
@@ -765,20 +765,17 @@ class AvailableAttr : public DeclAttribute {
   /// has been resolved successfully.
   bool hasCachedDomain() const { return DomainOrIdentifier.isDomain(); }
 
+  AvailabilityDomainOrIdentifier getDomainOrIdentifier() const {
+    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();
   }
 
-  /// 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();
-  }
-
   SourceLoc getDomainLoc() const { return DomainLoc; }
 
   /// Returns the parsed version for `introduced:`.
diff --git a/include/swift/AST/AvailabilityDomain.h b/include/swift/AST/AvailabilityDomain.h
@@ -248,12 +248,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:
@@ -327,6 +334,8 @@ class AvailabilityDomainOrIdentifier {
   /// 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
diff --git a/include/swift/AST/AvailabilitySpec.h b/include/swift/AST/AvailabilitySpec.h
@@ -56,9 +56,6 @@ class AvailabilitySpec : public ASTAllocated<AvailabilitySpec> {
       : Storage(Storage), SrcRange(SrcRange), Version(Version),
         VersionStartLoc(VersionStartLoc) {}
 
-  AvailabilityDomainOrIdentifier getDomainOrIdentifier() const {
-    return Storage.getPointer();
-  }
 
 public:
   /// Creates a wildcard availability specification that guards execution
@@ -101,6 +98,10 @@ class AvailabilitySpec : public ASTAllocated<AvailabilitySpec> {
     return false;
   }
 
+  AvailabilityDomainOrIdentifier getDomainOrIdentifier() const {
+    return Storage.getPointer();
+  }
+
   std::optional<AvailabilityDomain> getDomain() const {
     return getDomainOrIdentifier().getAsDomain();
   }
@@ -123,8 +124,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 +161,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
diff --git a/lib/AST/ASTDumper.cpp b/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:
diff --git a/lib/AST/Availability.cpp b/lib/AST/Availability.cpp
@@ -828,7 +828,7 @@ SemanticAvailableAttrRequest::evaluate(swift::Evaluator &evaluator,
   auto domain = attr->getCachedDomain();
 
   if (!domain) {
-    auto domainIdentifier = attr->getDomainIdentifier();
+    auto domainIdentifier = attr->getDomainOrIdentifier().getAsIdentifier();
     ASSERT(domainIdentifier);
 
     // Attempt to resolve the domain specified for the attribute and diagnose
diff --git a/lib/AST/AvailabilityDomain.cpp b/lib/AST/AvailabilityDomain.cpp
@@ -168,6 +168,10 @@ AvailabilityDomain AvailabilityDomain::getABICompatibilityDomain() const {
   return *this;
 }
 
+void AvailabilityDomain::print(llvm::raw_ostream &os) const {
+  os << getNameForAttributePrinting();
+}
+
 AvailabilityDomain AvailabilityDomain::copy(ASTContext &ctx) const {
   switch (getKind()) {
   case Kind::Universal:
@@ -205,3 +209,10 @@ AvailabilityDomainOrIdentifier::copy(ASTContext &ctx) const {
   DEBUG_ASSERT(isDomain());
   return getAsDomain()->copy(ctx);
 }
+
+void AvailabilityDomainOrIdentifier::print(llvm::raw_ostream &os) const {
+  if (auto identifier = getAsIdentifier())
+    os << identifier->str();
+  else
+    getAsDomain()->print(os);
+}
diff --git a/lib/AST/AvailabilitySpec.cpp b/lib/AST/AvailabilitySpec.cpp
@@ -54,6 +54,13 @@ AvailabilitySpec *AvailabilitySpec::clone(ASTContext &ctx) const {
                        SrcRange, Version, VersionStartLoc);
 }
 
+void AvailabilitySpec::print(llvm::raw_ostream &os) const {
+  getDomainOrIdentifier().print(os);
+
+  if (!getRawVersion().empty())
+    os << " " << getRawVersion().getAsString();
+}
+
 llvm::VersionTuple SemanticAvailabilitySpec::getVersion() const {
   // For macOS Big Sur, we canonicalize 10.16 to 11.0 for compile-time
   // checking since clang canonicalizes availability markup. However, to
