[NFC] Hide LoadedModules From Clients of ASTContext

There's no reason clients need to be able to access this data directly.
It obscures where module loading is actually happening, and makes it too
easy to accidentally register a module with the wrong identifier in the
context.

Hide the registration operations behind opaque accessors.

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