[C++20][Modules] Improve namespace look-up performance for modules. (Attempt #2)

clang/include/clang/Serialization/ASTWriter.h

@@ -773,13 +773,18 @@ class ASTWriter : public ASTDeserializationListener,
   /// Is this a local declaration (that is, one that will be written to
   /// our AST file)? This is the case for declarations that are neither imported
   /// from another AST file nor predefined.
-  bool IsLocalDecl(const Decl *D) {
+  bool IsLocalDecl(const Decl *D) const {
     if (D->isFromASTFile())
       return false;
     auto I = DeclIDs.find(D);
     return (I == DeclIDs.end() || I->second >= clang::NUM_PREDEF_DECL_IDS);
   };
 
+  /// Collect the first declaration from each module file that provides a
+  /// declaration of D.
+  llvm::MapVector<serialization::ModuleFile *, const Decl *>
+  CollectFirstDeclFromEachModule(const Decl *D, bool IncludeLocal);
+
   void AddLookupOffsets(const LookupBlockOffsets &Offsets,
                         RecordDataImpl &Record);
 

clang/lib/Serialization/ASTReader.cpp

@@ -555,7 +555,25 @@ namespace {
 
 using MacroDefinitionsMap =
     llvm::StringMap<std::pair<StringRef, bool /*IsUndef*/>>;
-using DeclsMap = llvm::DenseMap<DeclarationName, SmallVector<NamedDecl *, 8>>;
+
+class DeclsSet {
+  SmallVector<NamedDecl *, 64> Decls;
+  llvm::SmallPtrSet<NamedDecl *, 8> Found;
+
+public:
+  operator ArrayRef<NamedDecl *>() const { return Decls; }
+
+  bool empty() const { return Decls.empty(); }
+
+  bool insert(NamedDecl *ND) {
+    auto [_, Inserted] = Found.insert(ND);
+    if (Inserted)
+      Decls.push_back(ND);
+    return Inserted;
+  }
+};
+
+using DeclsMap = llvm::DenseMap<DeclarationName, DeclsSet>;
 
 } // namespace
 
@@ -8729,14 +8747,23 @@ bool ASTReader::FindExternalVisibleDeclsByName(const DeclContext *DC,
     return false;
 
   // Load the list of declarations.
-  SmallVector<NamedDecl *, 64> Decls;
-  llvm::SmallPtrSet<NamedDecl *, 8> Found;
+  DeclsSet DS;
 
   auto Find = [&, this](auto &&Table, auto &&Key) {
     for (GlobalDeclID ID : Table.find(Key)) {
       NamedDecl *ND = cast<NamedDecl>(GetDecl(ID));
-      if (ND->getDeclName() == Name && Found.insert(ND).second)
-        Decls.push_back(ND);
+      if (ND->getDeclName() != Name)
+        continue;
+      // Special case for namespaces: There can be a lot of redeclarations of
+      // some namespaces, and we import a "key declaration" per imported module.
+      // Since all declarations of a namespace are essentially interchangeable,
+      // we can optimize namespace look-up by only storing the key declaration
+      // of the current TU, rather than storing N key declarations where N is
+      // the # of imported modules that declare that namespace.
+      // TODO: Try to generalize this optimization to other redeclarable decls.
+      if (isa<NamespaceDecl>(ND))
+        ND = cast<NamedDecl>(getKeyDeclaration(ND));
+      DS.insert(ND);
     }
   };
 
@@ -8771,8 +8798,8 @@ bool ASTReader::FindExternalVisibleDeclsByName(const DeclContext *DC,
     Find(It->second.Table, Name);
   }
 
-  SetExternalVisibleDeclsForName(DC, Name, Decls);
-  return !Decls.empty();
+  SetExternalVisibleDeclsForName(DC, Name, DS);
+  return !DS.empty();
 }
 
 void ASTReader::completeVisibleDeclsMap(const DeclContext *DC) {
@@ -8790,7 +8817,16 @@ void ASTReader::completeVisibleDeclsMap(const DeclContext *DC) {
 
     for (GlobalDeclID ID : It->second.Table.findAll()) {
       NamedDecl *ND = cast<NamedDecl>(GetDecl(ID));
-      Decls[ND->getDeclName()].push_back(ND);
+      // Special case for namespaces: There can be a lot of redeclarations of
+      // some namespaces, and we import a "key declaration" per imported module.
+      // Since all declarations of a namespace are essentially interchangeable,
+      // we can optimize namespace look-up by only storing the key declaration
+      // of the current TU, rather than storing N key declarations where N is
+      // the # of imported modules that declare that namespace.
+      // TODO: Try to generalize this optimization to other redeclarable decls.
+      if (isa<NamespaceDecl>(ND))
+        ND = cast<NamedDecl>(getKeyDeclaration(ND));
+      Decls[ND->getDeclName()].insert(ND);
     }
 
     // FIXME: Why a PCH test is failing if we remove the iterator after findAll?
@@ -8800,9 +8836,9 @@ void ASTReader::completeVisibleDeclsMap(const DeclContext *DC) {
   findAll(ModuleLocalLookups, NumModuleLocalVisibleDeclContexts);
   findAll(TULocalLookups, NumTULocalVisibleDeclContexts);
 
-  for (DeclsMap::iterator I = Decls.begin(), E = Decls.end(); I != E; ++I) {
-    SetExternalVisibleDeclsForName(DC, I->first, I->second);
-  }
+  for (auto &[Name, DS] : Decls)
+    SetExternalVisibleDeclsForName(DC, Name, DS);
+
   const_cast<DeclContext *>(DC)->setHasExternalVisibleStorage(false);
 }
 

clang/lib/Serialization/ASTWriter.cpp

@@ -4399,20 +4399,20 @@ class ASTDeclContextNameLookupTrait
 
   template <typename Coll> data_type getData(const Coll &Decls) {
     unsigned Start = DeclIDs.size();
-    for (NamedDecl *D : Decls) {
+    auto AddDecl = [this](NamedDecl *D) {
       NamedDecl *DeclForLocalLookup =
           getDeclForLocalLookup(Writer.getLangOpts(), D);
 
       if (Writer.getDoneWritingDeclsAndTypes() &&
           !Writer.wasDeclEmitted(DeclForLocalLookup))
-        continue;
+        return;
 
       // Try to avoid writing internal decls to reduced BMI.
       // See comments in ASTWriter::WriteDeclContextLexicalBlock for details.
       if (Writer.isGeneratingReducedBMI() &&
           !DeclForLocalLookup->isFromExplicitGlobalModule() &&
           IsInternalDeclFromFileContext(DeclForLocalLookup))
-        continue;
+        return;
 
       auto ID = Writer.GetDeclRef(DeclForLocalLookup);
 
@@ -4426,7 +4426,7 @@ class ASTDeclContextNameLookupTrait
             ModuleLocalDeclsMap.insert({Key, DeclIDsTy{ID}});
           else
             Iter->second.push_back(ID);
-          continue;
+          return;
         }
         break;
       case LookupVisibility::TULocal: {
@@ -4435,14 +4435,32 @@ class ASTDeclContextNameLookupTrait
           TULocalDeclsMap.insert({D->getDeclName(), DeclIDsTy{ID}});
         else
           Iter->second.push_back(ID);
-        continue;
+        return;
       }
       case LookupVisibility::GenerallyVisibile:
         // Generally visible decls go into the general lookup table.
         break;
       }
 
       DeclIDs.push_back(ID);
+    };
+    ASTReader *Chain = Writer.getChain();
+    for (NamedDecl *D : Decls) {
+      if (Chain && isa<NamespaceDecl>(D) && D->isFromASTFile() &&
+          D == Chain->getKeyDeclaration(D)) {
+        // In ASTReader, we stored only the key declaration of a namespace decl
+        // for this TU. If we have an external namespace decl, this is that
+        // key declaration and we need to re-expand it to write out the first
+        // decl from each module.
+        //
+        // See comment 'ASTReader::FindExternalVisibleDeclsByName' for details.
+        auto Firsts =
+            Writer.CollectFirstDeclFromEachModule(D, /*IncludeLocal=*/false);
+        for (const auto &[_, First] : Firsts)
+          AddDecl(cast<NamedDecl>(const_cast<Decl *>(First)));
+      } else {
+        AddDecl(D);
+      }
     }
     return std::make_pair(Start, DeclIDs.size());
   }
@@ -6916,6 +6934,19 @@ TypeID ASTWriter::GetOrCreateTypeID(ASTContext &Context, QualType T) {
   });
 }
 
+llvm::MapVector<ModuleFile *, const Decl *>
+ASTWriter::CollectFirstDeclFromEachModule(const Decl *D, bool IncludeLocal) {
+  llvm::MapVector<ModuleFile *, const Decl *> Firsts;
+  // FIXME: We can skip entries that we know are implied by others.
+  for (const Decl *R = D->getMostRecentDecl(); R; R = R->getPreviousDecl()) {
+    if (R->isFromASTFile())
+      Firsts[Chain->getOwningModuleFile(R)] = R;
+    else if (IncludeLocal)
+      Firsts[nullptr] = R;
+  }
+  return Firsts;
+}
+
 void ASTWriter::AddLookupOffsets(const LookupBlockOffsets &Offsets,
                                  RecordDataImpl &Record) {
   Record.push_back(Offsets.LexicalOffset);

clang/lib/Serialization/ASTWriterDecl.cpp

@@ -194,30 +194,13 @@ namespace clang {
       Record.AddSourceLocation(typeParams->getRAngleLoc());
     }
 
-    /// Collect the first declaration from each module file that provides a
-    /// declaration of D.
-    void CollectFirstDeclFromEachModule(
-        const Decl *D, bool IncludeLocal,
-        llvm::MapVector<ModuleFile *, const Decl *> &Firsts) {
-
-      // FIXME: We can skip entries that we know are implied by others.
-      for (const Decl *R = D->getMostRecentDecl(); R; R = R->getPreviousDecl()) {
-        if (R->isFromASTFile())
-          Firsts[Writer.Chain->getOwningModuleFile(R)] = R;
-        else if (IncludeLocal)
-          Firsts[nullptr] = R;
-      }
-    }
-
     /// Add to the record the first declaration from each module file that
     /// provides a declaration of D. The intent is to provide a sufficient
     /// set such that reloading this set will load all current redeclarations.
     void AddFirstDeclFromEachModule(const Decl *D, bool IncludeLocal) {
-      llvm::MapVector<ModuleFile *, const Decl *> Firsts;
-      CollectFirstDeclFromEachModule(D, IncludeLocal, Firsts);
-
-      for (const auto &F : Firsts)
-        Record.AddDeclRef(F.second);
+      auto Firsts = Writer.CollectFirstDeclFromEachModule(D, IncludeLocal);
+      for (const auto &[_, First] : Firsts)
+        Record.AddDeclRef(First);
     }
 
     template <typename T> bool shouldSkipWritingSpecializations(T *Spec) {
@@ -272,18 +255,17 @@ namespace clang {
       assert((isa<ClassTemplateSpecializationDecl>(D) ||
               isa<VarTemplateSpecializationDecl>(D) || isa<FunctionDecl>(D)) &&
              "Must not be called with other decls");
-      llvm::MapVector<ModuleFile *, const Decl *> Firsts;
-      CollectFirstDeclFromEachModule(D, /*IncludeLocal*/ true, Firsts);
-
-      for (const auto &F : Firsts) {
-        if (shouldSkipWritingSpecializations(F.second))
+      auto Firsts =
+          Writer.CollectFirstDeclFromEachModule(D, /*IncludeLocal=*/true);
+      for (const auto &[_, First] : Firsts) {
+        if (shouldSkipWritingSpecializations(First))
           continue;
 
         if (isa<ClassTemplatePartialSpecializationDecl,
-                VarTemplatePartialSpecializationDecl>(F.second))
-          PartialSpecsInMap.push_back(F.second);
+                VarTemplatePartialSpecializationDecl>(First))
+          PartialSpecsInMap.push_back(First);
         else
-          SpecsInMap.push_back(F.second);
+          SpecsInMap.push_back(First);
       }
     }
 

clang/test/Modules/pr171769.cpp

@@ -0,0 +1,63 @@
+// RUN: rm -rf %t
+// RUN: mkdir -p %t
+// RUN: split-file %s %t
+// RUN: cd %t
+//
+
+// RUN: %clang_cc1 -fmodule-name=A -fno-cxx-modules -xc++ -emit-module \
+// RUN:   -fmodules A.cppmap -o A.pcm
+// RUN: %clang_cc1 -fmodule-name=B -fno-cxx-modules -xc++ -emit-module \
+// RUN:   -fmodules -fmodule-file=A.pcm B.cppmap -o B.pcm
+// RUN: %clang_cc1 -fmodule-name=C -fno-cxx-modules -xc++ -emit-module \
+// RUN:   -fmodules C.cppmap -o C.pcm
+// RUN: %clang_cc1 -fmodule-name=D -fno-cxx-modules -xc++ -emit-module \
+// RUN:   -fmodules -fmodule-file=C.pcm D.cppmap -o D.pcm
+// RUN: %clang_cc1 -fmodule-name=E -fno-cxx-modules -xc++ -emit-module \
+// RUN:   -fmodules -fmodule-file=B.pcm E.cppmap -o E.pcm
+// RUN: %clang_cc1 -fmodule-name=F -fno-cxx-modules -xc++ -emit-module \
+// RUN:   -fmodules -fmodule-file=D.pcm F.cppmap -o F.pcm
+// RUN: %clang_cc1 -fmodule-name=G -fno-cxx-modules -xc++ -emit-module \
+// RUN:   -fmodules -fmodule-file=E.pcm -fmodule-file=F.pcm G.cppmap -o G.pcm
+// RUN: %clang_cc1 -fno-cxx-modules -fmodules -fmodule-file=G.pcm src.cpp \
+// RUN:   -o /dev/null
+
+//--- A.cppmap
+module "A" { header "A.h" }
+
+//--- A.h
+int x;
+
+//--- B.cppmap
+module "B" {}
+
+//--- C.cppmap
+module "C" { header "C.h" }
+
+//--- C.h
+namespace xyz {}
+
+//--- D.cppmap
+module "D" {}
+
+//--- E.cppmap
+module "E" {}
+
+//--- F.cppmap
+module "F" { header "F.h" }
+
+//--- F.h
+namespace xyz { inline void func() {} }
+
+//--- G.cppmap
+module "G" { header "G.h" }
+
+//--- G.h
+#include "F.h"
+namespace { void func2() { xyz::func(); } }
+
+//--- hdr.h
+#include "F.h"
+namespace xyz_ns = xyz;
+
+//--- src.cpp
+#include "hdr.h"

clang/unittests/Serialization/CMakeLists.txt

@@ -2,6 +2,7 @@ add_clang_unittest(SerializationTests
   ForceCheckFileInputTest.cpp
   InMemoryModuleCacheTest.cpp
   ModuleCacheTest.cpp
+  NamespaceLookupTest.cpp
   NoCommentsTest.cpp
   PreambleInNamedModulesTest.cpp
   LoadSpecLazilyTest.cpp