AST/Sema: Model custom availability domains

include/swift/AST/Attr.h

@@ -3319,7 +3319,7 @@ class SemanticAvailableAttr final {
     return attr->getRawIntroduced();
   }
 
-  /// The source range of the `introduced:` component.
+  /// The source range of the `introduced:` version component.
   SourceRange getIntroducedSourceRange() const { return attr->IntroducedRange; }
 
   /// Returns the effective range in which the declaration with this attribute
@@ -3331,11 +3331,17 @@ class SemanticAvailableAttr final {
     return attr->getRawDeprecated();
   }
 
+  /// The source range of the `deprecated:` version component.
+  SourceRange getDeprecatedSourceRange() const { return attr->DeprecatedRange; }
+
   /// The version tuple written in source for the `obsoleted:` component.
   std::optional<llvm::VersionTuple> getObsoleted() const {
     return attr->getRawObsoleted();
   }
 
+  /// The source range of the `obsoleted:` version component.
+  SourceRange getObsoletedSourceRange() const { return attr->ObsoletedRange; }
+
   /// Returns the `message:` field of the attribute, or an empty string.
   StringRef getMessage() const { return attr->Message; }
 

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())
@@ -172,6 +190,10 @@ class AvailabilityDomain final {
     return PlatformKind::none;
   }
 
+  /// Returns true if availability for this domain can be specified in terms of
+  /// version ranges.
+  bool isVersioned() const;
+
   /// Returns true if this domain is considered active in the current
   /// compilation context.
   bool isActive(const ASTContext &ctx) const;
@@ -183,6 +205,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.
@@ -196,28 +221,39 @@ class AvailabilityDomain final {
     return !(*this == other);
   }
 
-  /// A total, stable ordering on domains.
-  bool operator<(const AvailabilityDomain &other) const {
-    if (getKind() != other.getKind())
-      return getKind() < other.getKind();
-
-    switch (getKind()) {
-    case Kind::Universal:
-    case Kind::SwiftLanguage:
-    case Kind::PackageDescription:
-    case Kind::Embedded:
-      // These availability domains are singletons.
-      return false;
-    case Kind::Platform:
-      return getPlatformKind() < other.getPlatformKind();
-    }
-  }
-
   void Profile(llvm::FoldingSetNodeID &ID) const {
     ID.AddPointer(getOpaqueValue());
   }
 };
 
+/// 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 {
@@ -249,6 +285,10 @@ struct DenseMapInfo<AvailabilityDomain> {
   static inline AvailabilityDomain getTombstoneKey() {
     return DenseMapInfo<AvailabilityDomain::Storage>::getTombstoneKey();
   }
+  static inline unsigned getHashValue(AvailabilityDomain domain) {
+    return DenseMapInfo<AvailabilityDomain::Storage>::getHashValue(
+        domain.storage);
+  }
   static bool isEqual(const AvailabilityDomain LHS,
                       const AvailabilityDomain RHS) {
     return LHS == RHS;

include/swift/AST/DiagnosticsSema.def

@@ -6731,6 +6731,13 @@ WARNING(attr_availability_cannot_be_used_for_domain, none,
 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))
+WARNING(attr_availability_unexpected_version,none,
+       "unexpected version number in '%0' attribute for '%1'",
+       (const DeclAttribute, StringRef))
 
 ERROR(availability_decl_unavailable, none,
       "%0 is unavailable%select{ in %2|}1%select{|: %3}3",

include/swift/AST/Module.h

@@ -19,6 +19,7 @@
 
 #include "swift/AST/AccessNotes.h"
 #include "swift/AST/AttrKind.h"
+#include "swift/AST/AvailabilityDomain.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/DeclContext.h"
 #include "swift/AST/Identifier.h"
@@ -41,8 +42,8 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MD5.h"
 #include <optional>
-#include <unordered_map>
 #include <set>
+#include <unordered_map>
 
 namespace clang {
   class Module;
@@ -51,20 +52,16 @@ namespace clang {
 namespace swift {
   enum class ArtificialMainKind : uint8_t;
   class ASTContext;
-  class ASTScope;
   class ASTWalker;
-  class AvailabilityScope;
-  class BraceStmt;
+  class CustomAvailabilityDomain;
   class Decl;
   class DeclAttribute;
   class TypeDecl;
   enum class DeclKind : uint8_t;
-  class ExtensionDecl;
   class DebuggerClient;
   class DeclName;
   class FileUnit;
   class FuncDecl;
-  class InfixOperatorDecl;
   enum class LibraryLevel : uint8_t;
   class LinkLibrary;
   class ModuleLoader;
@@ -74,20 +71,12 @@ namespace swift {
   class PostfixOperatorDecl;
   class PrefixOperatorDecl;
   class ProtocolConformance;
-  class ProtocolDecl;
   struct PrintOptions;
   class SourceLookupCache;
-  class Token;
-  class TupleType;
   class Type;
   class ValueDecl;
-  class VarDecl;
   class VisibleDeclConsumer;
 
-namespace ast_scope {
-class ASTSourceFileScope;
-}
-
 /// Discriminator for file-units.
 enum class FileUnitKind {
   /// For a .swift source file.
@@ -361,6 +350,10 @@ class ModuleDecl
   /// Used by the debugger to bypass resilient access to fields.
   bool BypassResilience = false;
 
+  using AvailabilityDomainMap =
+      llvm::SmallDenseMap<Identifier, CustomAvailabilityDomain *>;
+  AvailabilityDomainMap AvailabilityDomains;
+
 public:
   using PopulateFilesFn = llvm::function_ref<void(
       ModuleDecl *, llvm::function_ref<void(FileUnit *)>)>;
@@ -1233,6 +1226,15 @@ class ModuleDecl
   /// An empty `Version` is returned if the information is not available.
   version::Version getLanguageVersionBuiltWith() const;
 
+  /// Returns the custom availability domain defined by this module with the
+  /// given identifier, if one exists.
+  std::optional<AvailabilityDomain>
+  getAvailabilityDomainForIdentifier(Identifier identifier) const;
+
+  void setAvailabilityDomains(const AvailabilityDomainMap &&map) {
+    AvailabilityDomains = std::move(map);
+  }
+
   static bool classof(const DeclContext *DC) {
     if (auto D = DC->getAsDecl())
       return classof(D);

include/swift/AST/SourceFile.h

@@ -25,9 +25,11 @@
 #include "llvm/ADT/STLExtras.h"
 
 namespace swift {
-
+class ASTScope;
+class AvailabilityScope;
 class PersistentParserState;
 struct SourceFileExtras;
+class Token;
 
 /// Kind of import affecting how a decl can be reexported.
 ///

lib/AST/Availability.cpp

@@ -201,9 +201,46 @@ void AvailabilityInference::applyInferredAvailableAttrs(
     Decl *ToDecl, ArrayRef<const Decl *> InferredFromDecls) {
   auto &Context = ToDecl->getASTContext();
 
+  /// A wrapper for AvailabilityDomain that implements a stable, total ordering for
+  /// domains. This is needed to ensure that the inferred attributes are added to
+  /// the declaration in a consistent order, preserving interface printing output
+  /// stability across compilations.
+  class OrderedAvailabilityDomain {
+  public:
+    AvailabilityDomain domain;
+
+    OrderedAvailabilityDomain(AvailabilityDomain domain) : domain(domain) {}
+
+    bool operator<(const OrderedAvailabilityDomain &other) const {
+      auto kind = domain.getKind();
+      auto otherKind = other.domain.getKind();
+      if (kind != otherKind)
+        return kind < otherKind;
+
+      switch (kind) {
+      case AvailabilityDomain::Kind::Universal:
+      case AvailabilityDomain::Kind::SwiftLanguage:
+      case AvailabilityDomain::Kind::PackageDescription:
+      case AvailabilityDomain::Kind::Embedded:
+        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();
+      }
+      }
+    }
+  };
+
   // Iterate over the declarations and infer required availability on
   // a per-platform basis.
-  std::map<AvailabilityDomain, InferredAvailability> Inferred;
+  std::map<OrderedAvailabilityDomain, InferredAvailability> Inferred;
   for (const Decl *D : InferredFromDecls) {
     llvm::SmallVector<SemanticAvailableAttr, 8> MergedAttrs;
 
@@ -238,7 +275,8 @@ void AvailabilityInference::applyInferredAvailableAttrs(
   // Create an availability attribute for each observed platform and add
   // to ToDecl.
   for (auto &Pair : Inferred) {
-    if (auto Attr = createAvailableAttr(Pair.first, Pair.second, Context))
+    if (auto Attr =
+            createAvailableAttr(Pair.first.domain, Pair.second, Context))
       Attrs.add(Attr);
   }
 }