Merge pull request #32418 from CodaFi/load-time-lift

[NFC] Hide LoadedModules From Clients of ASTContext

Author: CodaFi

include/swift/AST/ASTContext.h

@@ -254,11 +254,6 @@ class ASTContext final {
   /// The request-evaluator that is used to process various requests.
   Evaluator evaluator;
 
-  /// The set of top-level modules we have loaded.
-  /// This map is used for iteration, therefore it's a MapVector and not a
-  /// DenseMap.
-  llvm::MapVector<Identifier, ModuleDecl*> LoadedModules;
-
   /// The builtin module.
   ModuleDecl * const TheBuiltinModule;
 
@@ -787,8 +782,24 @@ class ASTContext final {
   /// The loader is owned by the AST context.
   ClangModuleLoader *getDWARFModuleLoader() const;
 
+public:
   namelookup::ImportCache &getImportCache() const;
 
+  /// Returns an iterator over the modules that are known by this context
+  /// to be loaded.
+  ///
+  /// Iteration order is guaranteed to match the order in which
+  /// \c addLoadedModule was called to register the loaded module
+  /// with this context.
+  iterator_range<llvm::MapVector<Identifier, ModuleDecl *>::const_iterator>
+  getLoadedModules() const;
+
+  /// Returns the number of loaded modules known by this context to be loaded.
+  unsigned getNumLoadedModules() const {
+    auto eltRange = getLoadedModules();
+    return std::distance(eltRange.begin(), eltRange.end());
+  }
+
   /// Asks every module loader to verify the ASTs it has loaded.
   ///
   /// Does nothing in non-asserts (NDEBUG) builds.
@@ -828,6 +839,11 @@ class ASTContext final {
     return const_cast<ASTContext *>(this)->getStdlibModule(false);
   }
 
+  /// Insert an externally-sourced module into the set of known loaded modules
+  /// in this context.
+  void addLoadedModule(ModuleDecl *M);
+
+public:
   /// Retrieve the current generation number, which reflects the
   /// number of times a module import has caused mass invalidation of
   /// lookup tables.

lib/AST/ASTContext.cpp

@@ -156,6 +156,11 @@ struct ASTContext::Implementation {
   /// The set of cleanups to be called when the ASTContext is destroyed.
   std::vector<std::function<void(void)>> Cleanups;
 
+  /// The set of top-level modules we have loaded.
+  /// This map is used for iteration, therefore it's a MapVector and not a
+  /// DenseMap.
+  llvm::MapVector<Identifier, ModuleDecl *> LoadedModules;
+
   // FIXME: This is a StringMap rather than a StringSet because StringSet
   // doesn't allow passing in a pre-existing allocator.
   llvm::StringMap<Identifier::Aligner, llvm::BumpPtrAllocator&>
@@ -1538,8 +1543,18 @@ ModuleDecl *ASTContext::getLoadedModule(
   return nullptr;
 }
 
+iterator_range<llvm::MapVector<Identifier, ModuleDecl *>::const_iterator>
+ASTContext::getLoadedModules() const {
+  return {getImpl().LoadedModules.begin(), getImpl().LoadedModules.end()};
+}
+
 ModuleDecl *ASTContext::getLoadedModule(Identifier ModuleName) const {
-  return LoadedModules.lookup(ModuleName);
+  return getImpl().LoadedModules.lookup(ModuleName);
+}
+
+void ASTContext::addLoadedModule(ModuleDecl *M) {
+  assert(M);
+  getImpl().LoadedModules[M->getName()] = M;
 }
 
 static AllocationArena getArena(GenericSignature genericSig) {

lib/ClangImporter/ClangImporter.cpp

@@ -1825,9 +1825,9 @@ ModuleDecl *ClangImporter::Implementation::finishLoadingClangModule(
   if (clangModule->isSubModule()) {
     finishLoadingClangModule(clangModule->getTopLevelModule(), importLoc);
   } else {
-    ModuleDecl *&loaded = SwiftContext.LoadedModules[result->getName()];
-    if (!loaded)
-      loaded = result;
+
+    if (!SwiftContext.getLoadedModule(result->getName()))
+      SwiftContext.addLoadedModule(result);
   }
 
   return result;
@@ -3368,10 +3368,12 @@ ModuleDecl *ClangModuleUnit::getOverlayModule() const {
     if (overlay == M) {
       overlay = nullptr;
     } else {
-      auto &sharedModuleRef = Ctx.LoadedModules[M->getName()];
+      // FIXME: This bizarre and twisty invariant is due to nested
+      // re-entrancy in both clang module loading and overlay module loading.
+      auto *sharedModuleRef = Ctx.getLoadedModule(M->getName());
       assert(!sharedModuleRef || sharedModuleRef == overlay ||
              sharedModuleRef == M);
-      sharedModuleRef = overlay;
+      Ctx.addLoadedModule(overlay);
     }
 
     auto mutableThis = const_cast<ClangModuleUnit *>(this);

lib/ClangImporter/DWARFImporter.cpp

@@ -121,9 +121,8 @@ ModuleDecl *ClangImporter::Implementation::loadModuleDWARF(
   (void) namelookup::getAllImports(decl);
 
   // Register the module with the ASTContext so it is available for lookups.
-  ModuleDecl *&loaded = SwiftContext.LoadedModules[name];
-  if (!loaded)
-    loaded = decl;
+  if (!SwiftContext.getLoadedModule(name))
+    SwiftContext.addLoadedModule(decl);
 
   return decl;
 }

lib/Frontend/Frontend.cpp

@@ -751,7 +751,8 @@ ModuleDecl *CompilerInstance::getMainModule() const {
     if (Invocation.getFrontendOptions().EnableLibraryEvolution)
       MainModule->setResilienceStrategy(ResilienceStrategy::Resilient);
 
-    Context->LoadedModules[MainModule->getName()] = MainModule;
+    // Register the main module with the AST context.
+    Context->addLoadedModule(MainModule);
 
     // Create and add the module's files.
     SmallVector<FileUnit *, 16> files;
@@ -773,7 +774,7 @@ ModuleDecl *CompilerInstance::getMainModule() const {
 void CompilerInstance::setMainModule(ModuleDecl *newMod) {
   assert(newMod->isMainModule());
   MainModule = newMod;
-  Context->LoadedModules[newMod->getName()] = newMod;
+  Context->addLoadedModule(newMod);
 }
 
 void CompilerInstance::performParseAndResolveImportsOnly() {

lib/FrontendTool/FrontendTool.cpp

@@ -436,7 +436,7 @@ static bool emitLoadedModuleTraceIfNeeded(ModuleDecl *mainModule,
     importedModules.insert(import.importedModule);
 
   llvm::DenseMap<StringRef, ModuleDecl *> pathToModuleDecl;
-  for (auto &module : ctxt.LoadedModules) {
+  for (const auto &module : ctxt.getLoadedModules()) {
     ModuleDecl *loadedDecl = module.second;
     if (!loadedDecl)
       llvm::report_fatal_error("Expected loaded modules to be non-null.");
@@ -690,7 +690,7 @@ static void countStatsPostSema(UnifiedStatsReporter &Stats,
   C.NumLinkLibraries = Instance.getLinkLibraries().size();
 
   auto const &AST = Instance.getASTContext();
-  C.NumLoadedModules = AST.LoadedModules.size();
+  C.NumLoadedModules = AST.getNumLoadedModules();
 
   if (auto *D = Instance.getDependencyTracker()) {
     C.NumDependencies = D->getDependencies().size();
@@ -1229,9 +1229,9 @@ static void performEndOfPipelineActions(CompilerInstance &Instance) {
   auto action = Instance.getInvocation().getFrontendOptions().RequestedAction;
   if (FrontendOptions::shouldActionOnlyParse(action) &&
       action != FrontendOptions::ActionType::EmitImportedModules) {
-    assert(ctx.LoadedModules.size() == 1 &&
+    assert(ctx.getNumLoadedModules() == 1 &&
            "Loaded a module during parse-only");
-    assert(ctx.LoadedModules.front().second == Instance.getMainModule());
+    assert(ctx.getLoadedModules().begin()->second == Instance.getMainModule());
   }
 
   // Verify the AST for all the modules we've loaded.

lib/Sema/SourceLoader.cpp

@@ -118,7 +118,7 @@ ModuleDecl *SourceLoader::loadModule(SourceLoc importLoc,
   auto *importMod = ModuleDecl::create(moduleID.Item, Ctx, importInfo);
   if (EnableLibraryEvolution)
     importMod->setResilienceStrategy(ResilienceStrategy::Resilient);
-  Ctx.LoadedModules[moduleID.Item] = importMod;
+  Ctx.addLoadedModule(importMod);
 
   auto *importFile =
       new (Ctx) SourceFile(*importMod, SourceFileKind::Library, bufferID,

lib/Serialization/SerializedModuleLoader.cpp

@@ -959,7 +959,7 @@ SerializedModuleLoaderBase::loadModule(SourceLoc importLoc,
 
   auto M = ModuleDecl::create(moduleID.Item, Ctx);
   M->setIsSystemModule(isSystemModule);
-  Ctx.LoadedModules[moduleID.Item] = M;
+  Ctx.addLoadedModule(M);
   SWIFT_DEFER { M->setHasResolvedImports(); };
 
   StringRef moduleInterfacePathStr =
@@ -1010,7 +1010,7 @@ MemoryBufferSerializedModuleLoader::loadModule(SourceLoc importLoc,
     return nullptr;
 
   M->addFile(*file);
-  Ctx.LoadedModules[moduleID.Item] = M;
+  Ctx.addLoadedModule(M);
   return M;
 }
 

lib/Serialization/SerializedSILLoader.cpp

@@ -27,7 +27,7 @@ SerializedSILLoader::SerializedSILLoader(
 
   // Get a list of SerializedModules from ASTContext.
   // FIXME: Iterating over LoadedModules is not a good way to do this.
-  for (auto &Entry : Ctx.LoadedModules) {
+  for (const auto &Entry : Ctx.getLoadedModules()) {
     for (auto File : Entry.second->getFiles()) {
       if (auto LoadedAST = dyn_cast<SerializedASTFile>(File)) {
         auto Des = new SILDeserializer(&LoadedAST->File, *SILMod, callbacks);

unittests/AST/TestContext.cpp

@@ -40,7 +40,7 @@ TestContext::TestContext(ShouldDeclareOptionalTypes optionals)
   registerTypeCheckerRequestFunctions(Ctx.evaluator);
   auto stdlibID = Ctx.getIdentifier(STDLIB_NAME);
   auto *module = ModuleDecl::create(stdlibID, Ctx);
-  Ctx.LoadedModules[stdlibID] = module;
+  Ctx.addLoadedModule(module);
 
   FileForLookups = new (Ctx) SourceFile(*module, SourceFileKind::Library,
                                         /*buffer*/ None);