VitaNuo updated this revision to Diff 478203.
VitaNuo added a comment.

Add a comment.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D138797/new/

https://reviews.llvm.org/D138797

Files:
  clang-tools-extra/include-cleaner/lib/Record.cpp
  clang-tools-extra/include-cleaner/unittests/RecordTest.cpp

Index: clang-tools-extra/include-cleaner/unittests/RecordTest.cpp
===================================================================
--- clang-tools-extra/include-cleaner/unittests/RecordTest.cpp
+++ 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));
Index: clang-tools-extra/include-cleaner/lib/Record.cpp
===================================================================
--- clang-tools-extra/include-cleaner/lib/Record.cpp
+++ 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,20 @@
       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 =
@@ -256,23 +268,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 +290,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 +304,16 @@
   };
   // 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) {
_______________________________________________
cfe-commits mailing list
cfe-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to