diff --git a/clang-tools-extra/clangd/CodeComplete.cpp b/clang-tools-extra/clangd/CodeComplete.cpp --- a/clang-tools-extra/clangd/CodeComplete.cpp +++ b/clang-tools-extra/clangd/CodeComplete.cpp @@ -1271,8 +1271,11 @@ // deserialized only when mentioned. // Force them to be deserialized so SemaCodeComplete sees them. loadMainFilePreambleMacros(Clang->getPreprocessor(), Input.Preamble); - if (Includes) - Clang->getPreprocessor().addPPCallbacks(Includes->collect(*Clang)); + if (Includes) { + auto Collector = Includes->collect(*Clang); + Clang->getPreprocessor().addCommentHandler(Collector.get()); + Clang->getPreprocessor().addPPCallbacks(move(Collector)); + } if (llvm::Error Err = Action.Execute()) { log("Execute() failed when running codeComplete for {0}: {1}", Input.FileName, toString(std::move(Err))); diff --git a/clang-tools-extra/clangd/Headers.h b/clang-tools-extra/clangd/Headers.h --- a/clang-tools-extra/clangd/Headers.h +++ b/clang-tools-extra/clangd/Headers.h @@ -20,6 +20,7 @@ #include "clang/Frontend/CompilerInstance.h" #include "clang/Lex/HeaderSearch.h" #include "clang/Lex/PPCallbacks.h" +#include "clang/Lex/Preprocessor.h" #include "clang/Tooling/Inclusions/HeaderIncludes.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DenseSet.h" @@ -64,6 +65,7 @@ int HashLine = 0; // Line number containing the directive, 0-indexed. SrcMgr::CharacteristicKind FileKind = SrcMgr::C_User; llvm::Optional HeaderID; + bool BehindPragmaKeep = false; // Has IWYU pragma: keep right after. }; llvm::raw_ostream &operator<<(llvm::raw_ostream &, const Inclusion &); bool operator==(const Inclusion &LHS, const Inclusion &RHS); @@ -123,9 +125,21 @@ RealPathNames.emplace_back(); } + class IncludeCollector : public PPCallbacks, public CommentHandler {}; // Returns a PPCallback that visits all inclusions in the main file and - // populates the structure. - std::unique_ptr collect(const CompilerInstance &CI); + // populates the structure. IncludeCollector can also scan the comments for + // IWYU pragmas but it needs to be explicitly added to as a preprocessor + // comment handler. PP wants to own the PPCallbacks, so the typical way to + // do both is: + // + // auto Collector = collect(SM); + // PP.addCommentHandler(Colletor.get()); + // PP.addPPCallbacks(move(Collector)); + // + // When IncludeCollector is not added as the CommentHandler, resulting + // include structure won't have the IWYU pragmas information. + std::unique_ptr + collect(const CompilerInstance &CI); // HeaderID identifies file in the include graph. It corresponds to a // FileEntry rather than a FileID, but stays stable across preamble & main diff --git a/clang-tools-extra/clangd/Headers.cpp b/clang-tools-extra/clangd/Headers.cpp --- a/clang-tools-extra/clangd/Headers.cpp +++ b/clang-tools-extra/clangd/Headers.cpp @@ -17,17 +17,21 @@ #include "clang/Frontend/CompilerInvocation.h" #include "clang/Frontend/FrontendActions.h" #include "clang/Lex/HeaderSearch.h" +#include "clang/Lex/Preprocessor.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/Path.h" namespace clang { namespace clangd { -class IncludeStructure::RecordHeaders : public PPCallbacks { +const char IWYUPragmaKeep[] = "// IWYU pragma: keep"; + +class IncludeStructure::RecordHeaders + : public IncludeStructure::IncludeCollector { public: - RecordHeaders(const SourceManager &SM, HeaderSearch &HeaderInfo, - IncludeStructure *Out) - : SM(SM), HeaderInfo(HeaderInfo), Out(Out) {} + RecordHeaders(const CompilerInstance &CI, IncludeStructure *Out) + : SM(CI.getSourceManager()), + HeaderInfo(CI.getPreprocessor().getHeaderSearchInfo()), Out(Out) {} // Record existing #includes - both written and resolved paths. Only #includes // in the main file are collected. @@ -58,6 +62,8 @@ Inc.Directive = IncludeTok.getIdentifierInfo()->getPPKeywordID(); if (File) Inc.HeaderID = static_cast(Out->getOrCreateID(File)); + if (PragmaKeepInMainFile == Inc.HashLine) + Inc.BehindPragmaKeep = true; } // Record include graph (not just for main-file includes) @@ -102,6 +108,34 @@ } } + // Given: + // + // #include "foo.h" + // #include "bar.h" // IWYU pragma: keep + // + // The order in which the callbacks will be triggered: + // + // 1. InclusionDirective("foo.h") + // 2. HandleComment("// IWYU pragma: keep") + // 3. InclusionDirective("bar.h") + // + // HandleComment will store the last location of "IWYU pragma: keep" comment + // in the main file, so that when InclusionDirective is called, it will know + // that the next inclusion is behind the IWYU pragma. + bool HandleComment(Preprocessor &PP, SourceRange Range) override { + llvm::StringRef Text = + Lexer::getSourceText(CharSourceRange::getCharRange(Range), + PP.getSourceManager(), PP.getLangOpts()); + if (!Text.consume_front(IWYUPragmaKeep)) + return false; + if (!isInsideMainFile(Range.getBegin(), SM)) + return false; + unsigned Offset = SM.getFileOffset(Range.getBegin()); + PragmaKeepInMainFile = + SM.getLineNumber(SM.getFileID(Range.getBegin()), Offset) - 1; + return false; + } + private: const SourceManager &SM; HeaderSearch &HeaderInfo; @@ -111,6 +145,9 @@ bool InBuiltinFile = false; IncludeStructure *Out; + + // The last line "IWYU pragma: keep" was seen in the main file. + int PragmaKeepInMainFile = -1; }; bool isLiteralInclude(llvm::StringRef Include) { @@ -157,12 +194,9 @@ return Headers; } -std::unique_ptr +std::unique_ptr IncludeStructure::collect(const CompilerInstance &CI) { - auto &SM = CI.getSourceManager(); - MainFileEntry = SM.getFileEntryForID(SM.getMainFileID()); - return std::make_unique( - SM, CI.getPreprocessor().getHeaderSearchInfo(), this); + return std::make_unique(CI, this); } llvm::Optional diff --git a/clang-tools-extra/clangd/ParsedAST.cpp b/clang-tools-extra/clangd/ParsedAST.cpp --- a/clang-tools-extra/clangd/ParsedAST.cpp +++ b/clang-tools-extra/clangd/ParsedAST.cpp @@ -446,7 +446,9 @@ // Important: collectIncludeStructure is registered *after* ReplayPreamble! // Otherwise we would collect the replayed includes again... // (We can't *just* use the replayed includes, they don't have Resolved path). - Clang->getPreprocessor().addPPCallbacks(Includes.collect(*Clang)); + auto Collector = Includes.collect(*Clang); + Clang->getPreprocessor().addCommentHandler(Collector.get()); + Clang->getPreprocessor().addPPCallbacks(move(Collector)); // Copy over the macros in the preamble region of the main file, and combine // with non-preamble macros below. MainFileMacros Macros; 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 @@ -288,7 +288,9 @@ return error("failed BeginSourceFile"); Preprocessor &PP = Clang->getPreprocessor(); IncludeStructure Includes; - PP.addPPCallbacks(Includes.collect(*Clang)); + auto Collector = Includes.collect(*Clang); + PP.addCommentHandler(Collector.get()); + PP.addPPCallbacks(move(Collector)); ScannedPreamble SP; SP.Bounds = Bounds; PP.addPPCallbacks( diff --git a/clang-tools-extra/clangd/unittests/HeadersTests.cpp b/clang-tools-extra/clangd/unittests/HeadersTests.cpp --- a/clang-tools-extra/clangd/unittests/HeadersTests.cpp +++ b/clang-tools-extra/clangd/unittests/HeadersTests.cpp @@ -80,7 +80,9 @@ EXPECT_TRUE( Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0])); IncludeStructure Includes; - Clang->getPreprocessor().addPPCallbacks(Includes.collect(*Clang)); + auto Collector = Includes.collect(*Clang); + Clang->getPreprocessor().addCommentHandler(Collector.get()); + Clang->getPreprocessor().addPPCallbacks(move(Collector)); EXPECT_FALSE(Action.Execute()); Action.EndSourceFile(); return Includes; @@ -142,6 +144,7 @@ MATCHER_P(Resolved, Name, "") { return arg.Resolved == Name; } MATCHER_P(IncludeLine, N, "") { return arg.HashLine == N; } MATCHER_P(Directive, D, "") { return arg.Directive == D; } +MATCHER_P(HasPragmaKeep, H, "") { return arg.BehindPragmaKeep == H; } MATCHER_P2(Distance, File, D, "") { if (arg.getFirst() != File) @@ -257,6 +260,18 @@ Directive(tok::pp_include_next))); } +TEST_F(HeadersTest, IWYUPragmaKeep) { + FS.Files[MainFile] = R"cpp( +#include "bar.h" // IWYU pragma: keep +#include "foo.h" +)cpp"; + + EXPECT_THAT( + collectIncludes().MainFileIncludes, + UnorderedElementsAre(AllOf(Written("\"foo.h\""), HasPragmaKeep(false)), + AllOf(Written("\"bar.h\""), HasPragmaKeep(true)))); +} + TEST_F(HeadersTest, InsertInclude) { std::string Path = testPath("sub/bar.h"); FS.Files[Path] = "";