[clang] [WIP][C++20][Modules] Lazily, but fully load 'HeaderFileInfo' table into memory. (PR #140867)

[Michael Park via cfe-commits]

https://github.com/mpark updated https://github.com/llvm/llvm-project/pull/140867

>From 1013ab2297935cc3a7eccf8a812984d9b454a27f Mon Sep 17 00:00:00 2001
From: Michael Park <mcypark at gmail.com>
Date: Tue, 20 May 2025 17:18:36 -0700
Subject: [PATCH] Lazily, but fully load 'HeaderFileInfo' table into memory.

---
 clang/include/clang/Serialization/ASTReader.h |  4 +
 clang/lib/Serialization/ASTReader.cpp         | 80 ++++++++++++-------
 clang/lib/Serialization/ASTReaderInternals.h  |  3 +-
 3 files changed, 54 insertions(+), 33 deletions(-)

diff --git a/clang/include/clang/Serialization/ASTReader.h b/clang/include/clang/Serialization/ASTReader.h
index 57b0266af26bb..9b63ad1a7cde7 100644
--- a/clang/include/clang/Serialization/ASTReader.h
+++ b/clang/include/clang/Serialization/ASTReader.h
@@ -654,6 +654,10 @@ class ASTReader
   /// Map from the TU to its lexical contents from each module file.
   std::vector<std::pair<ModuleFile*, LexicalContents>> TULexicalDecls;
 
+  unsigned HeaderFileInfoIdx = 0;
+
+  llvm::DenseMap<FileEntryRef, HeaderFileInfo> HeaderFileInfoLookup;
+
   /// Map from a DeclContext to its lookup tables.
   llvm::DenseMap<const DeclContext *,
                  serialization::reader::DeclContextLookupTable> Lookups;
diff --git a/clang/lib/Serialization/ASTReader.cpp b/clang/lib/Serialization/ASTReader.cpp
index d068f5e163176..93766be18f775 100644
--- a/clang/lib/Serialization/ASTReader.cpp
+++ b/clang/lib/Serialization/ASTReader.cpp
@@ -4155,6 +4155,7 @@ llvm::Error ASTReader::ReadASTBlock(ModuleFile &F,
       if (Record[0]) {
         F.HeaderFileInfoTable = HeaderFileInfoLookupTable::Create(
             (const unsigned char *)F.HeaderFileInfoTableData + Record[0],
+            (const unsigned char *)F.HeaderFileInfoTableData + sizeof(uint32_t),
             (const unsigned char *)F.HeaderFileInfoTableData,
             HeaderFileInfoTrait(*this, F));
 
@@ -6831,43 +6832,60 @@ std::optional<bool> ASTReader::isPreprocessedEntityInFileID(unsigned Index,
     return false;
 }
 
-namespace {
-
-  /// Visitor used to search for information about a header file.
-  class HeaderFileInfoVisitor {
-  FileEntryRef FE;
-    std::optional<HeaderFileInfo> HFI;
-
-  public:
-    explicit HeaderFileInfoVisitor(FileEntryRef FE) : FE(FE) {}
-
-    bool operator()(ModuleFile &M) {
-      HeaderFileInfoLookupTable *Table
-        = static_cast<HeaderFileInfoLookupTable *>(M.HeaderFileInfoTable);
-      if (!Table)
-        return false;
+static void mergeHeaderFileInfoModuleBits(HeaderFileInfo &HFI,
+                                          bool isModuleHeader,
+                                          bool isTextualModuleHeader) {
+  HFI.isModuleHeader |= isModuleHeader;
+  if (HFI.isModuleHeader)
+    HFI.isTextualModuleHeader = false;
+  else
+    HFI.isTextualModuleHeader |= isTextualModuleHeader;
+}
 
-      // Look in the on-disk hash table for an entry for this file name.
-      HeaderFileInfoLookupTable::iterator Pos = Table->find(FE);
-      if (Pos == Table->end())
-        return false;
+/// Merge the header file info provided by \p OtherHFI into the current
+/// header file info (\p HFI)
+static void mergeHeaderFileInfo(HeaderFileInfo &HFI,
+                                const HeaderFileInfo &OtherHFI) {
+  assert(OtherHFI.External && "expected to merge external HFI");
 
-      HFI = *Pos;
-      return true;
-    }
+  HFI.isImport |= OtherHFI.isImport;
+  HFI.isPragmaOnce |= OtherHFI.isPragmaOnce;
+  mergeHeaderFileInfoModuleBits(HFI, OtherHFI.isModuleHeader,
+                                OtherHFI.isTextualModuleHeader);
 
-    std::optional<HeaderFileInfo> getHeaderFileInfo() const { return HFI; }
-  };
+  if (!HFI.LazyControllingMacro.isValid())
+    HFI.LazyControllingMacro = OtherHFI.LazyControllingMacro;
 
-} // namespace
+  HFI.DirInfo = OtherHFI.DirInfo;
+  HFI.External = (!HFI.IsValid || HFI.External);
+  HFI.IsValid = true;
+}
 
 HeaderFileInfo ASTReader::GetHeaderFileInfo(FileEntryRef FE) {
-  HeaderFileInfoVisitor Visitor(FE);
-  ModuleMgr.visit(Visitor);
-  if (std::optional<HeaderFileInfo> HFI = Visitor.getHeaderFileInfo())
-      return *HFI;
-
-  return HeaderFileInfo();
+  for (auto Iter = ModuleMgr.begin() + HeaderFileInfoIdx, End = ModuleMgr.end();
+       Iter != End; ++Iter) {
+    if (auto *Table = static_cast<HeaderFileInfoLookupTable *>(
+            Iter->HeaderFileInfoTable)) {
+      auto &Info = Table->getInfoObj();
+      for (auto Iter = Table->data_begin(), End = Table->data_end();
+           Iter != End; ++Iter) {
+        const auto *Item = Iter.getItem();
+        // Determine the length of the key and the data.
+        const auto& [KeyLen, DataLen] =
+            HeaderFileInfoTrait::ReadKeyDataLength(Item);
+        // Read the key.
+        const auto &Key = Info.ReadKey(Item, KeyLen);
+        if (auto EKey = Info.getFile(Key)) {
+          auto Data = Info.ReadData(Key, Item + KeyLen, DataLen);
+          auto [Iter, Inserted] = HeaderFileInfoLookup.try_emplace(*EKey, Data);
+          if (!Inserted)
+            mergeHeaderFileInfo(Iter->second, Data);
+        }
+      }
+    }
+  }
+  HeaderFileInfoIdx = ModuleMgr.size();
+  return HeaderFileInfoLookup.lookup(FE);
 }
 
 void ASTReader::ReadPragmaDiagnosticMappings(DiagnosticsEngine &Diag) {
diff --git a/clang/lib/Serialization/ASTReaderInternals.h b/clang/lib/Serialization/ASTReaderInternals.h
index 4a7794889b039..d78d9fee6042d 100644
--- a/clang/lib/Serialization/ASTReaderInternals.h
+++ b/clang/lib/Serialization/ASTReaderInternals.h
@@ -402,13 +402,12 @@ class HeaderFileInfoTrait {
 
   data_type ReadData(internal_key_ref,const unsigned char *d, unsigned DataLen);
 
-private:
   OptionalFileEntryRef getFile(const internal_key_type &Key);
 };
 
 /// The on-disk hash table used for known header files.
 using HeaderFileInfoLookupTable =
-    llvm::OnDiskChainedHashTable<HeaderFileInfoTrait>;
+    llvm::OnDiskIterableChainedHashTable<HeaderFileInfoTrait>;
 
 } // namespace reader