ClangImporter: Look up availability domains defined in bridging headers.

This is very brittle in this first iteration. For now we require the
declaration representing the availability domain be deserialized before it can
be looked up by name since Clang does not have a lookup table for availabilty
domains in its module representation. As a result, it only works for bridging
headers that are not precompiled.

Part of rdar://138441266.

Author: tshortli

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

@@ -350,6 +350,11 @@ void SwiftLookupTable::addCategory(clang::ObjCCategoryDecl *category) {
   Categories.push_back(category);
 }
 
+void SwiftLookupTable::addAvailabilityDomainDecl(StringRef name,
+                                                 clang::VarDecl *decl) {
+  AvailabilityDomains.insert({name, StoredSingleEntry(decl)});
+}
+
 bool SwiftLookupTable::resolveUnresolvedEntries(
     SmallVectorImpl<SingleEntry> &unresolved) {
   // Common case: nothing left to resolve.
@@ -712,6 +717,17 @@ SwiftLookupTable::lookupGlobalsAsMembers(SerializedSwiftName baseName,
   return lookupGlobalsAsMembersImpl(baseName, *storedContext);
 }
 
+clang::VarDecl *SwiftLookupTable::lookupAvailabilityDomainDecl(StringRef name) {
+  // FIXME: [availability] Remove this once Clang has a lookup table.
+  auto result = AvailabilityDomains.find(name);
+  if (result == AvailabilityDomains.end())
+    return nullptr;
+
+  auto &entry = result->second;
+  DEBUG_ASSERT(entry.isASTNodeEntry());
+  return static_cast<clang::VarDecl *>(entry.getASTNode());
+}
+
 SmallVector<SwiftLookupTable::SingleEntry, 4>
 SwiftLookupTable::allGlobalsAsMembersInContext(EffectiveClangContext context) {
   if (!context) return { };
@@ -1868,9 +1884,9 @@ SwiftNameLookupExtension::hashExtension(ExtensionHashBuilder &HBuilder) const {
 void importer::addEntryToLookupTable(SwiftLookupTable &table,
                                      clang::NamedDecl *named,
                                      NameImporter &nameImporter) {
+  auto &clangContext = nameImporter.getClangContext();
   clang::PrettyStackTraceDecl trace(
-      named, named->getLocation(),
-      nameImporter.getClangContext().getSourceManager(),
+      named, named->getLocation(), clangContext.getSourceManager(),
       "while adding SwiftName lookup table entries for clang declaration");
 
   // Determine whether this declaration is suppressed in Swift.
@@ -2009,6 +2025,16 @@ void importer::addEntryToLookupTable(SwiftLookupTable &table,
         addEntryToLookupTable(table, usingShadowDecl, nameImporter);
     }
   }
+
+  // If this decl represents an availability domain, add it to the lookup table
+  // as one.
+  // FIXME: [availability] Remove this once Clang has a lookup table.
+  if (auto varDecl = dyn_cast<clang::VarDecl>(named)) {
+    auto mutableVar = const_cast<clang::VarDecl *>(varDecl);
+    auto domainInfo = clangContext.getFeatureAvailInfo(mutableVar);
+    if (!domainInfo.first.empty())
+      table.addAvailabilityDomainDecl(domainInfo.first, mutableVar);
+  }
 }
 
 /// Returns the nearest parent of \p module that is marked \c explicit in its

lib/ClangImporter/SwiftLookupTable.h

@@ -486,6 +486,10 @@ class SwiftLookupTable {
   /// The list of Objective-C categories and extensions.
   llvm::SmallVector<clang::ObjCCategoryDecl *, 4> Categories;
 
+  /// A mapping from availability domain name strings to their corresponding
+  /// declarations.
+  llvm::SmallDenseMap<StringRef, StoredSingleEntry> AvailabilityDomains;
+
   /// A mapping from stored contexts to the set of global declarations that
   /// are mapped to members within that context.
   ///
@@ -550,6 +554,9 @@ class SwiftLookupTable {
   /// Add an Objective-C category or extension to the table.
   void addCategory(clang::ObjCCategoryDecl *category);
 
+  /// Add an entry for a Clang availability domain.
+  void addAvailabilityDomainDecl(StringRef name, clang::VarDecl *decl);
+
   /// Resolve any unresolved entries.
   ///
   /// \param unresolved Will be populated with the list of entries
@@ -632,6 +639,10 @@ class SwiftLookupTable {
   lookupGlobalsAsMembers(SerializedSwiftName baseName,
                          EffectiveClangContext searchContext);
 
+  /// Retrieve the `VarDecl` that represents the availability domain with the
+  /// given name, or `nullptr` if there isn't one.
+  clang::VarDecl *lookupAvailabilityDomainDecl(StringRef name);
+
   SmallVector<SingleEntry, 4>
   allGlobalsAsMembersInContext(EffectiveClangContext context);