[include-cleaner] Implement IWYU begin_keep/end_keep pragma support.

[VitaNuo]

Summary

Keep the includes contained between IWYU begin_keep and end_keep pragmas.

[VitaNuo]

Remove unnecessary code.

[hokein]

clang-tools-extra/include-cleaner/lib/Record.cpp

299

using a simple variable is not enough to handle the nested case like

// IWYU pragma: begin_keep
#include "keep1.h" 
// IWYU pragma: begin_keep
#include "keep2.h" 
#include "keep3.h"
// IWYU pragma: end_keep
#include "keep4.h"
// IWYU pragma: end_keep

now keep4.h is marked as unkeeped. I think we need to maintain a stack for Keep here (we already have a similar pattern for export)

[VitaNuo]

clang-tools-extra/include-cleaner/lib/Record.cpp

299

Ok, I had no idea this was valid usage :(

clang-tools-extra/include-cleaner/lib/Record.cpp

void InclusionDirective(SourceLocation HashLoc, const Token &IncludeTok,
                           llvm::StringRef /*SearchPath*/,
                           llvm::StringRef /*RelativePath*/,
                           const clang::Module * /*Imported*/,
                           SrcMgr::CharacteristicKind FileKind) override {
     FileID HashFID = SM.getFileID(HashLoc);
     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,
                       const FileEntry *IncludedHeader) {
     if (ExportStack.empty())
       return;
     auto &Top = ExportStack.back();
     if (Top.SeenAtFile != IncludingFile)
if (Pragma->startswith("export")) {
       // the begin_exports and end_exports is in the same file.
       if (!ExportStack.empty()) {
         assert(ExportStack.back().Block);
         ExportStack.pop_back();
       }
     }
 
     if (InMainFile) {
-      if (!Pragma->startswith("keep"))
-        return false;
+      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.
+        // 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;
   }
 
 private:
   StringRef save(llvm::StringRef S) { return UniqueStrings.save(S); }
 
   bool InMainFile = false;
   const SourceManager &SM;
   HeaderSearch &HeaderInfo;
   PragmaIncludes *Out;
   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.
     // The file where we saw the pragma.
     FileID SeenAtFile;
     // FullPath of the file SeenAtFile.
     StringRef FullPath;
     // true if it is a block begin/end_exports pragma; false if it is a

clang-tools-extra/include-cleaner/unittests/RecordTest.cpp

Inputs.Code = R"cpp(// Line 1
 
     #include "export1.h" // IWYU pragma: export // line 5
     // IWYU pragma: begin_exports
     #include "export2.h" // Line 7
     #include "export3.h"
     // 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));
   EXPECT_TRUE(PI.shouldKeep(7));
   EXPECT_TRUE(PI.shouldKeep(8));
   EXPECT_FALSE(PI.shouldKeep(6)); // no # directive
   EXPECT_FALSE(PI.shouldKeep(9)); // no # directive