ClangImporter: Look up availability domains in Clang modules

include/swift/AST/DeclContext.h

@@ -20,6 +20,7 @@
 #define SWIFT_DECLCONTEXT_H
 
 #include "swift/AST/ASTAllocated.h"
+#include "swift/AST/AvailabilityDomain.h"
 #include "swift/AST/Identifier.h"
 #include "swift/AST/LookupKinds.h"
 #include "swift/AST/ResilienceExpansion.h"
@@ -667,6 +668,12 @@ class alignas(1 << DeclContextAlignInBits) DeclContext
          ObjCSelector selector,
          SmallVectorImpl<AbstractFunctionDecl *> &results) const;
 
+  /// Look up the custom availability domains matching the given identifier that
+  /// are visible from this context.
+  void
+  lookupAvailabilityDomains(Identifier identifier,
+                            SmallVectorImpl<AvailabilityDomain> &results) const;
+
   /// Looks up an infix operator with a given \p name.
   ///
   /// This returns a vector of results, as it's possible to find multiple infix

include/swift/AST/FileUnit.h

@@ -126,6 +126,14 @@ class FileUnit : public DeclContext, public ASTAllocated<FileUnit> {
                             const ModuleDecl *importedModule,
                             llvm::SmallSetVector<Identifier, 4> &spiGroups) const {};
 
+  /// Find all availability domains defined in this module with the given
+  /// identifier.
+  ///
+  /// This does a simple local lookup, not recursively looking through imports.
+  virtual void lookupAvailabilityDomains(
+      Identifier identifier,
+      SmallVectorImpl<AvailabilityDomain> &results) const {};
+
   virtual std::optional<Fingerprint>
   loadFingerprint(const IterableDeclContext *IDC) const {
     return std::nullopt;

include/swift/AST/Module.h

@@ -970,6 +970,12 @@ class ModuleDecl
                          const ModuleDecl *importedModule,
                          llvm::SmallSetVector<Identifier, 4> &spiGroups) const;
 
+  /// Finds the custom availability domain defined by this module with the
+  /// given identifier and if one exists adds it to results.
+  void
+  lookupAvailabilityDomains(Identifier identifier,
+                            SmallVectorImpl<AvailabilityDomain> &results) const;
+
   // Is \p attr accessible as an explicitly imported SPI from this module?
   bool isImportedAsSPI(const SpecializeAttr *attr,
                        const ValueDecl *targetDecl) const;
@@ -1234,11 +1240,6 @@ 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);
   }

include/swift/ClangImporter/ClangModule.h

@@ -90,6 +90,10 @@ class ClangModuleUnit final : public LoadedFile {
          ObjCSelector selector,
          SmallVectorImpl<AbstractFunctionDecl *> &results) const override;
 
+  void lookupAvailabilityDomains(
+      Identifier identifier,
+      SmallVectorImpl<AvailabilityDomain> &results) const override;
+
   virtual void getTopLevelDecls(SmallVectorImpl<Decl*> &results) const override;
 
   virtual void getDisplayDecls(SmallVectorImpl<Decl*> &results, bool recursive = false) const override;

lib/AST/AvailabilityDomain.cpp

@@ -288,21 +288,6 @@ void CustomAvailabilityDomain::Profile(llvm::FoldingSetNodeID &ID,
   ID.AddInteger(static_cast<unsigned>(kind));
 }
 
-static std::optional<AvailabilityDomain>
-getAvailabilityDomainForName(Identifier identifier,
-                             const DeclContext *declContext) {
-  if (auto builtinDomain = AvailabilityDomain::builtinDomainForString(
-          identifier.str(), declContext))
-    return builtinDomain;
-
-  auto &ctx = declContext->getASTContext();
-  if (auto customDomain =
-          ctx.MainModule->getAvailabilityDomainForIdentifier(identifier))
-    return customDomain;
-
-  return std::nullopt;
-}
-
 std::optional<AvailabilityDomain>
 AvailabilityDomainOrIdentifier::lookUpInDeclContext(
     SourceLoc loc, const DeclContext *declContext) const {
@@ -312,7 +297,13 @@ AvailabilityDomainOrIdentifier::lookUpInDeclContext(
   auto &diags = ctx.Diags;
   std::optional<AvailabilityDomain> domain;
   auto identifier = getAsIdentifier().value();
-  domain = getAvailabilityDomainForName(identifier, declContext);
+
+  llvm::SmallVector<AvailabilityDomain> results;
+  declContext->lookupAvailabilityDomains(identifier, results);
+  if (results.size() > 0) {
+    // FIXME: [availability] Diagnose ambiguity if necessary.
+    domain = results.front();
+  }
 
   if (!domain) {
     auto domainString = identifier.str();

lib/AST/Module.cpp

@@ -1090,6 +1090,18 @@ void ModuleDecl::lookupImportedSPIGroups(
   FORWARD(lookupImportedSPIGroups, (importedModule, spiGroups));
 }
 
+void ModuleDecl::lookupAvailabilityDomains(
+    Identifier identifier,
+    llvm::SmallVectorImpl<AvailabilityDomain> &results) const {
+  auto iter = AvailabilityDomains.find(identifier);
+  if (iter != AvailabilityDomains.end()) {
+    results.push_back(AvailabilityDomain::forCustom(iter->getSecond()));
+    return;
+  }
+
+  FORWARD(lookupAvailabilityDomains, (identifier, results));
+}
+
 void BuiltinUnit::lookupValue(DeclName name, NLKind lookupKind,
                               OptionSet<ModuleLookupFlags> Flags,
                               SmallVectorImpl<ValueDecl*> &result) const {
@@ -4172,15 +4184,6 @@ version::Version ModuleDecl::getLanguageVersionBuiltWith() const {
   return version::Version();
 }
 
-std::optional<AvailabilityDomain>
-ModuleDecl::getAvailabilityDomainForIdentifier(Identifier identifier) const {
-  auto iter = AvailabilityDomains.find(identifier);
-  if (iter == AvailabilityDomains.end())
-    return std::nullopt;
-
-  return AvailabilityDomain::forCustom(iter->getSecond());
-}
-
 //===----------------------------------------------------------------------===//
 //                            MARK: SwiftSettings
 //===----------------------------------------------------------------------===//

lib/AST/NameLookup.cpp

@@ -2893,6 +2893,21 @@ void DeclContext::lookupAllObjCMethods(
     results.end());
 }
 
+void DeclContext::lookupAvailabilityDomains(
+    Identifier identifier, SmallVectorImpl<AvailabilityDomain> &results) const {
+  if (auto builtinDomain =
+          AvailabilityDomain::builtinDomainForString(identifier.str(), this)) {
+    results.push_back(*builtinDomain);
+    return;
+  }
+
+  // FIXME: [availability] Find the file/module scope decl context and look up
+  // the domain in that context using a request to cache the results.
+  for (auto import : namelookup::getAllImports(this)) {
+    import.importedModule->lookupAvailabilityDomains(identifier, results);
+  }
+}
+
 /// Given a set of type declarations, find all of the nominal type declarations
 /// that they reference, looking through typealiases as appropriate.
 static TinyPtrVector<NominalTypeDecl *>

lib/ClangImporter/ClangImporter.cpp

@@ -4103,6 +4103,39 @@ void ClangModuleUnit::lookupObjCMethods(
   }
 }
 
+void ClangModuleUnit::lookupAvailabilityDomains(
+    Identifier identifier, SmallVectorImpl<AvailabilityDomain> &results) const {
+  auto lookupTable = owner.findLookupTable(clangModule);
+  if (!lookupTable)
+    return;
+
+  auto varDecl = lookupTable->lookupAvailabilityDomainDecl(identifier.str());
+  if (!varDecl)
+    return;
+
+  auto featureInfo = getClangASTContext().getFeatureAvailInfo(varDecl);
+  if (featureInfo.first.empty())
+    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);
+}
+
 void ClangModuleUnit::collectLinkLibraries(
     ModuleDecl::LinkLibraryCallback callback) const {
   if (!clangModule)

lib/ClangImporter/ImportDecl.cpp

@@ -9186,6 +9186,34 @@ void ClangImporter::Implementation::importAttributes(
       MappedDecl->getAttrs().add(AvAttr);
     }
 
+    // __attribute__((availability(domain:)))
+    //
+    if (auto avail = dyn_cast<clang::DomainAvailabilityAttr>(*AI)) {
+      auto *declContext = MappedDecl->getInnermostDeclContext();
+
+      // FIXME: [availability] Don't look up the availability domain. Clang
+      // should be serializing the resolved VarDecl for the availability domain
+      // it found when type checking the attribute.
+      auto domainIdentifier = SwiftContext.getIdentifier(avail->getDomain());
+      llvm::SmallVector<AvailabilityDomain, 4> results;
+      declContext->lookupAvailabilityDomains(domainIdentifier, results);
+
+      if (results.size() > 0) {
+        // FIXME: [availability] Diagnose ambiguous availabilty domain name?
+        auto AttrKind = avail->getUnavailable()
+                            ? AvailableAttr::Kind::Unavailable
+                            : AvailableAttr::Kind::Default;
+
+        auto avAttr = new (C) AvailableAttr(
+            SourceLoc(), SourceRange(), results.front(), SourceLoc(), AttrKind,
+            /*Message=*/"", /*Rename=*/"", /*Introduced=*/{}, SourceRange(),
+            /*Deprecated=*/{}, SourceRange(), /*Obsoleted=*/{}, SourceRange(),
+            /*Implicit=*/false, /*IsSPI=*/false);
+
+        MappedDecl->getAttrs().add(avAttr);
+      }
+    }
+
     // __attribute__((swift_attr("attribute"))) are handled by
     // importSwiftAttrAttributes(). Other attributes are ignored.
   }