diff --git a/clang-tools-extra/clangd/Preamble.cpp b/clang-tools-extra/clangd/Preamble.cpp --- a/clang-tools-extra/clangd/Preamble.cpp +++ b/clang-tools-extra/clangd/Preamble.cpp @@ -671,10 +671,10 @@ // We are only interested in newly added includes, record the ones in // Baseline for exclusion. llvm::DenseMap, - /*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) @@ -685,7 +685,14 @@ // Include already present in the baseline preamble. Set resolved path and // put into preamble includes. if (It != ExistingIncludes.end()) { - Inc.Resolved = It->second.str(); + // Copy everything from existing include, apart from the location, when + // it's coming from baseline preamble. + if (It->second) { + Inc.Resolved = It->second->Resolved; + Inc.HeaderID = It->second->HeaderID; + Inc.BehindPragmaKeep = It->second->BehindPragmaKeep; + Inc.FileKind = It->second->FileKind; + } PP.PreambleIncludes.push_back(Inc); continue; } @@ -696,6 +703,10 @@ Patch << llvm::formatv( "#{0} {1}\n", spellingForIncDirective(Inc.Directive), Inc.Written); } + } else { + // Includes haven't changed or we're not patching. + // So, copy them over from the baseline. + PP.PreambleIncludes = Baseline.Includes.MainFileIncludes; } if (DirectivesChanged) { diff --git a/clang-tools-extra/clangd/unittests/PreambleTests.cpp b/clang-tools-extra/clangd/unittests/PreambleTests.cpp --- a/clang-tools-extra/clangd/unittests/PreambleTests.cpp +++ b/clang-tools-extra/clangd/unittests/PreambleTests.cpp @@ -15,6 +15,7 @@ #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" @@ -23,6 +24,7 @@ #include "llvm/Support/MemoryBuffer.h" #include "llvm/Support/VirtualFileSystem.h" #include "gmock/gmock.h" +#include "gtest/gtest-matchers.h" #include "gtest/gtest.h" #include #include @@ -166,7 +168,7 @@ 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\""; @@ -185,9 +187,14 @@ *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 createPatchedAST(llvm::StringRef Baseline, @@ -225,6 +232,20 @@ .str(); } +TEST(PreamblePatchTest, IncludesArePreserved) { + llvm::StringLiteral Baseline = "//error-ok\n#include \n#include \n"; + auto Modified = Baseline + "#define FOO\n"; + + 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 {