[clangd] Fix bugs in main-file include patching for stale preambles

[kadircet]

Summary

• Make sure main file includes are present even when they're not patched (because they didn't change or we're explicitly not patching them).
• Populate extra fields for includes, which can be used by include-cleaner.

[kadircet]

• Also populate additional fields

[sammccall]

Code looks reasonable, but I don't understand why the changes are being made - can you explain/link to a bug in the commit message?

clang-tools-extra/clangd/Preamble.cpp

687–688

This doesn't seem related to the patch description, can you update it?

689

why the null check?

690

seems a bit more future-proof to copy the whole thing and just overwrite the fields you want to preserve. (also would avoid the need for the comment)

706

this just echoes the code - say why instead?

(would also be good to have the reason in the commit message)

[kadircet]

• Address comments & update commit message

clang-tools-extra/clangd/Preamble.cpp

689

as discussed offline, we might have includes coming from scanned contents whose are not part of the preamble, e.g. because they're in a disabled PP region.

[sammccall]

clang-tools-extra/clangd/Preamble.cpp

691

if It->second is null, then all the #includes of this header from the baseline preamble were in disabled sections, so it's *very* likely this one is too.

I think we're better not pushing onto PP.PreambleIncludes at all in this case, rather than pushing an unresolved one - this is how the MainFileIncludes looks like when an #include is in a disabled section and there's no patching happening.

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

236

rawstrings?

[kadircet]

• use rawstrings in test

[kadircet]

clang-tools-extra/clangd/Preamble.cpp

691

agreed. as discussed sent out D143597

clang-tools-extra/clangd/Preamble.cpp

PreamblePatch PreamblePatch::create(llvm::StringRef FileName,
   // might lead to problems on windows especially.
   escapeBackslashAndQuotes(FileName, Patch);
   Patch << "\"\n";
 
   if (IncludesChanged && PatchType == PatchType::All) {
     // We are only interested in newly added includes, record the ones in
     // Baseline for exclusion.
     llvm::DenseMap<std::pair<tok::PPKeywordKind, llvm::StringRef>,
-                   /*Resolved=*/llvm::StringRef>
+                   const Inclusion *>
         ExistingIncludes;
     for (const auto &Inc : Baseline.Includes.MainFileIncludes)
-      ExistingIncludes[{Inc.Directive, Inc.Written}] = Inc.Resolved;
+      ExistingIncludes[{Inc.Directive, Inc.Written}] = &Inc;
     // There might be includes coming from disabled regions, record these for
     // exclusion too. note that we don't have resolved paths for those.
     for (const auto &Inc : BaselineScan->Includes)
       ExistingIncludes.try_emplace({Inc.Directive, Inc.Written});
     // Calculate extra includes that needs to be inserted.
     for (auto &Inc : ModifiedScan->Includes) {
       auto It = ExistingIncludes.find({Inc.Directive, Inc.Written});
       // Include already present in the baseline preamble. Set resolved path and
       // put into preamble includes.
       if (It != ExistingIncludes.end()) {
-        Inc.Resolved = It->second.str();
-        PP.PreambleIncludes.push_back(Inc);
+        auto &PatchedInc = PP.PreambleIncludes.emplace_back();
+        // Copy everything from existing include, apart from the location, when
+        // it's coming from baseline preamble.
+        if (It->second)
+          PatchedInc = *It->second;
+        PatchedInc.HashLine = Inc.HashLine;
+        PatchedInc.HashOffset = Inc.HashOffset;
         continue;
       }
       // Include is new in the modified preamble. Inject it into the patch and
       // use #line to set the presumed location to where it is spelled.
       auto LineCol = offsetToClangLineColumn(Modified.Contents, Inc.HashOffset);
       Patch << llvm::formatv("#line {0}\n", LineCol.first);
       Patch << llvm::formatv(
           "#{0} {1}\n", spellingForIncDirective(Inc.Directive), Inc.Written);
     }
+  } else {
+    // Make sure we have the full set of includes available even when we're not
+    // patching. As these are used by features we provide afterwards like hover,
+    // go-to-def or include-cleaner when preamble is stale.
+    PP.PreambleIncludes = Baseline.Includes.MainFileIncludes;
   }
 
   if (DirectivesChanged) {
     // We need to patch all the directives, since they are order dependent. e.g:
     // #define BAR(X) NEW(X) // Newly introduced in Modified
     // #define BAR(X) OLD(X) // Exists in the Baseline
     //
     // If we've patched only the first directive, the macro definition would've

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

 #include "Compiler.h"
 #include "Headers.h"
 #include "Hover.h"
 #include "Preamble.h"
 #include "SourceCode.h"
 #include "TestFS.h"
 #include "TestTU.h"
 #include "XRefs.h"
+#include "clang/Basic/SourceManager.h"
 #include "clang/Format/Format.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include "gmock/gmock.h"
+#include "gtest/gtest-matchers.h"
 #include "gtest/gtest.h"
 #include <memory>
 #include <optional>
 #include <string>
 #include <vector>
 
 using testing::Contains;
 using testing::Field;
TEST(PreamblePatchTest, MainFileIsEscaped) {
   EXPECT_THAT(Includes, ElementsAre(AllOf(Field(&Inclusion::Written, "<a.h>"),
                                           Field(&Inclusion::HashLine, 0))));
 }
 
 TEST(PreamblePatchTest, PatchesPreambleIncludes) {
   MockFS FS;
   IgnoreDiagnostics Diags;
   auto TU = TestTU::withCode(R"cpp(
-    #include "a.h"
+    #include "a.h" // IWYU pragma: keep
     #include "c.h"
   )cpp");
   TU.AdditionalFiles["a.h"] = "#include \"b.h\"";
   TU.AdditionalFiles["b.h"] = "";
   TU.AdditionalFiles["c.h"] = "";
   auto PI = TU.inputs(FS);
   auto BaselinePreamble = buildPreamble(
       TU.Filename, *buildCompilerInvocation(PI, Diags), PI, true, nullptr);
   // We drop c.h from modified and add a new header. Since the latter is patched
   // we should only get a.h in preamble includes.
   TU.Code = R"cpp(
     #include "a.h"
     #include "b.h"
   )cpp";
   auto PP = PreamblePatch::createFullPatch(testPath(TU.Filename), TU.inputs(FS),
                                            *BaselinePreamble);
   // Only a.h should exists in the preamble, as c.h has been dropped and b.h was
   // newly introduced.
-  EXPECT_THAT(PP.preambleIncludes(),
-              ElementsAre(AllOf(Field(&Inclusion::Written, "\"a.h\""),
-                                Field(&Inclusion::Resolved, testPath("a.h")))));
+  EXPECT_THAT(
+      PP.preambleIncludes(),
+      ElementsAre(AllOf(
+          Field(&Inclusion::Written, "\"a.h\""),
+          Field(&Inclusion::Resolved, testPath("a.h")),
+          Field(&Inclusion::HeaderID, testing::Not(testing::Eq(std::nullopt))),
+          Field(&Inclusion::BehindPragmaKeep, true),
+          Field(&Inclusion::FileKind, SrcMgr::CharacteristicKind::C_User))));
 }
 
 std::optional<ParsedAST> createPatchedAST(llvm::StringRef Baseline,
                                           llvm::StringRef Modified) {
   auto BaselinePreamble = TestTU::withCode(Baseline).preamble();
   if (!BaselinePreamble) {
     ADD_FAILURE() << "Failed to build baseline preamble";
     return std::nullopt;
std::string getPreamblePatch(llvm::StringRef Baseline,
   MockFS FS;
   auto TU = TestTU::withCode(Modified);
   return PreamblePatch::createFullPatch(testPath("main.cpp"), TU.inputs(FS),
                                         *BaselinePreamble)
       .text()
       .str();
 }
 
+TEST(PreamblePatchTest, IncludesArePreserved) {
+  llvm::StringLiteral Baseline = R"(//error-ok
+#include <foo>
+#include <bar>
+)";
+  llvm::StringLiteral Modified = R"(//error-ok
+#include <foo>
+#include <bar>
+#define FOO)";
+
+  auto Includes = createPatchedAST(Baseline, Modified.str())
+                      ->getIncludeStructure()
+                      .MainFileIncludes;
+  EXPECT_TRUE(!Includes.empty());
+  EXPECT_EQ(Includes, TestTU::withCode(Baseline)
+                          .build()
+                          .getIncludeStructure()
+                          .MainFileIncludes);
+}
+
 TEST(PreamblePatchTest, Define) {
   // BAR should be defined while parsing the AST.
   struct {
     const char *const Contents;
     const char *const ExpectedPatch;
   } Cases[] = {
       {
           R"cpp(