AST: Optimize ModuleDecl::isImportedAsWeakLinked()

include/swift/AST/FileUnit.h

@@ -125,17 +125,6 @@ class FileUnit : public DeclContext, public ASTAllocated<FileUnit> {
                             const ModuleDecl *importedModule,
                             llvm::SmallSetVector<Identifier, 4> &spiGroups) const {};
 
-  /// Checks whether this file imports \c module as \c @_weakLinked.
-  virtual bool importsModuleAsWeakLinked(const ModuleDecl *module) const {
-    // For source files, this should be overridden to inspect the import
-    // declarations in the file. Other kinds of file units, like serialized
-    // modules, can just use this default implementation since the @_weakLinked
-    // attribute is not transitive. If module C is imported @_weakLinked by
-    // module B, that does not imply that module A imports module C @_weakLinked
-    // if it imports module B.
-    return false;
-  }
-
   virtual std::optional<Fingerprint>
   loadFingerprint(const IterableDeclContext *IDC) const {
     return std::nullopt;

include/swift/AST/ImportCache.h

@@ -127,13 +127,18 @@ class alignas(ImportedModule) ImportCache {
   llvm::DenseMap<std::tuple<const ModuleDecl *,
                             const DeclContext *>,
                  bool> SwiftOnlyCache;
+  llvm::DenseMap<const ModuleDecl *, ArrayRef<ModuleDecl *>> WeakCache;
 
   ImportPath::Access EmptyAccessPath;
 
   ArrayRef<ImportPath::Access> allocateArray(
       ASTContext &ctx,
       SmallVectorImpl<ImportPath::Access> &results);
 
+  ArrayRef<ModuleDecl *> allocateArray(
+      ASTContext &ctx,
+      llvm::SetVector<ModuleDecl *> &results);
+
   ImportSet &getImportSet(ASTContext &ctx,
                           ArrayRef<ImportedModule> topLevelImports);
 
@@ -167,6 +172,13 @@ class alignas(ImportedModule) ImportCache {
                                  const ModuleDecl *other,
                                  const DeclContext *dc);
 
+  /// Returns all weak-linked imported modules.
+  ArrayRef<ModuleDecl *>
+  getWeakImports(const ModuleDecl *mod);
+
+  bool isWeakImportedBy(const ModuleDecl *mod,
+                        const ModuleDecl *from);
+
   /// This is a hack to cope with main file parsing and REPL parsing, where
   /// we can add ImportDecls after import resolution.
   void clear() {

include/swift/AST/SourceFile.h

@@ -478,9 +478,6 @@ class SourceFile final : public FileUnit {
                 const ModuleDecl *importedModule,
                 llvm::SmallSetVector<Identifier, 4> &spiGroups) const override;
 
-  /// Is \p module imported as \c @_weakLinked by this file?
-  bool importsModuleAsWeakLinked(const ModuleDecl *module) const override;
-
   // Is \p targetDecl accessible as an explicitly imported SPI from this file?
   bool isImportedAsSPI(const ValueDecl *targetDecl) const;
 

lib/AST/ImportCache.cpp

@@ -21,6 +21,7 @@
 #include "swift/AST/FileUnit.h"
 #include "swift/AST/ImportCache.h"
 #include "swift/AST/Module.h"
+#include "swift/AST/SourceFile.h"
 #include "swift/Basic/Assertions.h"
 
 using namespace swift;
@@ -392,3 +393,56 @@ swift::namelookup::getAllImports(const DeclContext *dc) {
   return dc->getASTContext().getImportCache().getImportSet(dc)
     .getAllImports();
 }
+
+ArrayRef<ModuleDecl *> ImportCache::allocateArray(
+    ASTContext &ctx,
+    llvm::SetVector<ModuleDecl *> &results) {
+  if (results.empty())
+    return {};
+  else
+    return ctx.AllocateCopy(results.getArrayRef());
+}
+
+ArrayRef<ModuleDecl *>
+ImportCache::getWeakImports(const ModuleDecl *mod) {
+  auto found = WeakCache.find(mod);
+  if (found != WeakCache.end())
+    return found->second;
+
+  llvm::SetVector<ModuleDecl *> result;
+
+  for (auto file : mod->getFiles()) {
+    auto *sf = dyn_cast<SourceFile>(file);
+    // Other kinds of file units, like serialized modules, can just use this
+    // default implementation since the @_weakLinked attribute is not
+    // transitive. If module C is imported @_weakLinked by module B, that does
+    // not imply that module A imports module C @_weakLinked if it imports
+    // module B.
+    if (!sf)
+      continue;
+
+    for (auto &import : sf->getImports()) {
+      if (!import.options.contains(ImportFlags::WeakLinked))
+        continue;
+
+      ModuleDecl *importedModule = import.module.importedModule;
+      result.insert(importedModule);
+
+      auto reexportedModules = getImportSet(importedModule).getAllImports();
+      for (auto reexportedModule : reexportedModules) {
+        result.insert(reexportedModule.importedModule);
+      }
+    }
+  }
+
+  auto resultArray = allocateArray(mod->getASTContext(), result);
+  WeakCache[mod] = resultArray;
+  return resultArray;
+}
+
+bool ImportCache::isWeakImportedBy(const ModuleDecl *mod,
+                                   const ModuleDecl *from) {
+  auto weakImports = getWeakImports(from);
+  return std::find(weakImports.begin(), weakImports.end(), mod)
+      != weakImports.end();
+}

lib/AST/Module.cpp

@@ -3121,35 +3121,6 @@ bool SourceFile::isImportedAsSPI(const ValueDecl *targetDecl) const {
   return false;
 }
 
-bool SourceFile::importsModuleAsWeakLinked(const ModuleDecl *module) const {
-  for (auto &import : *Imports) {
-    if (!import.options.contains(ImportFlags::WeakLinked))
-      continue;
-
-    const ModuleDecl *importedModule = import.module.importedModule;
-    if (module == importedModule)
-      return true;
-
-    // Also check whether the target module is actually the underlyingClang
-    // module for this @_weakLinked import.
-    const ModuleDecl *clangModule =
-        importedModule->getUnderlyingModuleIfOverlay();
-    if (module == clangModule)
-      return true;
-
-    // Traverse the exported modules of this weakly-linked module to ensure
-    // that we weak-link declarations from its exported peers.
-    SmallVector<ImportedModule, 8> reexportedModules;
-    importedModule->getImportedModules(reexportedModules,
-                                       ModuleDecl::ImportFilterKind::Exported);
-    for (const ImportedModule &reexportedModule : reexportedModules) {
-      if (module == reexportedModule.importedModule)
-        return true;
-    }
-  }
-  return false;
-}
-
 bool ModuleDecl::isImportedAsSPI(const SpecializeAttr *attr,
                                  const ValueDecl *targetDecl) const {
   auto declSPIGroups = attr->getSPIGroups();
@@ -3181,11 +3152,7 @@ bool ModuleDecl::isImportedAsSPI(Identifier spiGroup,
 }
 
 bool ModuleDecl::isImportedAsWeakLinked(const ModuleDecl *module) const {
-  for (auto file : getFiles()) {
-    if (file->importsModuleAsWeakLinked(module))
-      return true;
-  }
-  return false;
+  return getASTContext().getImportCache().isWeakImportedBy(module, this);
 }
 
 bool Decl::isSPI() const {