AST: Allow multiple unavailable domains in AvailabilityContext

include/swift/AST/AvailabilityContext.h

@@ -48,13 +48,6 @@ class AvailabilityContext {
     assert(storage);
   };
 
-  /// Retrieves an `AvailabilityContext` with the given platform availability
-  /// parameters.
-  static AvailabilityContext
-  get(const AvailabilityRange &platformAvailability,
-      std::optional<AvailabilityDomain> unavailableDomain, bool deprecated,
-      ASTContext &ctx);
-
 public:
   /// Retrieves an `AvailabilityContext` constrained by the given platform
   /// availability range.
@@ -119,6 +112,9 @@ class AvailabilityContext {
 
   void print(llvm::raw_ostream &os) const;
   SWIFT_DEBUG_DUMP;
+
+  /// Returns true if all internal invariants are satisfied.
+  bool verify(ASTContext &ctx) const;
 };
 
 } // end namespace swift

include/swift/AST/AvailabilityContextStorage.h

@@ -31,8 +31,9 @@ class AvailabilityContext::Info {
   /// The introduction version.
   AvailabilityRange Range;
 
-  /// The broadest unavailable domain.
-  std::optional<AvailabilityDomain> UnavailableDomain;
+  /// A sorted collection of disjoint domains that are known to be
+  /// unavailable in this context.
+  llvm::SmallVector<AvailabilityDomain, 1> UnavailableDomains;
 
   /// Whether or not the context is considered deprecated on the current
   /// platform.
@@ -49,20 +50,20 @@ class AvailabilityContext::Info {
   bool constrainWith(const DeclAvailabilityConstraints &constraints,
                      ASTContext &ctx);
 
-  bool constrainUnavailability(std::optional<AvailabilityDomain> domain);
+  bool constrainUnavailability(
+      const llvm::SmallVectorImpl<AvailabilityDomain> &domains);
+  bool constrainUnavailability(AvailabilityDomain domain) {
+    return constrainUnavailability(
+        llvm::SmallVector<AvailabilityDomain>{domain});
+  }
 
   /// Returns true if `other` is as available or is more available.
   bool isContainedIn(const Info &other) const;
 
-  void Profile(llvm::FoldingSetNodeID &ID) const {
-    Range.getRawVersionRange().Profile(ID);
-    if (UnavailableDomain) {
-      UnavailableDomain->Profile(ID);
-    } else {
-      ID.AddPointer(nullptr);
-    }
-    ID.AddBoolean(IsDeprecated);
-  }
+  void Profile(llvm::FoldingSetNodeID &ID) const;
+
+  /// Returns true if all internal invariants are satisfied.
+  bool verify(ASTContext &ctx) const;
 };
 
 /// As an implementation detail, the values that make up an `Availability`

include/swift/AST/AvailabilityDomain.h

@@ -265,6 +265,14 @@ inline void simple_display(llvm::raw_ostream &os,
   domain.print(os);
 }
 
+/// A comparator that implements a stable, total ordering on
+/// `AvailabilityDomain` that can be used for sorting in contexts where the
+/// result must be stable and deterministic across compilations.
+struct StableAvailabilityDomainComparator {
+  bool operator()(const AvailabilityDomain &lhs,
+                  const AvailabilityDomain &rhs) const;
+};
+
 /// Represents an availability domain that has been defined in a module.
 class CustomAvailabilityDomain : public ASTAllocated<CustomAvailabilityDomain> {
 public:

lib/AST/Availability.cpp

@@ -168,46 +168,11 @@ 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<OrderedAvailabilityDomain, InferredAvailability> Inferred;
+  // a per-domain basis.
+  std::map<AvailabilityDomain, InferredAvailability,
+           StableAvailabilityDomainComparator>
+      Inferred;
   for (const Decl *D : InferredFromDecls) {
     llvm::SmallVector<SemanticAvailableAttr, 8> MergedAttrs;
 
@@ -242,8 +207,7 @@ 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.domain, Pair.second, Context))
+    if (auto Attr = createAvailableAttr(Pair.first, Pair.second, Context))
       Attrs.add(Attr);
   }
 }

lib/AST/AvailabilityContext.cpp

@@ -38,27 +38,10 @@ static bool constrainRange(AvailabilityRange &existing,
   return true;
 }
 
-static bool constrainUnavailableDomain(
-    std::optional<AvailabilityDomain> &domain,
-    const std::optional<AvailabilityDomain> &otherDomain) {
-  // If the other domain is absent or is the same domain, it's a noop.
-  if (!otherDomain || domain == otherDomain)
-    return false;
-
-  // Check if the other domain is a superset and constrain to it if it is.
-  if (!domain || otherDomain->contains(*domain)) {
-    domain = otherDomain;
-    return true;
-  }
-
-  return false;
-}
-
 bool AvailabilityContext::Info::constrainWith(const Info &other) {
   bool isConstrained = false;
   isConstrained |= constrainRange(Range, other.Range);
-  if (other.UnavailableDomain)
-    isConstrained |= constrainUnavailability(other.UnavailableDomain);
+  isConstrained |= constrainUnavailability(other.UnavailableDomains);
   isConstrained |= CONSTRAIN_BOOL(IsDeprecated, other.IsDeprecated);
 
   return isConstrained;
@@ -89,24 +72,72 @@ bool AvailabilityContext::Info::constrainWith(
   return isConstrained;
 }
 
+/// Returns true if `domain` is not already contained in `unavailableDomains`.
+/// Also, removes domains from `unavailableDomains` that are contained in
+/// `domain`.
+static bool shouldConstrainUnavailableDomains(
+    AvailabilityDomain domain,
+    llvm::SmallVectorImpl<AvailabilityDomain> &unavailableDomains) {
+  bool didRemove = false;
+  for (auto iter = unavailableDomains.rbegin(), end = unavailableDomains.rend();
+       iter != end; ++iter) {
+    auto const &existingDomain = *iter;
+
+    // Check if the domain is already unavailable.
+    if (existingDomain.contains(domain)) {
+      ASSERT(!didRemove); // This would indicate that the context is malformed.
+      return false;
+    }
+
+    // Check if the existing domain would be absorbed by the new domain.
+    if (domain.contains(existingDomain)) {
+      unavailableDomains.erase((iter + 1).base());
+      didRemove = true;
+    }
+  }
+
+  return true;
+}
+
 bool AvailabilityContext::Info::constrainUnavailability(
-    std::optional<AvailabilityDomain> domain) {
-  return constrainUnavailableDomain(UnavailableDomain, domain);
+    const llvm::SmallVectorImpl<AvailabilityDomain> &domains) {
+  llvm::SmallVector<AvailabilityDomain, 2> domainsToAdd;
+
+  for (auto domain : domains) {
+    if (shouldConstrainUnavailableDomains(domain, UnavailableDomains))
+      domainsToAdd.push_back(domain);
+  }
+
+  if (domainsToAdd.size() < 1)
+    return false;
+
+  // Add the candidate domain and then re-sort.
+  for (auto domain : domainsToAdd)
+    UnavailableDomains.push_back(domain);
+
+  llvm::sort(UnavailableDomains, StableAvailabilityDomainComparator());
+  return true;
 }
 
 bool AvailabilityContext::Info::isContainedIn(const Info &other) const {
   // The available versions range be the same or smaller.
   if (!Range.isContainedIn(other.Range))
     return false;
 
-  // The set of unavailable domains should be the same or larger.
-  if (auto otherUnavailableDomain = other.UnavailableDomain) {
-    if (!UnavailableDomain)
-      return false;
-
-    if (!UnavailableDomain->contains(otherUnavailableDomain.value()))
-      return false;
-  }
+  // Every unavailable domain in the other context should be contained in some
+  // unavailable domain in this context.
+  bool disjointUnavailability = llvm::any_of(
+      other.UnavailableDomains,
+      [&](const AvailabilityDomain &otherUnavailableDomain) {
+        return llvm::none_of(
+            UnavailableDomains,
+            [&otherUnavailableDomain](const AvailabilityDomain &domain) {
+              return domain.contains(otherUnavailableDomain);
+            });
+      });
+
+  if (disjointUnavailability)
+    return false;
 
   // The set of deprecated domains should be the same or larger.
   if (!IsDeprecated && other.IsDeprecated)
@@ -115,10 +146,29 @@ bool AvailabilityContext::Info::isContainedIn(const Info &other) const {
   return true;
 }
 
+void AvailabilityContext::Info::Profile(llvm::FoldingSetNodeID &ID) const {
+  Range.getRawVersionRange().Profile(ID);
+  ID.AddInteger(UnavailableDomains.size());
+  for (auto domain : UnavailableDomains) {
+    domain.Profile(ID);
+  }
+  ID.AddBoolean(IsDeprecated);
+}
+
+bool AvailabilityContext::Info::verify(ASTContext &ctx) const {
+  // Unavailable domains must be sorted to ensure folding set node lookups yield
+  // consistent results.
+  if (!llvm::is_sorted(UnavailableDomains,
+                       StableAvailabilityDomainComparator()))
+    return false;
+
+  return true;
+}
+
 AvailabilityContext
 AvailabilityContext::forPlatformRange(const AvailabilityRange &range,
                                       ASTContext &ctx) {
-  Info info{range, /*UnavailableDomain*/ std::nullopt,
+  Info info{range, /*UnavailableDomains*/ {},
             /*IsDeprecated*/ false};
   return AvailabilityContext(Storage::get(info, ctx));
 }
@@ -133,27 +183,20 @@ AvailabilityContext AvailabilityContext::forDeploymentTarget(ASTContext &ctx) {
       AvailabilityRange::forDeploymentTarget(ctx), ctx);
 }
 
-AvailabilityContext
-AvailabilityContext::get(const AvailabilityRange &platformAvailability,
-                         std::optional<AvailabilityDomain> unavailableDomain,
-                         bool deprecated, ASTContext &ctx) {
-  Info info{platformAvailability, unavailableDomain, deprecated};
-  return AvailabilityContext(Storage::get(info, ctx));
-}
-
 AvailabilityRange AvailabilityContext::getPlatformRange() const {
   return storage->info.Range;
 }
 
 bool AvailabilityContext::isUnavailable() const {
-  return storage->info.UnavailableDomain.has_value();
+  return storage->info.UnavailableDomains.size() > 0;
 }
 
 bool AvailabilityContext::containsUnavailableDomain(
     AvailabilityDomain domain) const {
-  if (auto unavailableDomain = storage->info.UnavailableDomain)
-    return unavailableDomain->contains(domain);
-
+  for (auto unavailableDomain : storage->info.UnavailableDomains) {
+    if (unavailableDomain.contains(domain))
+      return true;
+  }
   return false;
 }
 
@@ -217,10 +260,7 @@ void AvailabilityContext::constrainWithDeclAndPlatformRange(
 }
 
 bool AvailabilityContext::isContainedIn(const AvailabilityContext other) const {
-  if (!storage->info.isContainedIn(other.storage->info))
-    return false;
-
-  return true;
+  return storage->info.isContainedIn(other.storage->info);
 }
 
 static std::string
@@ -236,11 +276,19 @@ stringForAvailability(const AvailabilityRange &availability) {
 void AvailabilityContext::print(llvm::raw_ostream &os) const {
   os << "version=" << stringForAvailability(getPlatformRange());
 
-  if (auto unavailableDomain = storage->info.UnavailableDomain)
-    os << " unavailable=" << unavailableDomain->getNameForAttributePrinting();
+  if (storage->info.UnavailableDomains.size() > 0) {
+    os << " unavailable=";
+    llvm::interleave(
+        storage->info.UnavailableDomains, os,
+        [&](const AvailabilityDomain &domain) { domain.print(os); }, ",");
+  }
 
   if (isDeprecated())
     os << " deprecated";
 }
 
 void AvailabilityContext::dump() const { print(llvm::errs()); }
+
+bool AvailabilityContext::verify(ASTContext &ctx) const {
+  return storage->info.verify(ctx);
+}