Lazily, but fully load 'HeaderFileInfo' table into memory.

Author: mpark

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;

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

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