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 @@ -186,9 +186,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, @@ -212,6 +210,19 @@ ExportStack.pop_back(); } + void checkForKeep(int HashLine) { + if (!InMainFile || KeepStack.empty()) + return; + KeepPragma &Top = KeepStack.back(); + 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 = @@ -256,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; } @@ -287,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. @@ -302,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 @@ -322,15 +322,42 @@ // IWYU pragma: end_exports #include "normal.h" // Line 11 + + // IWYU pragma: begin_keep // Line 13 + #include "keep3.h" + #include "keep4.h" + // IWYU pragma: end_keep + + // IWYU pragma: begin_keep // Line 18 + #include "keep5.h" + // IWYU pragma: begin_keep + #include "keep6.h" + #include "keep7.h" + // IWYU pragma: end_keep + #include "keep8.h" + // IWYU pragma: end_keep )cpp"; - createEmptyFiles({"keep1.h", "keep2.h", "export1.h", "export2.h", "export3.h", - "normal.h"}); + createEmptyFiles({"keep1.h", "keep2.h", "keep3.h", "keep4.h", "keep5.h", + "keep6.h", "keep7.h", "keep8.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_FALSE(PI.shouldKeep(13)); + EXPECT_TRUE(PI.shouldKeep(14)); + EXPECT_TRUE(PI.shouldKeep(15)); + EXPECT_FALSE(PI.shouldKeep(16)); + EXPECT_FALSE(PI.shouldKeep(18)); + EXPECT_TRUE(PI.shouldKeep(19)); + EXPECT_FALSE(PI.shouldKeep(20)); + EXPECT_TRUE(PI.shouldKeep(21)); + EXPECT_TRUE(PI.shouldKeep(22)); + EXPECT_FALSE(PI.shouldKeep(23)); + EXPECT_TRUE(PI.shouldKeep(24)); + EXPECT_FALSE(PI.shouldKeep(25)); // Exports EXPECT_TRUE(PI.shouldKeep(5));