[Serialization] Read main file's HeaderFileInfo from PCH even if changed

[sammccall]

Summary

Because file size is part of the cache key, overriding the main file
content leaves us unable to load its HeaderFileInfo from a preamble.

The fix is a little messy: write a second copy with size 0 in the key.
During deserialization, if a load fails and the main filename matches,
make a second attempt with this size-0 key.
In both cases, we only do this for preambles, not for modules.

Fixes https://github.com/clangd/clangd/issues/337

clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp

runSymbolCollector(R"cpp(
 
     #endif // Header guard is recognized here.
   )cpp",
                      "");
   EXPECT_THAT(Symbols, Not(Contains(qName("HEADER_GUARD_"))));
   EXPECT_THAT(Symbols, Each(includeHeader()));
 }
 
+TEST_F(SymbolCollectorTest, HeaderGuardDetectedPragmaInPreamble) {
+  // TestTU builds with a preamble.
+  auto TU = TestTU::withCode(R"cpp(
+    #pragma once
+
+    // Symbols are seen before the header guard is complete.
+    #define MACRO
+    int decl();
+    #define MACRO2
+  )cpp");
+  TU.HeaderFilename = "Foo.h";
+  auto Symbols = TU.headerSymbols();
+  EXPECT_THAT(Symbols, Not(Contains(qName("HEADER_GUARD_"))));
+  EXPECT_THAT(Symbols, Contains(qName("MACRO")));
+  EXPECT_THAT(Symbols, Contains(qName("MACRO2")));
+  EXPECT_THAT(Symbols, Contains(qName("decl")));
+}
+
+TEST_F(SymbolCollectorTest, HeaderGuardDetectedIfdefInPreamble) {
+  // TestTU builds with a preamble.
+  auto TU = TestTU::withCode(R"cpp(
+    #ifndef HEADER_GUARD_
+    #define HEADER_GUARD_
+
+    // Symbols are seen before the header guard is complete.
+    #define MACRO
+    int decl();
+    #define MACRO2
+
+    #endif // Header guard is recognized here.
+  )cpp");
+  TU.HeaderFilename = "Foo.h";
+  auto Symbols = TU.headerSymbols();
+  EXPECT_THAT(Symbols, Not(Contains(qName("HEADER_GUARD_"))));
+  EXPECT_THAT(Symbols, Contains(qName("MACRO")));
+  EXPECT_THAT(Symbols, Contains(qName("MACRO2")));
+  EXPECT_THAT(Symbols, Contains(qName("decl")));
+}
+
 TEST_F(SymbolCollectorTest, NonModularHeader) {
   auto TU = TestTU::withHeaderCode("int x();");
   EXPECT_THAT(TU.headerSymbols(), ElementsAre(includeHeader()));
 
   // Files missing include guards aren't eligible for insertion.
   TU.ImplicitHeaderGuard = false;
   EXPECT_THAT(TU.headerSymbols(), ElementsAre(Not(includeHeader())));
 

clang/lib/Serialization/ASTReader.cpp

public:
     bool operator()(ModuleFile &M) {
       HeaderFileInfoLookupTable *Table
         = static_cast<HeaderFileInfoLookupTable *>(M.HeaderFileInfoTable);
       if (!Table)
         return false;
 
       // Look in the on-disk hash table for an entry for this file name.
       HeaderFileInfoLookupTable::iterator Pos = Table->find(FE);
+      // Preambles may be reused with different main-file content.
+      // A second entry with size zero is stored for the main-file, try that.
+      // To avoid doing this on every miss, require the bare filename to match.
+      if (Pos == Table->end() && M.Kind == clang::serialization::MK_Preamble &&
+          llvm::sys::path::filename(FE->getName()) ==
+              llvm::sys::path::filename(M.OriginalSourceFileName)) {
+        auto InternalKey = Table->getInfoObj().GetInternalKey(FE);
+        InternalKey.Size = 0;
+        Pos = Table->find_hashed(InternalKey,
+                                 Table->getInfoObj().ComputeHash(InternalKey));
+      }
       if (Pos == Table->end())
         return false;
 
       HFI = *Pos;
       return true;
     }
 
     std::optional<HeaderFileInfo> getHeaderFileInfo() const { return HFI; }

clang/lib/Serialization/ASTWriter.cpp

while (!Worklist.empty()) {
       }
       auto SubmodulesRange = M->submodules();
       Worklist.append(SubmodulesRange.begin(), SubmodulesRange.end());
     }
   }
 
   SmallVector<const FileEntry *, 16> FilesByUID;
   HS.getFileMgr().GetUniqueIDMapping(FilesByUID);
+  const auto &SM = Context->getSourceManager();
+  unsigned MainFileUID = -1;
+  if (const auto *Entry = SM.getFileEntryForID(SM.getMainFileID()))
+    MainFileUID = Entry->getUID();
 
   if (FilesByUID.size() > HS.header_file_size())
     FilesByUID.resize(HS.header_file_size());
 
   for (unsigned UID = 0, LastUID = FilesByUID.size(); UID != LastUID; ++UID) {
     const FileEntry *File = FilesByUID[UID];
     if (!File)
       continue;
for (unsigned UID = 0, LastUID = FilesByUID.size(); UID != LastUID; ++UID) {
     HeaderFileInfoTrait::key_type Key = {
       Filename, File->getSize(), getTimestampForOutput(File)
     };
     HeaderFileInfoTrait::data_type Data = {
       *HFI, Included, HS.getModuleMap().findResolvedModulesForHeader(File), {}
     };
     Generator.insert(Key, Data, GeneratorTrait);
     ++NumHeaderSearchEntries;
+    // We may reuse a preamble even if the rest of the file is different, so
+    // allow looking up info for the main file with a zero size.
+    if (this->getASTContext().getLangOpts().CompilingPCH &&
+        File->getUID() == MainFileUID) {
+      Key.Size = 0;
+      Generator.insert(Key, Data, GeneratorTrait);
+      ++NumHeaderSearchEntries;
+    }
   }
 
   // Create the on-disk hash table in a buffer.
   SmallString<4096> TableData;
   uint32_t BucketOffset;
   {
     using namespace llvm::support;
 

llvm/include/llvm/Support/OnDiskHashTable.h

public:
   offset_type getNumEntries() const { return NumEntries; }
   const unsigned char *getBase() const { return Base; }
   const unsigned char *getBuckets() const { return Buckets; }
 
   bool isEmpty() const { return NumEntries == 0; }
 
   class iterator {
     internal_key_type Key;
-    const unsigned char *const Data;
-    const offset_type Len;
+    const unsigned char *Data;
+    offset_type Len;
     Info *InfoObj;
 
   public:
     iterator() : Key(), Data(nullptr), Len(0), InfoObj(nullptr) {}
     iterator(const internal_key_type K, const unsigned char *D, offset_type L,
              Info *InfoObj)
         : Key(K), Data(D), Len(L), InfoObj(InfoObj) {}