From e056c73fb1a449202b2ea4a3c16e52f1db9f780c Mon Sep 17 00:00:00 2001
From: Allan Shortlidge <shortlidge@apple.com>
Date: Thu, 1 Feb 2024 16:53:51 -0800
Subject: [PATCH] AST: Cache underlying clang module during import resolution.

As recommended in feedback on https://github.com/apple/swift/pull/71302, cache
the underlying clang module after loading it in `ImportResolver`, rather than
filtering it out of the overall set of resolved imports. This is more efficient
and results in less duplicated code that must identify the underlying clang
module.
---
 include/swift/AST/SourceFile.h |  7 +++++++
 lib/AST/Module.cpp             | 14 --------------
 lib/Sema/ImportResolution.cpp  | 15 +++++++++++++--
 3 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/include/swift/AST/SourceFile.h b/include/swift/AST/SourceFile.h
index e40b4e52a8b53..ee4874f5ae5b9 100644
--- a/include/swift/AST/SourceFile.h
+++ b/include/swift/AST/SourceFile.h
@@ -406,6 +406,13 @@ class SourceFile final : public FileUnit {
   /// resolution.
   void setImports(ArrayRef<AttributedImport<ImportedModule>> imports);
 
+  /// Set the imported underlying clang module for this source file. This gets
+  /// called by import resolution.
+  void setImportedUnderlyingModule(ModuleDecl *module) {
+    assert(!ImportedUnderlyingModule && "underlying module already set");
+    ImportedUnderlyingModule = module;
+  }
+
   /// Whether the given import has used @preconcurrency.
   bool hasImportUsedPreconcurrency(
       AttributedImport<ImportedModule> import) const;
diff --git a/lib/AST/Module.cpp b/lib/AST/Module.cpp
index 30386b3c08382..4da2e42b75abe 100644
--- a/lib/AST/Module.cpp
+++ b/lib/AST/Module.cpp
@@ -2566,20 +2566,6 @@ void
 SourceFile::setImports(ArrayRef<AttributedImport<ImportedModule>> imports) {
   assert(!Imports && "Already computed imports");
   Imports = getASTContext().AllocateCopy(imports);
-
-  // Find and cache the import of the underlying module, if present.
-  auto parentModuleName = getParentModule()->getName();
-  for (auto import : imports) {
-    if (!import.options.contains(ImportFlags::Exported))
-      continue;
-
-    auto importedModule = import.module.importedModule;
-    if (importedModule->getName() == parentModuleName &&
-        importedModule->findUnderlyingClangModule()) {
-      ImportedUnderlyingModule = import.module.importedModule;
-      break;
-    }
-  }
 }
 
 bool SourceFile::hasImportUsedPreconcurrency(
diff --git a/lib/Sema/ImportResolution.cpp b/lib/Sema/ImportResolution.cpp
index cd4fbb8312117..f3f07b7179be3 100644
--- a/lib/Sema/ImportResolution.cpp
+++ b/lib/Sema/ImportResolution.cpp
@@ -177,6 +177,10 @@ class ImportResolver final : public DeclVisitor<ImportResolver> {
   /// much, much smaller than \c crossImportableModules.
   SmallVector<AttributedImport<ImportedModule>, 16> crossImportDeclaringModules;
 
+  /// The underlying clang module of the source file's parent module, if
+  /// imported.
+  ModuleDecl *underlyingClangModule = nullptr;
+
   /// The index of the next module in \c visibleModules that should be
   /// cross-imported.
   size_t nextModuleToCrossImport = 0;
@@ -193,6 +197,10 @@ class ImportResolver final : public DeclVisitor<ImportResolver> {
     return boundImports;
   }
 
+  /// Retrieve the underlying clang module which will be cached if it was loaded
+  /// when resolving imports.
+  ModuleDecl *getUnderlyingClangModule() const { return underlyingClangModule; }
+
 private:
   // We only need to visit import decls.
   void visitImportDecl(ImportDecl *ID);
@@ -287,6 +295,7 @@ void swift::performImportResolution(SourceFile &SF) {
     resolver.visit(D);
 
   SF.setImports(resolver.getFinishedImports());
+  SF.setImportedUnderlyingModule(resolver.getUnderlyingClangModule());
 
   SF.ASTStage = SourceFile::ImportsResolved;
   verify(SF);
@@ -396,8 +405,10 @@ ImportResolver::getModule(ImportPath::Module modulePath) {
   // for clang overlays as well.
   if (ctx.getRealModuleName(moduleID.Item) == loadingModule->getName() &&
       modulePath.size() == 1) {
-    if (auto importer = ctx.getClangModuleLoader())
-      return importer->loadModule(moduleID.Loc, modulePath);
+    if (auto importer = ctx.getClangModuleLoader()) {
+      underlyingClangModule = importer->loadModule(moduleID.Loc, modulePath);
+      return underlyingClangModule;
+    }
     return nullptr;
   }