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 @@ -202,8 +202,14 @@ int HashLine = SM.getLineNumber(HashFID, SM.getFileOffset(HashLoc)); checkForExport(HashFID, HashLine, File ? &File->getFileEntry() : nullptr); + // one-line "IWYU pragma: keep" pragma handling. if (InMainFile && LastPragmaKeepInMainFileLine == HashLine) Out->ShouldKeep.insert(HashLine); + + // multiline "IWYU pragma: begin_keep"/"IWYU pragma: end_keep" handling. + if (InMainFile && InsidePragmaKeepBlock) { + Out->ShouldKeep.insert(HashLine); + } } void checkForExport(FileID IncludingFile, int HashLine, @@ -270,23 +276,28 @@ } 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")) + // 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->starts_with("begin_keep")) { + InsidePragmaKeepBlock = true; + } else if (Pragma->starts_with("end_keep")) { + InsidePragmaKeepBlock = false; + } } return false; } @@ -301,8 +312,12 @@ 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. + // Track the last line "IWYU pragma: keep" or "IWYU pragma: begin_keep" + // was seen in the main file, 1-based. int LastPragmaKeepInMainFileLine = -1; + // true if the location is inside a keep block (i.e., between "IWYU pragma: + // begin_keep" and "IWYU pragma: end_keep"). + bool InsidePragmaKeepBlock = false; struct ExportPragma { // The line number where we saw the begin_exports or export pragma. int SeenAtLine = 0; // 1-based line number. 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,36 @@ // IWYU pragma: end_exports #include "normal.h" // Line 11 + + // IWYU pragma: begin_keep + #include "keep3.h" + #include "keep4.h" + // IWYU pragma: end_keep + + // IWYU pragma: begin_keep // Line 18 + // IWYU pragma: end_keep + + // IWYU pragma: begin_keep // Line 21 + #include "keep5.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", + "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)); // line with "begin_keep" pragma + EXPECT_TRUE(PI.shouldKeep(14)); + EXPECT_TRUE(PI.shouldKeep(15)); + EXPECT_FALSE(PI.shouldKeep(16)); // line with "end_keep" pragma + EXPECT_FALSE(PI.shouldKeep(18)); // starts empty "begin_keep" block + EXPECT_FALSE(PI.shouldKeep(19)); // ends empty "begin_keep" block + EXPECT_FALSE(PI.shouldKeep(21)); // line with "begin_keep" pragma + EXPECT_TRUE(PI.shouldKeep(22)); + EXPECT_FALSE(PI.shouldKeep(23)); // line with "end_keep" pragma // Exports EXPECT_TRUE(PI.shouldKeep(5));