AST: Introduce a Custom availability domain kind.

Author: tshortli

include/swift/AST/AvailabilityDomain.h

@@ -18,6 +18,8 @@
 #ifndef SWIFT_AST_AVAILABILITY_DOMAIN_H
 #define SWIFT_AST_AVAILABILITY_DOMAIN_H
 
+#include "swift/AST/ASTAllocated.h"
+#include "swift/AST/Identifier.h"
 #include "swift/AST/PlatformKind.h"
 #include "swift/Basic/LLVM.h"
 #include "llvm/ADT/FoldingSet.h"
@@ -26,7 +28,9 @@
 
 namespace swift {
 class ASTContext;
+class CustomAvailabilityDomain;
 class DeclContext;
+class ModuleDecl;
 
 /// Represents a dimension of availability (e.g. macOS platform or Swift
 /// language mode).
@@ -48,6 +52,10 @@ class AvailabilityDomain final {
 
     /// Represents availability for a specific operating system platform.
     Platform,
+
+    /// Represents availability for an arbitrary domain that is defined at
+    /// compile time by a module.
+    Custom,
   };
 
 private:
@@ -86,13 +94,10 @@ class AvailabilityDomain final {
     PlatformKind getPlatform() { return platform; }
   };
 
-  /// This will eventually be a class storing information about externally
-  /// defined availability domains.
-  using ExternalDomain = void;
-
   using InlineDomainPtr =
       llvm::PointerEmbeddedInt<uint32_t, InlineDomain::ReprBits>;
-  using Storage = llvm::PointerUnion<ExternalDomain *, InlineDomainPtr>;
+  using Storage =
+      llvm::PointerUnion<CustomAvailabilityDomain *, InlineDomainPtr>;
   Storage storage;
 
   AvailabilityDomain(Kind kind)
@@ -103,6 +108,8 @@ class AvailabilityDomain final {
   AvailabilityDomain(PlatformKind platform)
       : storage(InlineDomain(Kind::Platform, platform).asInteger()) {};
 
+  AvailabilityDomain(CustomAvailabilityDomain *domain) : storage(domain) {};
+
   AvailabilityDomain(Storage storage) : storage(storage) {};
 
   static AvailabilityDomain fromOpaque(void *opaque) {
@@ -118,6 +125,11 @@ class AvailabilityDomain final {
                : std::nullopt;
   }
 
+  CustomAvailabilityDomain *getCustomDomain() const {
+    assert(isCustom());
+    return storage.get<CustomAvailabilityDomain *>();
+  }
+
 public:
   AvailabilityDomain() {}
 
@@ -141,6 +153,10 @@ class AvailabilityDomain final {
     return AvailabilityDomain(Kind::Embedded);
   }
 
+  static AvailabilityDomain forCustom(CustomAvailabilityDomain *domain) {
+    return AvailabilityDomain(domain);
+  }
+
   /// Returns the built-in availability domain identified by the given string.
   static std::optional<AvailabilityDomain>
   builtinDomainForString(StringRef string, const DeclContext *declContext);
@@ -149,7 +165,7 @@ class AvailabilityDomain final {
     if (auto inlineDomain = getInlineDomain())
       return inlineDomain->getKind();
 
-    llvm_unreachable("unimplemented");
+    return Kind::Custom;
   }
 
   bool isUniversal() const { return getKind() == Kind::Universal; }
@@ -164,6 +180,8 @@ class AvailabilityDomain final {
 
   bool isEmbedded() const { return getKind() == Kind::Embedded; }
 
+  bool isCustom() const { return getKind() == Kind::Custom; }
+
   /// Returns the platform kind for this domain if applicable.
   PlatformKind getPlatformKind() const {
     if (auto inlineDomain = getInlineDomain())
@@ -183,6 +201,9 @@ class AvailabilityDomain final {
   /// Returns the string to use when printing an `@available` attribute.
   llvm::StringRef getNameForAttributePrinting() const;
 
+  /// Returns the module that the domain belongs to, if it is a custom domain.
+  ModuleDecl *getModule() const;
+
   /// Returns true if availability in `other` is a subset of availability in
   /// this domain. The set of all availability domains form a lattice where the
   /// universal domain (`*`) is the bottom element.
@@ -201,6 +222,34 @@ class AvailabilityDomain final {
   }
 };
 
+/// Represents an availability domain that has been defined in a module.
+class CustomAvailabilityDomain : public ASTAllocated<CustomAvailabilityDomain> {
+public:
+  enum class Kind {
+    /// A domain that is known to be enabled at compile time.
+    Enabled,
+    /// A domain that is known to be disabled at compile time.
+    Disabled,
+    /// A domain with an enablement state that must be queried at runtime.
+    Dynamic,
+  };
+
+private:
+  Identifier name;
+  Kind kind;
+  ModuleDecl *mod;
+
+  CustomAvailabilityDomain(Identifier name, ModuleDecl *mod, Kind kind);
+
+public:
+  static CustomAvailabilityDomain *create(const ASTContext &ctx, StringRef name,
+                                          ModuleDecl *mod, Kind kind);
+
+  Identifier getName() const { return name; }
+  Kind getKind() const { return kind; }
+  ModuleDecl *getModule() const { return mod; }
+};
+
 } // end namespace swift
 
 namespace llvm {

lib/AST/Availability.cpp

@@ -221,6 +221,15 @@ void AvailabilityInference::applyInferredAvailableAttrs(
         return false;
       case AvailabilityDomain::Kind::Platform:
         return domain.getPlatformKind() < other.domain.getPlatformKind();
+      case AvailabilityDomain::Kind::Custom: {
+        auto mod = domain.getModule();
+        auto otherMod = other.domain.getModule();
+        if (mod != otherMod)
+          return mod->getName() < otherMod->getName();
+
+        return domain.getNameForAttributePrinting() <
+               other.domain.getNameForAttributePrinting();
+      }
       }
     }
   };

lib/AST/AvailabilityDomain.cpp

@@ -13,6 +13,7 @@
 #include "swift/AST/AvailabilityDomain.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/Decl.h"
+#include "swift/Basic/Assertions.h"
 #include "llvm/ADT/StringSwitch.h"
 
 using namespace swift;
@@ -48,6 +49,10 @@ bool AvailabilityDomain::isActive(const ASTContext &ctx) const {
     return true;
   case Kind::Platform:
     return isPlatformActive(getPlatformKind(), ctx.LangOpts);
+  case Kind::Custom:
+    // For now, custom domains are always active but it's conceivable that in
+    // the future someone might want to define a domain but leave it inactive.
+    return true;
   }
 }
 
@@ -63,6 +68,8 @@ llvm::StringRef AvailabilityDomain::getNameForDiagnostics() const {
     return "Embedded Swift";
   case Kind::Platform:
     return swift::prettyPlatformString(getPlatformKind());
+  case Kind::Custom:
+    return getCustomDomain()->getName().str();
   }
 }
 
@@ -78,9 +85,18 @@ llvm::StringRef AvailabilityDomain::getNameForAttributePrinting() const {
     return "Embedded";
   case Kind::Platform:
     return swift::platformString(getPlatformKind());
+  case Kind::Custom:
+    return getCustomDomain()->getName().str();
   }
 }
 
+ModuleDecl *AvailabilityDomain::getModule() const {
+  if (auto customDomain = getCustomDomain())
+    return customDomain->getModule();
+
+  return nullptr;
+}
+
 bool AvailabilityDomain::contains(const AvailabilityDomain &other) const {
   // FIXME: [availability] This currently implements something closer to a
   // total ordering instead of the more flexible partial ordering that it
@@ -100,5 +116,20 @@ bool AvailabilityDomain::contains(const AvailabilityDomain &other) const {
       return true;
     return inheritsAvailabilityFromPlatform(other.getPlatformKind(),
                                             getPlatformKind());
+  case Kind::Custom:
+    return getCustomDomain() == other.getCustomDomain();
   }
 }
+
+CustomAvailabilityDomain::CustomAvailabilityDomain(Identifier name,
+                                                   ModuleDecl *mod, Kind kind)
+    : name(name), kind(kind), mod(mod) {
+  ASSERT(!name.empty());
+  ASSERT(mod);
+}
+
+CustomAvailabilityDomain *
+CustomAvailabilityDomain::create(const ASTContext &ctx, StringRef name,
+                                 ModuleDecl *mod, Kind kind) {
+  return new (ctx) CustomAvailabilityDomain(ctx.getIdentifier(name), mod, kind);
+}

lib/Sema/TypeCheckAvailability.cpp

@@ -3006,6 +3006,7 @@ bool shouldHideDomainNameForConstraintDiagnostic(
   switch (constraint.getDomain().getKind()) {
   case AvailabilityDomain::Kind::Universal:
   case AvailabilityDomain::Kind::Embedded:
+  case AvailabilityDomain::Kind::Custom:
     return true;
   case AvailabilityDomain::Kind::Platform:
     return false;