[clang][modules] Avoid calling expensive `SourceManager::translateFile()` (#86216)

The `ASTWriter` algorithm for computing affecting module maps uses
`SourceManager::translateFile()` to get a `FileID` from a `FileEntry`.
This is slow (O(n)) since the function performs a linear walk over
`SLocEntries` until it finds one with a matching `FileEntry`.

This patch removes this use of `SourceManager::translateFile()` by
tracking `FileID` instead of `FileEntry` in couple of places in
`ModuleMap`, giving `ASTWriter` the desired `FileID` directly. There are
no changes required for clients that still want a `FileEntry` from
`ModuleMap`: the existing APIs internally use `SourceManager` to perform
the reverse `FileID` to `FileEntry` conversion in O(1).

1 files changed, 17 insertions, 2 deletions

diff --git a/clang/lib/Frontend/CompilerInstance.cpp b/clang/lib/Frontend/CompilerInstance.cpp
index 019f847..79ebb0a 100644
--- a/clang/lib/Frontend/CompilerInstance.cpp
+++ b/clang/lib/Frontend/CompilerInstance.cpp
@@ -1337,9 +1337,24 @@ static bool compileModule(CompilerInstance &ImportingInstance,
   // Get or create the module map that we'll use to build this module.
   ModuleMap &ModMap
     = ImportingInstance.getPreprocessor().getHeaderSearchInfo().getModuleMap();
+  SourceManager &SourceMgr = ImportingInstance.getSourceManager();
   bool Result;
-  if (OptionalFileEntryRef ModuleMapFile =
-          ModMap.getContainingModuleMapFile(Module)) {
+  if (FileID ModuleMapFID = ModMap.getContainingModuleMapFileID(Module);
+      ModuleMapFID.isValid()) {
+    // We want to use the top-level module map. If we don't, the compiling
+    // instance may think the containing module map is a top-level one, while
+    // the importing instance knows it's included from a parent module map via
+    // the extern directive. This mismatch could bite us later.
+    SourceLocation Loc = SourceMgr.getIncludeLoc(ModuleMapFID);
+    while (Loc.isValid() && isModuleMap(SourceMgr.getFileCharacteristic(Loc))) {
+      ModuleMapFID = SourceMgr.getFileID(Loc);
+      Loc = SourceMgr.getIncludeLoc(ModuleMapFID);
+    }
+
+    OptionalFileEntryRef ModuleMapFile =
+        SourceMgr.getFileEntryRefForID(ModuleMapFID);
+    assert(ModuleMapFile && "Top-level module map with no FileID");
+
     // Canonicalize compilation to start with the public module map. This is
     // vital for submodules declarations in the private module maps to be
     // correctly parsed when depending on a top level module in the public one.