AST: Cache underlying clang module during import resolution

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;

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(

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;
   }