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

[VitaNuo]

Summary

Implement support for begin_keep/end_keep pragmas.

[VitaNuo]

Add a comment.

[hokein]

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

219

nit: remove the brace for the if, the same below.

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

354

with more tests being added, these line-number-inlined tests now are hard to follow, even with the line comment in the example.

I think we can polish these tests by using the annotation point ^, something like.

llvm::Annotation Code = R"cpp(
$keep1^#include "keep1.h" // IWYU pragma: keep

$normal^%include "normal.h"
)cpp";

EXPECT_TRUE(PI.shouldKeep(offsetToLine(Code.point("keep1"))));
EXPECT_FALSE(PI.shouldKeep(offsetToLine(Code.point("normal"))));

We need to implement an offsetToLine function, it should be trivial to add (basically calculate the number of \ns to the offset point).

[VitaNuo]

Address review comments. Use annotation points instead of line numbers.

[VitaNuo]

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

354

Ok, thanks. Have a look now please.

[VitaNuo]

Fix formatting.

[hokein]

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

312

this is only used in IWYUKeep test, I'd make it as a local lambda (with MainFile captured) in the test:

auto OffsetToLineNum = [&MainFile](size_t Offset) { ... };

349

I think we can also do similar thing on $begin_keep_pragma^// IWYU pragma...

351

keep4.h can be removed to simplify the testcode.

358

keep7.h can be removed.

365

no need to repeatly test the begin_keep and end_keep pragmas.

[VitaNuo]

Address review comments.

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

312

Nice idea, thanks.

[hokein]

thanks, looks good.

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

312

this method can be removed.

[VitaNuo]

Remove excessive method.

[VitaNuo]

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

312

Oops, sorry.

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

void InclusionDirective(SourceLocation HashLoc, const Token &IncludeTok,
                           Optional<FileEntryRef> File,
                           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);
-
-    if (InMainFile && LastPragmaKeepInMainFileLine == HashLine)
-      Out->ShouldKeep.insert(HashLine);
+    checkForKeep(HashLine);
   }
 
   void checkForExport(FileID IncludingFile, int HashLine,
                       const FileEntry *IncludedHeader) {
     if (ExportStack.empty())
       return;
     auto &Top = ExportStack.back();
     if (Top.SeenAtFile != IncludingFile)
       return;
     // Make sure current include is covered by the export pragma.
     if ((Top.Block && HashLine > Top.SeenAtLine) ||
         Top.SeenAtLine == HashLine) {
       if (IncludedHeader)
         Out->IWYUExportBy[IncludedHeader->getUniqueID()].push_back(
             Top.FullPath);
       // main-file #include with export pragma should never be removed.
       if (Top.SeenAtFile == SM.getMainFileID())
         Out->ShouldKeep.insert(HashLine);
     }
     if (!Top.Block) // Pop immediately for single-line export pragma.
       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 =
         tooling::parseIWYUPragma(SM.getCharacterData(Range.getBegin()));
     if (!Pragma)
       return false;
 
     if (Pragma->consume_front("private")) {
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;
-      // 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;
   }
 
 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.
-  int LastPragmaKeepInMainFileLine = -1;
+
   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
     // single-line export pragma.
     bool Block = false;
   };
   // A stack for tracking all open begin_exports or single-line export.
   std::vector<ExportPragma> 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<KeepPragma> KeepStack;
 };
 
 void PragmaIncludes::record(const CompilerInstance &CI) {
   auto Record = std::make_unique<RecordPragma>(CI, this);
   CI.getPreprocessor().addCommentHandler(Record.get());
   CI.getPreprocessor().addPPCallbacks(std::move(Record));
 }
 

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

protected:
 
   TestAST build() { return TestAST(Inputs); }
 
   void createEmptyFiles(llvm::ArrayRef<StringRef> FileNames) {
     for (llvm::StringRef File : FileNames)
       Inputs.ExtraFiles[File] = "";
   }
 };
 
 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 */
 
-    #include "export1.h" // IWYU pragma: export // line 5
-    // IWYU pragma: begin_exports
-    #include "export2.h" // Line 7
-    #include "export3.h"
-    // IWYU pragma: end_exports
+    $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 "normal.h" // Line 11
-  )cpp";
-  createEmptyFiles({"keep1.h", "keep2.h", "export1.h", "export2.h", "export3.h",
+  auto OffsetToLineNum = [&MainFile](size_t Offset) {
+    int Count = MainFile.code().substr(0, Offset).count('\n');
+    return Count + 1;
+  };
+
+  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"))));
+
+  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"))));
 
   // 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_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(11));
+  EXPECT_FALSE(PI.shouldKeep(OffsetToLineNum(MainFile.point("normal"))));
 }
 
 TEST_F(PragmaIncludeTest, IWYUPrivate) {
   Inputs.Code = R"cpp(
     #include "public.h"
   )cpp";
   Inputs.ExtraFiles["public.h"] = R"cpp(
     #include "private.h"