[include-cleaner] Introduce support for always_keep pragma

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

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))
+    if (RecordedPI.shouldKeep(I.Line) || 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

@@ -68,21 +68,20 @@ bool isIgnored(llvm::StringRef HeaderPath, HeaderFilter IgnoreHeaders) {
   return false;
 }
 
-bool mayConsiderUnused(
-    const Inclusion &Inc, ParsedAST &AST,
-    const include_cleaner::PragmaIncludes *PI) {
-  if (PI && PI->shouldKeep(Inc.HashLine + 1))
-      return false;
-  // FIXME(kirillbobyrev): We currently do not support the umbrella headers.
-  // System headers are likely to be standard library headers.
-  // Until we have good support for umbrella headers, don't warn about them.
-  if (Inc.Written.front() == '<')
-    return tooling::stdlib::Header::named(Inc.Written).has_value();
+bool mayConsiderUnused(const Inclusion &Inc, ParsedAST &AST,
+                       const include_cleaner::PragmaIncludes *PI) {
   assert(Inc.HeaderID);
   auto HID = static_cast<IncludeStructure::HeaderID>(*Inc.HeaderID);
   auto FE = AST.getSourceManager().getFileManager().getFileRef(
       AST.getIncludeStructure().getRealPath(HID));
   assert(FE);
+  if (PI && (PI->shouldKeep(Inc.HashLine + 1) || PI->shouldKeep(*FE)))
+    return false;
+  // FIXME(kirillbobyrev): We currently do not support the umbrella headers.
+  // System headers are likely to be standard library headers.
+  // Until we have good support for umbrella headers, don't warn about them.
+  if (Inc.Written.front() == '<')
+    return tooling::stdlib::Header::named(Inc.Written).has_value();
   if (PI) {
     // 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/include/clang-include-cleaner/Record.h

@@ -59,9 +59,8 @@ public:
 
   /// Returns true if the given #include of the main-file should never be
   /// removed.
-  bool shouldKeep(unsigned HashLineNumber) const {
-    return ShouldKeep.contains(HashLineNumber);
-  }
+  bool shouldKeep(unsigned HashLineNumber) const;
+  bool shouldKeep(const FileEntry *FE) const;
 
   /// Returns the public mapping include for the given physical header file.
   /// Returns "" if there is none.
@@ -113,6 +112,8 @@ private:
 
   /// 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;
 
   /// 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))
+      if (PI->shouldKeep(I.Line) || 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

@@ -11,6 +11,9 @@
 #include "clang/AST/ASTConsumer.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/DeclGroup.h"
+#include "clang/Basic/FileEntry.h"
+#include "clang/Basic/LLVM.h"
+#include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/Specifiers.h"
 #include "clang/Frontend/CompilerInstance.h"
@@ -20,8 +23,17 @@
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/Inclusions/HeaderAnalysis.h"
 #include "clang/Tooling/Inclusions/StandardLibrary.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Allocator.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/StringSaver.h"
+#include <algorithm>
+#include <assert.h>
 #include <memory>
+#include <optional>
 #include <utility>
 #include <vector>
 
@@ -262,32 +274,14 @@ public:
     if (!Pragma)
       return false;
 
-    if (Pragma->consume_front("private")) {
-      auto *FE = SM.getFileEntryForID(SM.getFileID(Range.getBegin()));
-      if (!FE)
-        return false;
-      StringRef PublicHeader;
-      if (Pragma->consume_front(", include ")) {
-        // We always insert using the spelling from the pragma.
-        PublicHeader = save(Pragma->startswith("<") || Pragma->startswith("\"")
-                                ? (*Pragma)
-                                : ("\"" + *Pragma + "\"").str());
-      }
-      Out->IWYUPublic.insert({FE->getLastRef().getUniqueID(), PublicHeader});
-      return false;
-    }
-    FileID CommentFID = SM.getFileID(Range.getBegin());
-    int CommentLine = SM.getLineNumber(SM.getFileID(Range.getBegin()),
-                                       SM.getFileOffset(Range.getBegin()));
+    auto [CommentFID, CommentOffset] = SM.getDecomposedLoc(Range.getBegin());
+    int CommentLine = SM.getLineNumber(CommentFID, CommentOffset);
+    auto Filename = SM.getBufferName(Range.getBegin());
     // Record export pragma.
     if (Pragma->startswith("export")) {
-      ExportStack.push_back({CommentLine, CommentFID,
-                             save(SM.getFileEntryForID(CommentFID)->getName()),
-                             false});
+      ExportStack.push_back({CommentLine, CommentFID, save(Filename), false});
     } else if (Pragma->startswith("begin_exports")) {
-      ExportStack.push_back({CommentLine, CommentFID,
-                             save(SM.getFileEntryForID(CommentFID)->getName()),
-                             true});
+      ExportStack.push_back({CommentLine, CommentFID, save(Filename), true});
     } else if (Pragma->startswith("end_exports")) {
       // FIXME: be robust on unmatching cases. We should only pop the stack if
       // the begin_exports and end_exports is in the same file.
@@ -307,6 +301,29 @@ public:
         KeepStack.pop_back();
       }
     }
+
+    auto FE = SM.getFileEntryRefForID(CommentFID);
+    if (!FE) {
+      // FIXME: Support IWYU pragmas in virtual files. Our mappings rely on
+      // "persistent" UniqueIDs and that is not the case for virtual files.
+      return false;
+    }
+    auto CommentUID = FE->getUniqueID();
+    if (Pragma->consume_front("private")) {
+      StringRef PublicHeader;
+      if (Pragma->consume_front(", include ")) {
+        // We always insert using the spelling from the pragma.
+        PublicHeader = save(Pragma->startswith("<") || Pragma->startswith("\"")
+                                ? (*Pragma)
+                                : ("\"" + *Pragma + "\"").str());
+      }
+      Out->IWYUPublic.insert({CommentUID, PublicHeader});
+      return false;
+    }
+    if (Pragma->consume_front("always_keep")) {
+      Out->AlwaysKeep.insert(CommentUID);
+      return false;
+    }
     return false;
   }
 
@@ -401,6 +418,14 @@ 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());
+}
+
 namespace {
 template <typename T> bool isImplicitTemplateSpecialization(const Decl *D) {
   if (const auto *TD = dyn_cast<T>(D))

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

@@ -502,5 +502,20 @@ TEST_F(PragmaIncludeTest, SelfContained) {
   EXPECT_FALSE(PI.isSelfContained(FM.getFile("unguarded.h").get()));
 }
 
+TEST_F(PragmaIncludeTest, AlwaysKeep) {
+  Inputs.Code = R"cpp(
+  #include "always_keep.h"
+  #include "usual.h"
+  )cpp";
+  Inputs.ExtraFiles["always_keep.h"] = R"cpp(
+  #pragma once
+  // IWYU pragma: always_keep
+  )cpp";
+  Inputs.ExtraFiles["usual.h"] = "#pragma once";
+  TestAST Processed = build();
+  auto &FM = Processed.fileManager();
+  EXPECT_TRUE(PI.shouldKeep(FM.getFile("always_keep.h").get()));
+  EXPECT_FALSE(PI.shouldKeep(FM.getFile("usual.h").get()));
+}
 } // namespace
 } // namespace clang::include_cleaner