From a1dcfb75ea8c31dd39edb6bdab6f54cde81cad85 Mon Sep 17 00:00:00 2001
From: Nathan Sidwell <nathan@acm.org>
Date: Tue, 22 Mar 2022 10:49:08 -0700
Subject: [PATCH] [clang] Module global init mangling

C++20 modules require emission of an initializer function, which is
called by importers of the module.  This implements the mangling for
that function.  It is the one place the ABI exposes partition names in
symbols -- but fortunately only needed by other TUs of that same module.

Reviewed By: bruno

Differential Revision: https://reviews.llvm.org/D122741
---
 clang/include/clang/AST/Mangle.h |  2 ++
 clang/lib/AST/ItaniumMangle.cpp  | 31 ++++++++++++++++++++++++++-----
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/clang/include/clang/AST/Mangle.h b/clang/include/clang/AST/Mangle.h
index 9bca97a611d0..96cc8c90a8e8 100644
--- a/clang/include/clang/AST/Mangle.h
+++ b/clang/include/clang/AST/Mangle.h
@@ -199,6 +199,8 @@ public:
 
   virtual void mangleDynamicStermFinalizer(const VarDecl *D, raw_ostream &) = 0;
 
+  virtual void mangleModuleInitializer(const Module *Module, raw_ostream &) = 0;
+
   // This has to live here, otherwise the CXXNameMangler won't have access to
   // it.
   virtual DiscriminatorOverrideTy getDiscriminatorOverride() const = 0;
diff --git a/clang/lib/AST/ItaniumMangle.cpp b/clang/lib/AST/ItaniumMangle.cpp
index b380e02fc8f7..a9416397c90d 100644
--- a/clang/lib/AST/ItaniumMangle.cpp
+++ b/clang/lib/AST/ItaniumMangle.cpp
@@ -130,6 +130,8 @@ public:
 
   void mangleLambdaSig(const CXXRecordDecl *Lambda, raw_ostream &) override;
 
+  void mangleModuleInitializer(const Module *Module, raw_ostream &) override;
+
   bool getNextDiscriminator(const NamedDecl *ND, unsigned &disc) {
     // Lambda closure types are already numbered.
     if (isLambda(ND))
@@ -438,7 +440,7 @@ public:
   void mangleType(QualType T);
   void mangleNameOrStandardSubstitution(const NamedDecl *ND);
   void mangleLambdaSig(const CXXRecordDecl *Lambda);
-  void mangleModuleNamePrefix(StringRef Name);
+  void mangleModuleNamePrefix(StringRef Name, bool IsPartition = false);
 
 private:
 
@@ -1057,8 +1059,8 @@ void CXXNameMangler::mangleModuleName(const NamedDecl *ND) {
 //		 ::= <module-name> <module-subname>
 //	 	 ::= <substitution>
 // <module-subname> ::= W <source-name>
-//		    ::= W P <source-name> # not (yet) needed
-void CXXNameMangler::mangleModuleNamePrefix(StringRef Name) {
+//		    ::= W P <source-name>
+void CXXNameMangler::mangleModuleNamePrefix(StringRef Name, bool IsPartition) {
   //  <substitution> ::= S <seq-id> _
   auto It = ModuleSubstitutions.find(Name);
   if (It != ModuleSubstitutions.end()) {
@@ -1072,10 +1074,14 @@ void CXXNameMangler::mangleModuleNamePrefix(StringRef Name) {
   auto Parts = Name.rsplit('.');
   if (Parts.second.empty())
     Parts.second = Parts.first;
-  else
-    mangleModuleNamePrefix(Parts.first);
+  else {
+    mangleModuleNamePrefix(Parts.first, IsPartition);
+    IsPartition = false;
+  }
 
   Out << 'W';
+  if (IsPartition)
+    Out << 'P';
   Out << Parts.second.size() << Parts.second;
   ModuleSubstitutions.insert({Name, SeqID++});
 }
@@ -6533,6 +6539,21 @@ void ItaniumMangleContextImpl::mangleLambdaSig(const CXXRecordDecl *Lambda,
   Mangler.mangleLambdaSig(Lambda);
 }
 
+void ItaniumMangleContextImpl::mangleModuleInitializer(const Module *M,
+                                                       raw_ostream &Out) {
+  // <special-name> ::= GI <module-name>  # module initializer function
+  CXXNameMangler Mangler(*this, Out);
+  Mangler.getStream() << "_ZGI";
+  Mangler.mangleModuleNamePrefix(M->getPrimaryModuleInterfaceName());
+  if (M->isModulePartition()) {
+    // The partition needs including, as partitions can have them too.
+    auto Partition = M->Name.find(':');
+    Mangler.mangleModuleNamePrefix(
+        StringRef(&M->Name[Partition + 1], M->Name.size() - Partition - 1),
+        /*IsPartition*/ true);
+  }
+}
+
 ItaniumMangleContext *ItaniumMangleContext::create(ASTContext &Context,
                                                    DiagnosticsEngine &Diags,
                                                    bool IsAux) {