AST: Unique CustomAvailabilityDomain instances.

Store `CustomAvailabilityDomain` instances in a folding set on `ASTContext`.
This instances of custom domains to be created without needing to cache them in
disparate locations.

Author: tshortli

include/swift/AST/AvailabilityDomain.h

@@ -100,7 +100,7 @@ class AvailabilityDomain final {
   using InlineDomainPtr =
       llvm::PointerEmbeddedInt<uint32_t, InlineDomain::ReprBits>;
   using Storage =
-      llvm::PointerUnion<CustomAvailabilityDomain *, InlineDomainPtr>;
+      llvm::PointerUnion<const CustomAvailabilityDomain *, InlineDomainPtr>;
   Storage storage;
 
   AvailabilityDomain(Kind kind)
@@ -111,7 +111,8 @@ class AvailabilityDomain final {
   AvailabilityDomain(PlatformKind platform)
       : storage(InlineDomain(Kind::Platform, platform).asInteger()) {};
 
-  AvailabilityDomain(CustomAvailabilityDomain *domain) : storage(domain) {};
+  AvailabilityDomain(const CustomAvailabilityDomain *domain)
+      : storage(domain) {};
 
   AvailabilityDomain(Storage storage) : storage(storage) {};
 
@@ -122,9 +123,9 @@ class AvailabilityDomain final {
                : std::nullopt;
   }
 
-  CustomAvailabilityDomain *getCustomDomain() const {
+  const CustomAvailabilityDomain *getCustomDomain() const {
     ASSERT(isCustom());
-    return storage.get<CustomAvailabilityDomain *>();
+    return storage.get<const CustomAvailabilityDomain *>();
   }
 
 public:
@@ -152,7 +153,7 @@ class AvailabilityDomain final {
     return AvailabilityDomain(Kind::Embedded);
   }
 
-  static AvailabilityDomain forCustom(CustomAvailabilityDomain *domain) {
+  static AvailabilityDomain forCustom(const CustomAvailabilityDomain *domain) {
     return AvailabilityDomain(domain);
   }
 
@@ -294,7 +295,7 @@ struct StableAvailabilityDomainComparator {
 };
 
 /// Represents an availability domain that has been defined in a module.
-class CustomAvailabilityDomain : public ASTAllocated<CustomAvailabilityDomain> {
+class CustomAvailabilityDomain : public llvm::FoldingSetNode {
 public:
   enum class Kind {
     /// A domain that is known to be enabled at compile time.
@@ -313,12 +314,20 @@ class CustomAvailabilityDomain : public ASTAllocated<CustomAvailabilityDomain> {
   CustomAvailabilityDomain(Identifier name, ModuleDecl *mod, Kind kind);
 
 public:
-  static CustomAvailabilityDomain *create(const ASTContext &ctx, StringRef name,
-                                          ModuleDecl *mod, Kind kind);
+  static const CustomAvailabilityDomain *get(StringRef name, ModuleDecl *mod,
+                                             Kind kind, const ASTContext &ctx);
 
   Identifier getName() const { return name; }
   Kind getKind() const { return kind; }
   ModuleDecl *getModule() const { return mod; }
+
+  /// Uniquing for `ASTContext`.
+  static void Profile(llvm::FoldingSetNodeID &ID, Identifier name,
+                      ModuleDecl *mod, Kind kind);
+
+  void Profile(llvm::FoldingSetNodeID &ID) const {
+    Profile(ID, name, mod, kind);
+  }
 };
 
 /// Represents either a resolved availability domain or an identifier written

include/swift/AST/Module.h

@@ -351,7 +351,7 @@ class ModuleDecl
   bool BypassResilience = false;
 
   using AvailabilityDomainMap =
-      llvm::SmallDenseMap<Identifier, CustomAvailabilityDomain *>;
+      llvm::SmallDenseMap<Identifier, const CustomAvailabilityDomain *>;
   AvailabilityDomainMap AvailabilityDomains;
 
 public:

lib/AST/ASTContext.cpp

@@ -637,6 +637,9 @@ struct ASTContext::Implementation {
   /// The set of unique AvailabilityContexts (uniqued by their storage).
   llvm::FoldingSet<AvailabilityContext::Storage> AvailabilityContexts;
 
+  /// The set of unique custom availability domains.
+  llvm::FoldingSet<CustomAvailabilityDomain> CustomAvailabilityDomains;
+
   /// A cache of information about whether particular nominal types
   /// are representable in a foreign language.
   llvm::DenseMap<NominalTypeDecl *, ForeignRepresentationInfo>
@@ -886,6 +889,7 @@ void ASTContext::Implementation::dump(llvm::raw_ostream &os) const {
   SIZE(NormalConformances);
   SIZE(SelfConformances);
   SIZE(AvailabilityContexts);
+  SIZE(CustomAvailabilityDomains);
   SIZE_AND_BYTES(ForeignRepresentableCache);
   SIZE(SearchPathsSet);
 
@@ -5738,6 +5742,27 @@ const AvailabilityContext::Storage *AvailabilityContext::Storage::get(
   return newNode;
 }
 
+const CustomAvailabilityDomain *
+CustomAvailabilityDomain::get(StringRef name, ModuleDecl *mod, Kind kind,
+                              const ASTContext &ctx) {
+  auto identifier = ctx.getIdentifier(name);
+  llvm::FoldingSetNodeID id;
+  CustomAvailabilityDomain::Profile(id, identifier, mod, kind);
+
+  auto &foldingSet = ctx.getImpl().CustomAvailabilityDomains;
+  void *insertPos;
+  auto *existing = foldingSet.FindNodeOrInsertPos(id, insertPos);
+  if (existing)
+    return existing;
+
+  void *mem = ctx.Allocate(sizeof(CustomAvailabilityDomain),
+                           alignof(CustomAvailabilityDomain));
+  auto *newNode = ::new (mem) CustomAvailabilityDomain(identifier, mod, kind);
+  foldingSet.InsertNode(newNode, insertPos);
+
+  return newNode;
+}
+
 void GenericSignatureImpl::Profile(llvm::FoldingSetNodeID &ID,
                               ArrayRef<GenericTypeParamType *> genericParams,
                               ArrayRef<Requirement> requirements) {

lib/AST/Availability.cpp

@@ -738,7 +738,6 @@ Decl::getAvailableAttrForPlatformIntroduction(bool checkExtension) const {
 }
 
 AvailabilityRange AvailabilityInference::availableRange(const Decl *D) {
-  // ALLANXXX
   if (auto attr = D->getAvailableAttrForPlatformIntroduction())
     return attr->getIntroducedRange(D->getASTContext())
         .value_or(AvailabilityRange::alwaysAvailable());

lib/AST/AvailabilityDomain.cpp

@@ -280,10 +280,12 @@ CustomAvailabilityDomain::CustomAvailabilityDomain(Identifier name,
   ASSERT(mod);
 }
 
-CustomAvailabilityDomain *
-CustomAvailabilityDomain::create(const ASTContext &ctx, StringRef name,
-                                 ModuleDecl *mod, Kind kind) {
-  return new (ctx) CustomAvailabilityDomain(ctx.getIdentifier(name), mod, kind);
+void CustomAvailabilityDomain::Profile(llvm::FoldingSetNodeID &ID,
+                                       Identifier name, ModuleDecl *mod,
+                                       Kind kind) {
+  ID.AddPointer(name.getAsOpaquePointer());
+  ID.AddPointer(mod);
+  ID.AddInteger(static_cast<unsigned>(kind));
 }
 
 static std::optional<AvailabilityDomain>

lib/Frontend/Frontend.cpp

@@ -1390,11 +1390,11 @@ bool CompilerInstance::createFilesForMainModule(
 static void configureAvailabilityDomains(const ASTContext &ctx,
                                          const FrontendOptions &opts,
                                          ModuleDecl *mainModule) {
-  llvm::SmallDenseMap<Identifier, CustomAvailabilityDomain *> domainMap;
+  llvm::SmallDenseMap<Identifier, const CustomAvailabilityDomain *> domainMap;
   auto createAndInsertDomain = [&](const std::string &name,
                                    CustomAvailabilityDomain::Kind kind) {
-    auto *domain = CustomAvailabilityDomain::create(
-        ctx, name, mainModule, CustomAvailabilityDomain::Kind::Enabled);
+    auto *domain = CustomAvailabilityDomain::get(
+        name, mainModule, CustomAvailabilityDomain::Kind::Enabled, ctx);
     bool inserted = domainMap.insert({domain->getName(), domain}).second;
     ASSERT(inserted); // Domains must be unique.
   };