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 @@ -1272,7 +1272,7 @@ // Force them to be deserialized so SemaCodeComplete sees them. loadMainFilePreambleMacros(Clang->getPreprocessor(), Input.Preamble); if (Includes) - Clang->getPreprocessor().addPPCallbacks(Includes->collect(*Clang)); + Includes->collect(*Clang); 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,10 @@ RealPathNames.emplace_back(); } - // Returns a PPCallback that visits all inclusions in the main file and - // populates the structure. - std::unique_ptr collect(const CompilerInstance &CI); + // Inserts a PPCallback and CommentHandler that visits all includes in the + // main file and populates the structure. It will also scan for IWYU pragmas + // in comments. + void 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,22 @@ #include "clang/Frontend/CompilerInvocation.h" #include "clang/Frontend/FrontendActions.h" #include "clang/Lex/HeaderSearch.h" +#include "clang/Lex/PPCallbacks.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 PPCallbacks, + public CommentHandler { 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 +63,8 @@ Inc.Directive = IncludeTok.getIdentifierInfo()->getPPKeywordID(); if (File) Inc.HeaderID = static_cast(Out->getOrCreateID(File)); + if (LastPragmaKeepInMainFileLine == Inc.HashLine) + Inc.BehindPragmaKeep = true; } // Record include graph (not just for main-file includes) @@ -80,12 +87,14 @@ FileID PrevFID) override { switch (Reason) { case PPCallbacks::EnterFile: + ++Level; if (BuiltinFile.isInvalid() && SM.isWrittenInBuiltinFile(Loc)) { BuiltinFile = SM.getFileID(Loc); InBuiltinFile = true; } break; case PPCallbacks::ExitFile: { + --Level; if (PrevFID == BuiltinFile) InBuiltinFile = false; // At file exit time HeaderSearchInfo is valid and can be used to @@ -102,7 +111,38 @@ } } + // 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 { + if (!inMainFile() || Range.getBegin().isMacroID()) + return false; + bool Err = false; + llvm::StringRef Text = SM.getCharacterData(Range.getBegin(), &Err); + if (Err || !Text.consume_front(IWYUPragmaKeep)) + return false; + unsigned Offset = SM.getFileOffset(Range.getBegin()); + LastPragmaKeepInMainFileLine = + SM.getLineNumber(SM.getFileID(Range.getBegin()), Offset) - 1; + return false; + } + private: + // Keeps track of include depth for the current file. It's 1 for main file. + int Level = 0; + bool inMainFile() const { return Level == 1; } + const SourceManager &SM; HeaderSearch &HeaderInfo; // Set after entering the file. @@ -111,6 +151,9 @@ bool InBuiltinFile = false; IncludeStructure *Out; + + // The last line "IWYU pragma: keep" was seen in the main file, 0-indexed. + int LastPragmaKeepInMainFileLine = -1; }; bool isLiteralInclude(llvm::StringRef Include) { @@ -157,12 +200,12 @@ return Headers; } -std::unique_ptr -IncludeStructure::collect(const CompilerInstance &CI) { +void IncludeStructure::collect(const CompilerInstance &CI) { auto &SM = CI.getSourceManager(); MainFileEntry = SM.getFileEntryForID(SM.getMainFileID()); - return std::make_unique( - SM, CI.getPreprocessor().getHeaderSearchInfo(), this); + auto Collector = std::make_unique(CI, this); + CI.getPreprocessor().addCommentHandler(Collector.get()); + CI.getPreprocessor().addPPCallbacks(std::move(Collector)); } 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,7 @@ // 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)); + Includes.collect(*Clang); // 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 @@ -99,7 +99,7 @@ CanonIncludes.addSystemHeadersMapping(CI.getLangOpts()); LangOpts = &CI.getLangOpts(); SourceMgr = &CI.getSourceManager(); - Compiler = &CI; + Includes.collect(CI); } std::unique_ptr createPPCallbacks() override { @@ -107,10 +107,8 @@ "SourceMgr and LangOpts must be set at this point"); return std::make_unique( - Includes.collect(*Compiler), - std::make_unique( - std::make_unique(*SourceMgr, Macros), - collectPragmaMarksCallback(*SourceMgr, Marks))); + std::make_unique(*SourceMgr, Macros), + collectPragmaMarksCallback(*SourceMgr, Marks)); } CommentHandler *getCommentHandler() override { @@ -142,7 +140,6 @@ std::unique_ptr IWYUHandler = nullptr; const clang::LangOptions *LangOpts = nullptr; const SourceManager *SourceMgr = nullptr; - const CompilerInstance *Compiler = nullptr; }; // Represents directives other than includes, where basic textual information is @@ -288,7 +285,7 @@ return error("failed BeginSourceFile"); Preprocessor &PP = Clang->getPreprocessor(); IncludeStructure Includes; - PP.addPPCallbacks(Includes.collect(*Clang)); + Includes.collect(*Clang); 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,7 @@ EXPECT_TRUE( Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0])); IncludeStructure Includes; - Clang->getPreprocessor().addPPCallbacks(Includes.collect(*Clang)); + Includes.collect(*Clang); EXPECT_FALSE(Action.Execute()); Action.EndSourceFile(); return Includes; @@ -142,6 +142,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 +258,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] = ""; 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 @@ -82,7 +82,7 @@ return {}; } IncludeStructure Includes; - Clang->getPreprocessor().addPPCallbacks(Includes.collect(*Clang)); + Includes.collect(*Clang); if (llvm::Error Err = Action.Execute()) { ADD_FAILURE() << "failed to execute action: " << std::move(Err); return {};