diff --git a/clang-tools-extra/include-cleaner/lib/Record.cpp b/clang-tools-extra/include-cleaner/lib/Record.cpp --- a/clang-tools-extra/include-cleaner/lib/Record.cpp +++ b/clang-tools-extra/include-cleaner/lib/Record.cpp @@ -187,9 +187,7 @@ FileID HashFID = SM.getFileID(HashLoc); int HashLine = SM.getLineNumber(HashFID, SM.getFileOffset(HashLoc)); checkForExport(HashFID, HashLine, File ? &File->getFileEntry() : nullptr); - - if (InMainFile && LastPragmaKeepInMainFileLine == HashLine) - Out->ShouldKeep.insert(HashLine); + checkForKeep(HashLine); } void checkForExport(FileID IncludingFile, int HashLine, @@ -213,6 +211,18 @@ ExportStack.pop_back(); } + void checkForKeep(int HashLine) { + if (!InMainFile || KeepStack.empty()) + return; + KeepPragma &Top = KeepStack.back(); + // Check if the current include is covered by a keep pragma. + if ((Top.Block && HashLine > Top.SeenAtLine) || Top.SeenAtLine == HashLine) + Out->ShouldKeep.insert(HashLine); + + if (!Top.Block) + KeepStack.pop_back(); // Pop immediately for single-line keep pragma. + } + bool HandleComment(Preprocessor &PP, SourceRange Range) override { auto &SM = PP.getSourceManager(); auto Pragma = @@ -257,23 +267,14 @@ } if (InMainFile) { - if (!Pragma->startswith("keep")) - return false; - // Given: - // - // #include "foo.h" - // #include "bar.h" // IWYU pragma: keep - // - // The order in which the callbacks will be triggered: - // - // 1. InclusionDirective("foo.h") - // 2. handleCommentInMainFile("// IWYU pragma: keep") - // 3. InclusionDirective("bar.h") - // - // This code stores the last location of "IWYU pragma: keep" comment in - // the main file, so that when next InclusionDirective is called, it will - // know that the next inclusion is behind the IWYU pragma. - LastPragmaKeepInMainFileLine = CommentLine; + if (Pragma->startswith("keep")) { + KeepStack.push_back({CommentLine, false}); + } else if (Pragma->starts_with("begin_keep")) { + KeepStack.push_back({CommentLine, true}); + } else if (Pragma->starts_with("end_keep") && !KeepStack.empty()) { + assert(KeepStack.back().Block); + KeepStack.pop_back(); + } } return false; } @@ -288,8 +289,7 @@ llvm::BumpPtrAllocator Arena; /// Intern table for strings. Contents are on the arena. llvm::StringSaver UniqueStrings; - // Track the last line "IWYU pragma: keep" was seen in the main file, 1-based. - int LastPragmaKeepInMainFileLine = -1; + struct ExportPragma { // The line number where we saw the begin_exports or export pragma. int SeenAtLine = 0; // 1-based line number. @@ -303,6 +303,16 @@ }; // A stack for tracking all open begin_exports or single-line export. std::vector ExportStack; + + struct KeepPragma { + // The line number where we saw the begin_keep or keep pragma. + int SeenAtLine = 0; // 1-based line number. + // true if it is a block begin/end_keep pragma; false if it is a + // single-line keep pragma. + bool Block = false; + }; + // A stack for tracking all open begin_keep pragmas or single-line keeps. + std::vector KeepStack; }; void PragmaIncludes::record(const CompilerInstance &CI) { diff --git a/clang-tools-extra/include-cleaner/unittests/RecordTest.cpp b/clang-tools-extra/include-cleaner/unittests/RecordTest.cpp --- a/clang-tools-extra/include-cleaner/unittests/RecordTest.cpp +++ b/clang-tools-extra/include-cleaner/unittests/RecordTest.cpp @@ -311,35 +311,68 @@ }; TEST_F(PragmaIncludeTest, IWYUKeep) { - Inputs.Code = R"cpp(// Line 1 - #include "keep1.h" // IWYU pragma: keep - #include "keep2.h" /* IWYU pragma: keep */ + llvm::Annotations MainFile(R"cpp( + $keep1^#include "keep1.h" // IWYU pragma: keep + $keep2^#include "keep2.h" /* IWYU pragma: keep */ + + $export1^#include "export1.h" // IWYU pragma: export + $begin_exports^// IWYU pragma: begin_exports + $export2^#include "export2.h" + $export3^#include "export3.h" + $end_exports^// IWYU pragma: end_exports + + $normal^#include "normal.h" + + $begin_keep^// IWYU pragma: begin_keep + $keep3^#include "keep3.h" + $end_keep^// IWYU pragma: end_keep + + // IWYU pragma: begin_keep + $keep4^#include "keep4.h" + // IWYU pragma: begin_keep + $keep5^#include "keep5.h" + // IWYU pragma: end_keep + $keep6^#include "keep6.h" + // IWYU pragma: end_keep + )cpp"); - #include "export1.h" // IWYU pragma: export // line 5 - // IWYU pragma: begin_exports - #include "export2.h" // Line 7 - #include "export3.h" - // IWYU pragma: end_exports + auto OffsetToLineNum = [&MainFile](size_t Offset) { + int Count = MainFile.code().substr(0, Offset).count('\n'); + return Count + 1; + }; - #include "normal.h" // Line 11 - )cpp"; - createEmptyFiles({"keep1.h", "keep2.h", "export1.h", "export2.h", "export3.h", + Inputs.Code = MainFile.code(); + createEmptyFiles({"keep1.h", "keep2.h", "keep3.h", "keep4.h", "keep5.h", + "keep6.h", "export1.h", "export2.h", "export3.h", "normal.h"}); TestAST Processed = build(); EXPECT_FALSE(PI.shouldKeep(1)); + // Keep - EXPECT_TRUE(PI.shouldKeep(2)); - EXPECT_TRUE(PI.shouldKeep(3)); + EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("keep1")))); + EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("keep2")))); - // Exports - EXPECT_TRUE(PI.shouldKeep(5)); - EXPECT_TRUE(PI.shouldKeep(7)); - EXPECT_TRUE(PI.shouldKeep(8)); - EXPECT_FALSE(PI.shouldKeep(6)); // no # directive - EXPECT_FALSE(PI.shouldKeep(9)); // no # directive + EXPECT_FALSE(PI.shouldKeep( + OffsetToLineNum(MainFile.point("begin_keep")))); // no # directive + EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("keep3")))); + EXPECT_FALSE(PI.shouldKeep( + OffsetToLineNum(MainFile.point("end_keep")))); // no # directive + + EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("keep4")))); + EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("keep5")))); + EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("keep6")))); - EXPECT_FALSE(PI.shouldKeep(11)); + // Exports + EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("export1")))); + EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("export2")))); + EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("export3")))); + EXPECT_FALSE(PI.shouldKeep( + OffsetToLineNum(MainFile.point("begin_exports")))); // no # directive + EXPECT_FALSE(PI.shouldKeep( + OffsetToLineNum(MainFile.point("end_exports")))); // no # directive + + EXPECT_FALSE(PI.shouldKeep(OffsetToLineNum(MainFile.point("normal")))); } TEST_F(PragmaIncludeTest, IWYUPrivate) {