[include-cleaner] Unify always_keep with rest of the keep logic

Depends on D156122

Differential Revision: https://reviews.llvm.org/D156123

clang-tools-extra/clang-tidy/misc/IncludeCleanerCheck.cpp

@@ -143,7 +143,7 @@ void IncludeCleanerCheck::check(const MatchFinder::MatchResult &Result) {
        RecordedPreprocessor.Includes.all()) {
     if (Used.contains(&I) || !I.Resolved)
       continue;
-    if (RecordedPI.shouldKeep(I.Line) || RecordedPI.shouldKeep(*I.Resolved))
+    if (RecordedPI.shouldKeep(*I.Resolved))
       continue;
     // Check if main file is the public interface for a private header. If so
     // we shouldn't diagnose it as unused.

clang-tools-extra/clangd/IncludeCleaner.cpp

@@ -75,7 +75,7 @@ bool mayConsiderUnused(const Inclusion &Inc, ParsedAST &AST,
   auto FE = AST.getSourceManager().getFileManager().getFileRef(
       AST.getIncludeStructure().getRealPath(HID));
   assert(FE);
-  if (PI && (PI->shouldKeep(Inc.HashLine + 1) || PI->shouldKeep(*FE)))
+  if (PI && PI->shouldKeep(*FE))
     return false;
   // FIXME(kirillbobyrev): We currently do not support the umbrella headers.
   // System headers are likely to be standard library headers.

clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp

@@ -15,7 +15,6 @@
 #include "TestTU.h"
 #include "clang-include-cleaner/Analysis.h"
 #include "clang-include-cleaner/Types.h"
-#include "support/Context.h"
 #include "clang/AST/DeclBase.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Tooling/Syntax/Tokens.h"
@@ -26,13 +25,11 @@
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ScopedPrinter.h"
-#include "llvm/Testing/Support/SupportHelpers.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <cstddef>
 #include <optional>
 #include <string>
-#include <utility>
 #include <vector>
 
 namespace clang {

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h

@@ -59,7 +59,6 @@ class PragmaIncludes {
 
   /// Returns true if the given #include of the main-file should never be
   /// removed.
-  bool shouldKeep(unsigned HashLineNumber) const;
   bool shouldKeep(const FileEntry *FE) const;
 
   /// Returns the public mapping include for the given physical header file.
@@ -81,10 +80,6 @@ class PragmaIncludes {
 
 private:
   class RecordPragma;
-  /// 1-based Line numbers for the #include directives of the main file that
-  /// should always keep (e.g. has the `IWYU pragma: keep` or `IWYU pragma:
-  /// export` right after).
-  llvm::DenseSet</*LineNumber*/ unsigned> ShouldKeep;
 
   /// The public header mapping by the IWYU private pragma. For private pragmas
   //  without public mapping an empty StringRef is stored.
@@ -112,8 +107,10 @@ class PragmaIncludes {
 
   /// Contains all non self-contained files detected during the parsing.
   llvm::DenseSet<llvm::sys::fs::UniqueID> NonSelfContainedFiles;
-  // Files with an always_keep pragma.
-  llvm::DenseSet<llvm::sys::fs::UniqueID> AlwaysKeep;
+  // Files whose inclusions shouldn't be dropped. E.g. because they have an
+  // always_keep pragma or because user marked particular includes with
+  // keep/export pragmas in the main file.
+  llvm::DenseSet<llvm::sys::fs::UniqueID> ShouldKeep;
 
   /// Owns the strings.
   llvm::BumpPtrAllocator Arena;

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

@@ -91,7 +91,7 @@ analyze(llvm::ArrayRef<Decl *> ASTRoots,
         HeaderFilter(I.Resolved->getFileEntry().tryGetRealPathName()))
       continue;
     if (PI) {
-      if (PI->shouldKeep(I.Line) || PI->shouldKeep(*I.Resolved))
+      if (PI->shouldKeep(*I.Resolved))
         continue;
       // Check if main file is the public interface for a private header. If so
       // we shouldn't diagnose it as unused.

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

@@ -219,12 +219,13 @@ class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
       }
     if (!IncludedHeader && File)
       IncludedHeader = &File->getFileEntry();
-    checkForExport(HashFID, HashLine, std::move(IncludedHeader));
-    checkForKeep(HashLine);
+    checkForExport(HashFID, HashLine, std::move(IncludedHeader), File);
+    checkForKeep(HashLine, File);
   }
 
   void checkForExport(FileID IncludingFile, int HashLine,
-                      std::optional<Header> IncludedHeader) {
+                      std::optional<Header> IncludedHeader,
+                      OptionalFileEntryRef IncludedFile) {
     if (ExportStack.empty())
       return;
     auto &Top = ExportStack.back();
@@ -248,20 +249,22 @@ class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
         }
       }
       // main-file #include with export pragma should never be removed.
-      if (Top.SeenAtFile == SM.getMainFileID())
-        Out->ShouldKeep.insert(HashLine);
+      if (Top.SeenAtFile == SM.getMainFileID() && IncludedFile)
+        Out->ShouldKeep.insert(IncludedFile->getUniqueID());
     }
     if (!Top.Block) // Pop immediately for single-line export pragma.
       ExportStack.pop_back();
   }
 
-  void checkForKeep(int HashLine) {
+  void checkForKeep(int HashLine, OptionalFileEntryRef IncludedFile) {
     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 (IncludedFile && ((Top.Block && HashLine > Top.SeenAtLine) ||
+                         Top.SeenAtLine == HashLine)) {
+      Out->ShouldKeep.insert(IncludedFile->getUniqueID());
+    }
 
     if (!Top.Block)
       KeepStack.pop_back(); // Pop immediately for single-line keep pragma.
@@ -321,7 +324,7 @@ class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
       return false;
     }
     if (Pragma->consume_front("always_keep")) {
-      Out->AlwaysKeep.insert(CommentUID);
+      Out->ShouldKeep.insert(CommentUID);
       return false;
     }
     return false;
@@ -418,12 +421,8 @@ bool PragmaIncludes::isPrivate(const FileEntry *FE) const {
   return IWYUPublic.contains(FE->getUniqueID());
 }
 
-bool PragmaIncludes::shouldKeep(unsigned HashLineNumber) const {
-  return ShouldKeep.contains(HashLineNumber);
-}
-
 bool PragmaIncludes::shouldKeep(const FileEntry *FE) const {
-  return AlwaysKeep.contains(FE->getUniqueID());
+  return ShouldKeep.contains(FE->getUniqueID());
 }
 
 namespace {

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

@@ -13,10 +13,12 @@
 #include "clang/Testing/TestAST.h"
 #include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Testing/Annotations/Annotations.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include <cassert>
 
 namespace clang::include_cleaner {
 namespace {
@@ -309,68 +311,58 @@ class PragmaIncludeTest : public ::testing::Test {
 };
 
 TEST_F(PragmaIncludeTest, IWYUKeep) {
-  llvm::Annotations MainFile(R"cpp(
-    $keep1^#include "keep1.h" // IWYU pragma: keep
-    $keep2^#include "keep2.h" /* IWYU pragma: keep */
+  Inputs.Code = R"cpp(
+    #include "keep1.h" // IWYU pragma: keep
+    #include "keep2.h" /* IWYU pragma: keep */
 
-    $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
+    #include "export1.h" // IWYU pragma: export
+    // IWYU pragma: begin_exports
+    #include "export2.h"
+    #include "export3.h"
+    // IWYU pragma: end_exports
 
-    $normal^#include "normal.h"
+    #include "normal.h"
 
-    $begin_keep^// IWYU pragma: begin_keep 
-    $keep3^#include "keep3.h"
-    $end_keep^// IWYU pragma: end_keep
+    // IWYU pragma: begin_keep
+    #include "keep3.h"
+    // IWYU pragma: end_keep
 
-    // IWYU pragma: begin_keep 
-    $keep4^#include "keep4.h"
     // IWYU pragma: begin_keep
-    $keep5^#include "keep5.h"
+    #include "keep4.h"
+    // IWYU pragma: begin_keep
+    #include "keep5.h"
     // IWYU pragma: end_keep
-    $keep6^#include "keep6.h"
+    #include "keep6.h"
     // IWYU pragma: end_keep
-  )cpp");
-
-  auto OffsetToLineNum = [&MainFile](size_t Offset) {
-    int Count = MainFile.code().substr(0, Offset).count('\n');
-    return Count + 1;
-  };
-
-  Inputs.Code = MainFile.code();
+    #include <vector>
+    #include <map> // IWYU pragma: keep
+    #include <set> // IWYU pragma: export
+  )cpp";
   createEmptyFiles({"keep1.h", "keep2.h", "keep3.h", "keep4.h", "keep5.h",
                     "keep6.h", "export1.h", "export2.h", "export3.h",
-                    "normal.h"});
+                    "normal.h", "std/vector", "std/map", "std/set"});
 
+  Inputs.ExtraArgs.push_back("-isystemstd");
   TestAST Processed = build();
-  EXPECT_FALSE(PI.shouldKeep(1));
-
-  // Keep
-  EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("keep1"))));
-  EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("keep2"))));
+  auto &FM = Processed.fileManager();
 
-  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_FALSE(PI.shouldKeep(FM.getFile("normal.h").get()));
+  EXPECT_FALSE(PI.shouldKeep(FM.getFile("std/vector").get()));
 
-  EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("keep4"))));
-  EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("keep5"))));
-  EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("keep6"))));
+  // Keep
+  EXPECT_TRUE(PI.shouldKeep(FM.getFile("keep1.h").get()));
+  EXPECT_TRUE(PI.shouldKeep(FM.getFile("keep2.h").get()));
+  EXPECT_TRUE(PI.shouldKeep(FM.getFile("keep3.h").get()));
+  EXPECT_TRUE(PI.shouldKeep(FM.getFile("keep4.h").get()));
+  EXPECT_TRUE(PI.shouldKeep(FM.getFile("keep5.h").get()));
+  EXPECT_TRUE(PI.shouldKeep(FM.getFile("keep6.h").get()));
+  EXPECT_TRUE(PI.shouldKeep(FM.getFile("std/map").get()));
 
   // Exports
-  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(OffsetToLineNum(MainFile.point("normal"))));
+  EXPECT_TRUE(PI.shouldKeep(FM.getFile("export1.h").get()));
+  EXPECT_TRUE(PI.shouldKeep(FM.getFile("export2.h").get()));
+  EXPECT_TRUE(PI.shouldKeep(FM.getFile("export3.h").get()));
+  EXPECT_TRUE(PI.shouldKeep(FM.getFile("std/set").get()));
 }
 
 TEST_F(PragmaIncludeTest, IWYUPrivate) {