AST: Introduce StableAvailabilityDomainComparator.

Author: tshortli

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/AvailabilityDomain.cpp

@@ -14,6 +14,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/DiagnosticsSema.h"
+#include "swift/AST/Module.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/Basic/Assertions.h"
 #include "llvm/ADT/StringSwitch.h"
@@ -203,6 +204,33 @@ AvailabilityDomain AvailabilityDomain::copy(ASTContext &ctx) const {
   }
 }
 
+bool StableAvailabilityDomainComparator::operator()(
+    const AvailabilityDomain &lhs, const AvailabilityDomain &rhs) const {
+  auto lhsKind = lhs.getKind();
+  auto rhsKind = rhs.getKind();
+  if (lhsKind != rhsKind)
+    return lhsKind < rhsKind;
+
+  switch (lhsKind) {
+  case AvailabilityDomain::Kind::Universal:
+  case AvailabilityDomain::Kind::SwiftLanguage:
+  case AvailabilityDomain::Kind::PackageDescription:
+  case AvailabilityDomain::Kind::Embedded:
+    return false;
+  case AvailabilityDomain::Kind::Platform:
+    return lhs.getPlatformKind() < rhs.getPlatformKind();
+  case AvailabilityDomain::Kind::Custom: {
+    auto lhsMod = lhs.getModule();
+    auto rhsMod = rhs.getModule();
+    if (lhsMod != rhsMod)
+      return lhsMod->getName() < rhsMod->getName();
+
+    return lhs.getNameForAttributePrinting() <
+           rhs.getNameForAttributePrinting();
+  }
+  }
+}
+
 CustomAvailabilityDomain::CustomAvailabilityDomain(Identifier name,
                                                    ModuleDecl *mod, Kind kind)
     : name(name), kind(kind), mod(mod) {