Serialization: Encode custom availability domains.

When serializing `@available` attributes, if the attribute applies to a custom
domain include enough information to deserialize the reference to that domain.

Resolves rdar://138441265.

Author: tshortli

include/swift/AST/AvailabilityDomain.h

@@ -32,6 +32,7 @@
 namespace swift {
 class ASTContext;
 class CustomAvailabilityDomain;
+class Decl;
 class DeclContext;
 class ModuleDecl;
 
@@ -148,6 +149,11 @@ class AvailabilityDomain final {
     return AvailabilityDomain(Kind::Embedded);
   }
 
+  /// If `decl` represents an availability domain, returns the corresponding
+  /// `AvailabilityDomain` value. Otherwise, returns `std::nullopt`.
+  static std::optional<AvailabilityDomain> forCustom(Decl *decl,
+                                                     const ASTContext &ctx);
+
   static AvailabilityDomain forCustom(const CustomAvailabilityDomain *domain) {
     return AvailabilityDomain(domain);
   }
@@ -236,6 +242,10 @@ class AvailabilityDomain final {
   /// Returns the string to use when printing an `@available` attribute.
   llvm::StringRef getNameForAttributePrinting() const;
 
+  /// Returns the decl that represents the domain, or `nullptr` if the domain
+  /// does not have a decl.
+  Decl *getDecl() const;
+
   /// Returns the module that the domain belongs to, if it is a custom domain.
   ModuleDecl *getModule() const;
 
@@ -307,24 +317,26 @@ class CustomAvailabilityDomain : public llvm::FoldingSetNode {
   Identifier name;
   Kind kind;
   ModuleDecl *mod;
+  Decl *decl;
 
-  CustomAvailabilityDomain(Identifier name, ModuleDecl *mod, Kind kind);
+  CustomAvailabilityDomain(Identifier name, Kind kind, ModuleDecl *mod,
+                           Decl *decl);
 
 public:
-  static const CustomAvailabilityDomain *get(StringRef name, ModuleDecl *mod,
-                                             Kind kind, const ASTContext &ctx);
+  static const CustomAvailabilityDomain *get(StringRef name, Kind kind,
+                                             ModuleDecl *mod, Decl *decl,
+                                             const ASTContext &ctx);
 
   Identifier getName() const { return name; }
   Kind getKind() const { return kind; }
   ModuleDecl *getModule() const { return mod; }
+  Decl *getDecl() const { return decl; }
 
   /// Uniquing for `ASTContext`.
   static void Profile(llvm::FoldingSetNodeID &ID, Identifier name,
-                      ModuleDecl *mod, Kind kind);
+                      ModuleDecl *mod);
 
-  void Profile(llvm::FoldingSetNodeID &ID) const {
-    Profile(ID, name, mod, kind);
-  }
+  void Profile(llvm::FoldingSetNodeID &ID) const { Profile(ID, name, mod); }
 };
 
 /// Represents either a resolved availability domain or an identifier written

lib/AST/ASTContext.cpp

@@ -5866,11 +5866,11 @@ const AvailabilityContext::Storage *AvailabilityContext::Storage::get(
 }
 
 const CustomAvailabilityDomain *
-CustomAvailabilityDomain::get(StringRef name, ModuleDecl *mod, Kind kind,
-                              const ASTContext &ctx) {
+CustomAvailabilityDomain::get(StringRef name, Kind kind, ModuleDecl *mod,
+                              Decl *decl, const ASTContext &ctx) {
   auto identifier = ctx.getIdentifier(name);
   llvm::FoldingSetNodeID id;
-  CustomAvailabilityDomain::Profile(id, identifier, mod, kind);
+  CustomAvailabilityDomain::Profile(id, identifier, mod);
 
   auto &foldingSet = ctx.getImpl().CustomAvailabilityDomains;
   void *insertPos;
@@ -5880,7 +5880,8 @@ CustomAvailabilityDomain::get(StringRef name, ModuleDecl *mod, Kind kind,
 
   void *mem = ctx.Allocate(sizeof(CustomAvailabilityDomain),
                            alignof(CustomAvailabilityDomain));
-  auto *newNode = ::new (mem) CustomAvailabilityDomain(identifier, mod, kind);
+  auto *newNode =
+      ::new (mem) CustomAvailabilityDomain(identifier, kind, mod, decl);
   foldingSet.InsertNode(newNode, insertPos);
 
   return newNode;

lib/AST/AvailabilityDomain.cpp

@@ -17,10 +17,61 @@
 #include "swift/AST/Module.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/Basic/Assertions.h"
+#include "clang/AST/ASTContext.h"
+#include "clang/AST/Decl.h"
 #include "llvm/ADT/StringSwitch.h"
 
 using namespace swift;
 
+CustomAvailabilityDomain::Kind
+getCustomDomainKind(clang::FeatureAvailKind featureAvailKind) {
+  switch (featureAvailKind) {
+  case clang::FeatureAvailKind::None:
+    llvm_unreachable("unexpected kind");
+  case clang::FeatureAvailKind::Available:
+    return CustomAvailabilityDomain::Kind::Enabled;
+  case clang::FeatureAvailKind::Unavailable:
+    return CustomAvailabilityDomain::Kind::Disabled;
+  case clang::FeatureAvailKind::Dynamic:
+    return CustomAvailabilityDomain::Kind::Dynamic;
+  }
+}
+
+static const CustomAvailabilityDomain *
+customDomainForClangDecl(Decl *decl, const ASTContext &ctx) {
+  auto *clangDecl = decl->getClangDecl();
+  ASSERT(clangDecl);
+
+  auto featureInfo = clangDecl->getASTContext().getFeatureAvailInfo(
+      const_cast<clang::Decl *>(clangDecl));
+
+  // Ensure the decl actually represents an availability domain.
+  if (featureInfo.first.empty())
+    return nullptr;
+
+  if (featureInfo.second.Kind == clang::FeatureAvailKind::None)
+    return nullptr;
+
+  return CustomAvailabilityDomain::get(
+      featureInfo.first, getCustomDomainKind(featureInfo.second.Kind),
+      decl->getModuleContext(), decl, ctx);
+}
+
+std::optional<AvailabilityDomain>
+AvailabilityDomain::forCustom(Decl *decl, const ASTContext &ctx) {
+  if (!decl)
+    return std::nullopt;
+
+  if (decl->hasClangNode()) {
+    if (auto *customDomain = customDomainForClangDecl(decl, ctx))
+      return AvailabilityDomain::forCustom(customDomain);
+  } else {
+    // FIXME: [availability] Handle Swift availability domains decls.
+  }
+
+  return std::nullopt;
+}
+
 std::optional<AvailabilityDomain>
 AvailabilityDomain::builtinDomainForString(StringRef string,
                                            const DeclContext *declContext) {
@@ -166,6 +217,13 @@ llvm::StringRef AvailabilityDomain::getNameForAttributePrinting() const {
   }
 }
 
+Decl *AvailabilityDomain::getDecl() const {
+  if (auto *customDomain = getCustomDomain())
+    return customDomain->getDecl();
+
+  return nullptr;
+}
+
 ModuleDecl *AvailabilityDomain::getModule() const {
   if (auto customDomain = getCustomDomain())
     return customDomain->getModule();
@@ -268,19 +326,17 @@ bool StableAvailabilityDomainComparator::operator()(
   }
 }
 
-CustomAvailabilityDomain::CustomAvailabilityDomain(Identifier name,
-                                                   ModuleDecl *mod, Kind kind)
-    : name(name), kind(kind), mod(mod) {
+CustomAvailabilityDomain::CustomAvailabilityDomain(Identifier name, Kind kind,
+                                                   ModuleDecl *mod, Decl *decl)
+    : name(name), kind(kind), mod(mod), decl(decl) {
   ASSERT(!name.empty());
   ASSERT(mod);
 }
 
 void CustomAvailabilityDomain::Profile(llvm::FoldingSetNodeID &ID,
-                                       Identifier name, ModuleDecl *mod,
-                                       Kind kind) {
+                                       Identifier name, ModuleDecl *mod) {
   ID.AddPointer(name.getAsOpaquePointer());
   ID.AddPointer(mod);
-  ID.AddInteger(static_cast<unsigned>(kind));
 }
 
 std::optional<AvailabilityDomain>

lib/ClangImporter/ClangImporter.cpp

@@ -4115,49 +4115,28 @@ void ClangModuleUnit::lookupObjCMethods(
 void ClangModuleUnit::lookupAvailabilityDomains(
     Identifier identifier, SmallVectorImpl<AvailabilityDomain> &results) const {
   auto domainName = identifier.str();
+  auto &ctx = getASTContext();
   auto &clangASTContext = getClangASTContext();
 
-  // First, check the Clang AST context for a decl matching this availability
-  // domain name. If the domain was defined in the bridging header and that
-  // bridging header wasn't precompiled, this is where it will be found.
-  auto *domainDecl = clangASTContext.getFeatureAvailDecl(domainName);
-  if (!domainDecl) {
-    // Next, try external sources (.pcm/.pch).
-    auto externalSource = clangASTContext.getExternalSource();
-    if (externalSource)
-      domainDecl = dyn_cast_or_null<clang::VarDecl>(
-          externalSource->getAvailabilityDomainDecl(domainName));
-  }
+  auto domainInfo = clangASTContext.getFeatureAvailInfo(domainName);
+  if (domainInfo.Kind == clang::FeatureAvailKind::None)
+    return;
 
-  if (!domainDecl)
+  auto *varDecl = dyn_cast_or_null<clang::VarDecl>(domainInfo.Decl);
+  if (!varDecl)
     return;
 
   // The decl that was found may belong to a different Clang module.
-  if (domainDecl->getOwningModule() != getClangModule())
+  if (varDecl->getOwningModule() != getClangModule())
     return;
 
-  // Construct a CustomAvailabilityDomain for the decl that was found.
-  auto featureInfo = clangASTContext.getFeatureAvailInfo(domainDecl);
-  if (featureInfo.first.empty())
+  auto *imported = ctx.getClangModuleLoader()->importDeclDirectly(varDecl);
+  if (!imported)
     return;
 
-  auto getDomainKind = [](clang::FeatureAvailKind featureAvailKind) {
-    switch (featureAvailKind) {
-    case clang::FeatureAvailKind::None:
-      llvm_unreachable("unexpected kind");
-    case clang::FeatureAvailKind::Available:
-      return CustomAvailabilityDomain::Kind::Enabled;
-    case clang::FeatureAvailKind::Unavailable:
-      return CustomAvailabilityDomain::Kind::Disabled;
-    case clang::FeatureAvailKind::Dynamic:
-      return CustomAvailabilityDomain::Kind::Dynamic;
-    }
-  };
-
-  auto domain = AvailabilityDomain::forCustom(CustomAvailabilityDomain::get(
-      featureInfo.first, getParentModule(),
-      getDomainKind(featureInfo.second.Kind), getASTContext()));
-  results.push_back(domain);
+  auto customDomain = AvailabilityDomain::forCustom(imported, ctx);
+  ASSERT(customDomain);
+  results.push_back(*customDomain);
 }
 
 void ClangModuleUnit::collectLinkLibraries(

lib/Frontend/Frontend.cpp

@@ -1409,7 +1409,8 @@ static void configureAvailabilityDomains(const ASTContext &ctx,
   llvm::SmallDenseMap<Identifier, const CustomAvailabilityDomain *> domainMap;
   auto createAndInsertDomain = [&](const std::string &name,
                                    CustomAvailabilityDomain::Kind kind) {
-    auto *domain = CustomAvailabilityDomain::get(name, mainModule, kind, ctx);
+    auto *domain =
+        CustomAvailabilityDomain::get(name, kind, mainModule, nullptr, ctx);
     bool inserted = domainMap.insert({domain->getName(), domain}).second;
     ASSERT(inserted); // Domains must be unique.
   };

lib/Serialization/Deserialization.cpp

@@ -181,6 +181,8 @@ const char InvalidEnumValueError::ID = '\0';
 void InvalidEnumValueError::anchor() {}
 const char ConformanceXRefError::ID = '\0';
 void ConformanceXRefError::anchor() {}
+const char InavalidAvailabilityDomainError::ID = '\0';
+void InavalidAvailabilityDomainError::anchor() {}
 
 /// Skips a single record in the bitstream.
 ///
@@ -5720,35 +5722,79 @@ ModuleFile::getDeclChecked(
   return declOrOffset;
 }
 
+static std::optional<AvailabilityDomainKind>
+decodeDomainKind(uint8_t kind) {
+  switch (kind) {
+    case static_cast<uint8_t>(AvailabilityDomainKind::Universal):
+      return AvailabilityDomainKind::Universal;
+    case static_cast<uint8_t>(AvailabilityDomainKind::SwiftLanguage):
+      return AvailabilityDomainKind::SwiftLanguage;
+    case static_cast<uint8_t>(AvailabilityDomainKind::PackageDescription):
+      return AvailabilityDomainKind::PackageDescription;
+    case static_cast<uint8_t>(AvailabilityDomainKind::Embedded):
+      return AvailabilityDomainKind::Embedded;
+    case static_cast<uint8_t>(AvailabilityDomainKind::Platform):
+      return AvailabilityDomainKind::Platform;
+    case static_cast<uint8_t>(AvailabilityDomainKind::Custom):
+      return AvailabilityDomainKind::Custom;
+    default:
+      return std::nullopt;
+  }
+}
+
+static std::optional<AvailabilityDomain>
+decodeAvailabilityDomain(AvailabilityDomainKind domainKind,
+                         PlatformKind platformKind, Decl *decl,
+                         const ASTContext &ctx) {
+  switch (domainKind) {
+  case AvailabilityDomainKind::Universal:
+    return AvailabilityDomain::forUniversal();
+  case AvailabilityDomainKind::SwiftLanguage:
+    return AvailabilityDomain::forSwiftLanguage();
+  case AvailabilityDomainKind::PackageDescription:
+    return AvailabilityDomain::forPackageDescription();
+  case AvailabilityDomainKind::Embedded:
+    return AvailabilityDomain::forEmbedded();
+  case AvailabilityDomainKind::Platform:
+    return AvailabilityDomain::forPlatform(platformKind);
+  case AvailabilityDomainKind::Custom:
+    return AvailabilityDomain::forCustom(decl, ctx);
+  }
+}
+
 Expected<AvailableAttr *>
 DeclDeserializer::readAvailable_DECL_ATTR(SmallVectorImpl<uint64_t> &scratch,
                                           StringRef blobData) {
   bool isImplicit;
   bool isUnavailable;
   bool isDeprecated;
   bool isNoAsync;
-  bool isPackageDescriptionVersionSpecific;
   bool isSPI;
-  bool isForEmbedded;
+  uint8_t rawDomainKind;
+  unsigned rawPlatform;
+  DeclID domainDeclID;
   DEF_VER_TUPLE_PIECES(Introduced);
   DEF_VER_TUPLE_PIECES(Deprecated);
   DEF_VER_TUPLE_PIECES(Obsoleted);
-  unsigned rawPlatform, messageSize, renameSize;
+  unsigned messageSize, renameSize;
 
   // Decode the record, pulling the version tuple information.
   serialization::decls_block::AvailableDeclAttrLayout::readRecord(
-      scratch, isImplicit, isUnavailable, isDeprecated, isNoAsync,
-      isPackageDescriptionVersionSpecific, isSPI, isForEmbedded,
+      scratch, isImplicit, isUnavailable, isDeprecated, isNoAsync, isSPI,
+      rawDomainKind, rawPlatform, domainDeclID,
       LIST_VER_TUPLE_PIECES(Introduced), LIST_VER_TUPLE_PIECES(Deprecated),
-      LIST_VER_TUPLE_PIECES(Obsoleted), rawPlatform, messageSize,
-      renameSize);
+      LIST_VER_TUPLE_PIECES(Obsoleted), messageSize, renameSize);
+
+  auto maybeDomainKind = decodeDomainKind(rawDomainKind);
+  if (!maybeDomainKind)
+    return llvm::make_error<InvalidEnumValueError>(rawDomainKind, "AvailabilityDomainKind");
 
   auto maybePlatform = platformFromUnsigned(rawPlatform);
   if (!maybePlatform.has_value())
     return llvm::make_error<InvalidEnumValueError>(rawPlatform, "PlatformKind");
 
-  PlatformKind platform = maybePlatform.value();
-
+  AvailabilityDomainKind domainKind = *maybeDomainKind;
+  PlatformKind platform = *maybePlatform;
   StringRef message = blobData.substr(0, messageSize);
   blobData = blobData.substr(messageSize);
   StringRef rename = blobData.substr(0, renameSize);
@@ -5767,23 +5813,19 @@ DeclDeserializer::readAvailable_DECL_ATTR(SmallVectorImpl<uint64_t> &scratch,
   else
     kind = AvailableAttr::Kind::Default;
 
-  AvailabilityDomain domain;
-  if (platform != PlatformKind::none) {
-    domain = AvailabilityDomain::forPlatform(platform);
-  } else if (!Introduced.empty() || !Deprecated.empty() || !Obsoleted.empty()) {
-    domain = isPackageDescriptionVersionSpecific
-                 ? AvailabilityDomain::forPackageDescription()
-                 : AvailabilityDomain::forSwiftLanguage();
-  } else if (isForEmbedded) {
-    domain = AvailabilityDomain::forEmbedded();
-  } else {
-    domain = AvailabilityDomain::forUniversal();
+  Decl *domainDecl = nullptr;
+  if (domainDeclID) {
+    SET_OR_RETURN_ERROR(domainDecl, MF.getDeclChecked(domainDeclID));
   }
 
+  auto domain = decodeAvailabilityDomain(domainKind, platform, domainDecl, ctx);
+  if (!domain)
+    return llvm::make_error<InavalidAvailabilityDomainError>();
+
   auto attr = new (ctx)
-      AvailableAttr(SourceLoc(), SourceRange(), domain, SourceLoc(), kind, message, rename,
-                    Introduced, SourceRange(), Deprecated, SourceRange(),
-                    Obsoleted, SourceRange(), isImplicit, isSPI);
+      AvailableAttr(SourceLoc(), SourceRange(), *domain, SourceLoc(), kind,
+                    message, rename, Introduced, SourceRange(), Deprecated,
+                    SourceRange(), Obsoleted, SourceRange(), isImplicit, isSPI);
   return attr;
 }