6 files changed, 14 insertions, 28 deletions
diff --git a/clang/include/clang/AST/ASTContext.h b/clang/include/clang/AST/ASTContext.h
index 0238371..006063e 100644
--- a/clang/include/clang/AST/ASTContext.h
+++ b/clang/include/clang/AST/ASTContext.h
@@ -447,8 +447,9 @@ class ASTContext : public RefCountedBase<ASTContext> {
   };
   llvm::DenseMap<Module*, PerModuleInitializers*> ModuleInitializers;
 
-  /// For module code-gen cases, this is the top-level module we are building.
-  Module *TopLevelModule = nullptr;
+  /// For module code-gen cases, this is the top-level (C++20) Named module
+  /// we are building.
+  Module *TopLevelCXXNamedModule = nullptr;
 
   static constexpr unsigned ConstantArrayTypesLog2InitSize = 8;
   static constexpr unsigned GeneralTypesLog2InitSize = 9;
@@ -1051,10 +1052,10 @@ public:
   ArrayRef<Decl*> getModuleInitializers(Module *M);
 
   /// Set the (C++20) module we are building.
-  void setModuleForCodeGen(Module *M) { TopLevelModule = M; }
+  void setNamedModuleForCodeGen(Module *M) { TopLevelCXXNamedModule = M; }
 
   /// Get module under construction, nullptr if this is not a C++20 module.
-  Module *getModuleForCodeGen() const { return TopLevelModule; }
+  Module *getNamedModuleForCodeGen() const { return TopLevelCXXNamedModule; }
 
   TranslationUnitDecl *getTranslationUnitDecl() const {
     return TUDecl->getMostRecentDecl();
diff --git a/clang/lib/CodeGen/CGDeclCXX.cpp b/clang/lib/CodeGen/CGDeclCXX.cpp
index dcd811e..4168596 100644
--- a/clang/lib/CodeGen/CGDeclCXX.cpp
+++ b/clang/lib/CodeGen/CGDeclCXX.cpp
@@ -880,11 +880,11 @@ CodeGenModule::EmitCXXGlobalInitFunc() {
   // and makes sure these symbols appear lexicographically behind the symbols
   // with priority emitted above.
   llvm::Function *Fn;
-  if (CXX20ModuleInits && getContext().getModuleForCodeGen()) {
+  if (CXX20ModuleInits && getContext().getNamedModuleForCodeGen()) {
     SmallString<256> InitFnName;
     llvm::raw_svector_ostream Out(InitFnName);
     cast<ItaniumMangleContext>(getCXXABI().getMangleContext())
-        .mangleModuleInitializer(getContext().getModuleForCodeGen(), Out);
+        .mangleModuleInitializer(getContext().getNamedModuleForCodeGen(), Out);
     Fn = CreateGlobalInitOrCleanUpFunction(
         FTy, llvm::Twine(InitFnName), FI, SourceLocation(), false,
         llvm::GlobalVariable::ExternalLinkage);
diff --git a/clang/lib/CodeGen/CodeGenModule.cpp b/clang/lib/CodeGen/CodeGenModule.cpp
index 71ec831..71a2f61 100644
--- a/clang/lib/CodeGen/CodeGenModule.cpp
+++ b/clang/lib/CodeGen/CodeGenModule.cpp
@@ -509,7 +509,7 @@ static void setVisibilityFromDLLStorageClass(const clang::LangOptions &LO,
 }
 
 void CodeGenModule::Release() {
-  Module *Primary = getContext().getModuleForCodeGen();
+  Module *Primary = getContext().getNamedModuleForCodeGen();
   if (CXX20ModuleInits && Primary && !Primary->isHeaderLikeModule())
     EmitModuleInitializers(Primary);
   EmitDeferred();
diff --git a/clang/lib/Frontend/ASTUnit.cpp b/clang/lib/Frontend/ASTUnit.cpp
index 3b4f251..f2939a2 100644
--- a/clang/lib/Frontend/ASTUnit.cpp
+++ b/clang/lib/Frontend/ASTUnit.cpp
@@ -877,6 +877,10 @@ std::unique_ptr<ASTUnit> ASTUnit::LoadFromASTFile(
 
   PP.setCounterValue(Counter);
 
+  Module *M = HeaderInfo.lookupModule(AST->getLangOpts().CurrentModule);
+  if (M && AST->getLangOpts().isCompilingModule() && M->isModulePurview())
+    AST->Ctx->setNamedModuleForCodeGen(M);
+
   // Create an AST consumer, even though it isn't used.
   if (ToLoad >= LoadASTOnly)
     AST->Consumer.reset(new ASTConsumer);
diff --git a/clang/lib/Parse/ParseAST.cpp b/clang/lib/Parse/ParseAST.cpp
index f442b62..04b3f04 100644
--- a/clang/lib/Parse/ParseAST.cpp
+++ b/clang/lib/Parse/ParseAST.cpp
@@ -172,27 +172,6 @@ void clang::ParseAST(Sema &S, bool PrintStats, bool SkipFunctionBodies) {
   for (Decl *D : S.WeakTopLevelDecls())
     Consumer->HandleTopLevelDecl(DeclGroupRef(D));
 
-  // For C++20 modules, the codegen for module initializers needs to be altered
-  // and to be able to use a name based on the module name.
-
-  // At this point, we should know if we are building a non-header C++20 module.
-  if (S.getLangOpts().CPlusPlusModules) {
-    // If we are building the module from source, then the top level module
-    // will be here.
-    Module *CodegenModule = S.getCurrentModule();
-    bool Interface = true;
-    if (CodegenModule)
-      // We only use module initializers for importable module (including
-      // partition implementation units).
-      Interface = S.currentModuleIsInterface();
-    else if (S.getLangOpts().isCompilingModuleInterface())
-      // If we are building the module from a PCM file, then the module can be
-      // found here.
-      CodegenModule = S.getPreprocessor().getCurrentModule();
-
-    if (Interface && CodegenModule)
-      S.getASTContext().setModuleForCodeGen(CodegenModule);
-  }
   Consumer->HandleTranslationUnit(S.getASTContext());
 
   // Finalize the template instantiation observer chain.
diff --git a/clang/lib/Sema/SemaModule.cpp b/clang/lib/Sema/SemaModule.cpp
index 194239a..8f38dc8 100644
--- a/clang/lib/Sema/SemaModule.cpp
+++ b/clang/lib/Sema/SemaModule.cpp
@@ -409,6 +409,8 @@ Sema::ActOnModuleDecl(SourceLocation StartLoc, SourceLocation ModuleLoc,
     return ConvertDeclToDeclGroup(Import);
   }
 
+  getASTContext().setNamedModuleForCodeGen(Mod);
+
   // FIXME: Create a ModuleDecl.
   return nullptr;
 }