clang/Modules: Sink CompilerInstance::KnownModules into ModuleMap

Avoid use-after-frees when FrontendAction::BeginSourceFile is called
twice on the same CompilerInstance by sinking
CompilerInstance::KnownModules into ModuleMap.  On the way, rename the
map to CachedModuleLoads.  I considered (but rejected) merging this with
ModuleMap::Modules, since that only has top-level modules and this map
includes submodules.

This is an alternative to https://reviews.llvm.org/D58497.  Thanks to
nemanjai for the detailed analysis of the problem!

Author: dexonsmith

clang/include/clang/Frontend/CompilerInstance.h

@@ -126,10 +126,6 @@ class CompilerInstance : public ModuleLoader {
 
   std::vector<std::shared_ptr<DependencyCollector>> DependencyCollectors;
 
-  /// The set of top-level modules that has already been loaded,
-  /// along with the module map
-  llvm::DenseMap<const IdentifierInfo *, Module *> KnownModules;
-
   /// The set of top-level modules that has already been built on the
   /// fly as part of this overall compilation action.
   std::map<std::string, std::string> BuiltModules;

clang/include/clang/Lex/ModuleMap.h

@@ -14,17 +14,18 @@
 #ifndef LLVM_CLANG_LEX_MODULEMAP_H
 #define LLVM_CLANG_LEX_MODULEMAP_H
 
+#include "clang/Basic/IdentifierTable.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/Module.h"
 #include "clang/Basic/SourceLocation.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/PointerIntPair.h"
-#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/ADT/Twine.h"
 #include <ctime>
@@ -108,6 +109,10 @@ class ModuleMap {
   /// The top-level modules that are known.
   llvm::StringMap<Module *> Modules;
 
+  /// Module loading cache that includes submodules, indexed by IdentifierInfo.
+  /// nullptr is stored for modules that are known to fail to load.
+  llvm::DenseMap<const IdentifierInfo *, Module *> CachedModuleLoads;
+
   /// Shadow modules created while building this module map.
   llvm::SmallVector<Module*, 2> ShadowModules;
 
@@ -695,6 +700,16 @@ class ModuleMap {
 
   module_iterator module_begin() const { return Modules.begin(); }
   module_iterator module_end()   const { return Modules.end(); }
+
+  /// Cache a module load.  M might be nullptr.
+  void cacheModuleLoad(const IdentifierInfo &II, Module *M) {
+    CachedModuleLoads[&II] = M;
+  }
+
+  /// Return a cached module load.
+  Module *getCachedModuleLoad(const IdentifierInfo &II) {
+    return CachedModuleLoads.lookup(&II);
+  }
 };
 
 } // namespace clang

clang/lib/Frontend/CompilerInstance.cpp

@@ -1560,12 +1560,9 @@ bool CompilerInstance::loadModuleFile(StringRef FileName) {
     }
 
     void registerAll() {
-      for (auto *II : LoadedModules) {
-        CI.KnownModules[II] = CI.getPreprocessor()
-                                  .getHeaderSearchInfo()
-                                  .getModuleMap()
-                                  .findModule(II->getName());
-      }
+      ModuleMap &MM = CI.getPreprocessor().getHeaderSearchInfo().getModuleMap();
+      for (auto *II : LoadedModules)
+        MM.cacheModuleLoad(*II, MM.findModule(II->getName()));
       LoadedModules.clear();
     }
 
@@ -1654,14 +1651,11 @@ CompilerInstance::loadModule(SourceLocation ImportLoc,
     return LastModuleImportResult;
   }
 
-  clang::Module *Module = nullptr;
-
   // If we don't already have information on this module, load the module now.
-  llvm::DenseMap<const IdentifierInfo *, clang::Module *>::iterator Known
-    = KnownModules.find(Path[0].first);
-  if (Known != KnownModules.end()) {
-    // Retrieve the cached top-level module.
-    Module = Known->second;
+  ModuleMap &MM = getPreprocessor().getHeaderSearchInfo().getModuleMap();
+  clang::Module *Module = MM.getCachedModuleLoad(*Path[0].first);
+  if (Module) {
+    // Nothing to do here, we found it.
   } else if (ModuleName == getLangOpts().CurrentModule) {
     // This is the module we're building.
     Module = PP->getHeaderSearchInfo().lookupModule(
@@ -1675,7 +1669,7 @@ CompilerInstance::loadModule(SourceLocation ImportLoc,
     //  ModuleBuildFailed = true;
     //  return ModuleLoadResult();
     //}
-    Known = KnownModules.insert(std::make_pair(Path[0].first, Module)).first;
+    MM.cacheModuleLoad(*Path[0].first, Module);
   } else {
     // Search for a module with the given name.
     Module = PP->getHeaderSearchInfo().lookupModule(ModuleName, true,
@@ -1769,7 +1763,7 @@ CompilerInstance::loadModule(SourceLocation ImportLoc,
           getDiagnostics().Report(ModuleNameLoc, diag::err_module_prebuilt)
               << ModuleName;
           ModuleBuildFailed = true;
-          KnownModules[Path[0].first] = nullptr;
+          MM.cacheModuleLoad(*Path[0].first, nullptr);
           return ModuleLoadResult();
         }
       }
@@ -1783,7 +1777,7 @@ CompilerInstance::loadModule(SourceLocation ImportLoc,
         // necessarily even have a module map. Since ReadAST already produces
         // diagnostics for these two cases, we simply error out here.
         ModuleBuildFailed = true;
-        KnownModules[Path[0].first] = nullptr;
+        MM.cacheModuleLoad(*Path[0].first, nullptr);
         return ModuleLoadResult();
       }
 
@@ -1828,7 +1822,7 @@ CompilerInstance::loadModule(SourceLocation ImportLoc,
                "undiagnosed error in compileAndLoadModule");
         if (getPreprocessorOpts().FailedModules)
           getPreprocessorOpts().FailedModules->addFailed(ModuleName);
-        KnownModules[Path[0].first] = nullptr;
+        MM.cacheModuleLoad(*Path[0].first, nullptr);
         ModuleBuildFailed = true;
         return ModuleLoadResult();
       }
@@ -1851,19 +1845,19 @@ CompilerInstance::loadModule(SourceLocation ImportLoc,
       ModuleLoader::HadFatalFailure = true;
       // FIXME: The ASTReader will already have complained, but can we shoehorn
       // that diagnostic information into a more useful form?
-      KnownModules[Path[0].first] = nullptr;
+      MM.cacheModuleLoad(*Path[0].first, nullptr);
       return ModuleLoadResult();
 
     case ASTReader::Failure:
       ModuleLoader::HadFatalFailure = true;
       // Already complained, but note now that we failed.
-      KnownModules[Path[0].first] = nullptr;
+      MM.cacheModuleLoad(*Path[0].first, nullptr);
       ModuleBuildFailed = true;
       return ModuleLoadResult();
     }
 
     // Cache the result of this top-level module lookup for later.
-    Known = KnownModules.insert(std::make_pair(Path[0].first, Module)).first;
+    MM.cacheModuleLoad(*Path[0].first, Module);
   }
 
   // If we never found the module, fail.