[include-cleaner] Introduce support for always_keep pragma

[kadircet]

FWIW, details of pragma are explained in https://github.com/include-what-you-use/include-what-you-use/blob/master/docs/IWYUPragmas.md#iwyu-pragma-always_keep

[VitaNuo]

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

63

Why wouldn't you actually inline the implementation for both these functions? The implementations seem trivial enough.

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

38

I don't think you've added a dependency on all these headers in this patch. Nice if you've actually fixed the missing include violations all over the file, but please double-check that there are no debugging etc. artefacts left.

277

I believe you need to check FE for nullptr, similar to private pragma handling above.
Also, the code above does FE->getLastRef().getUniqueID(), is there any difference to FE->getUniqueID()? I wouldn't expect so, but then you could maybe unify this one with the above, this way or the other.

277

You might want to inline this call, it seems to have a single usage right below.

427

As mentioned above, consider inlining.

[kadircet]

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

63

most of the binaries are build with LTO nowadays, so inlining code into headers doesn't really gain much from performance perspective. moreover, these are not really hot functions, they're invoked once per #include directive inside the main file, at the end of analysis.

OTOH, when they're inlined and there needs to be a change to these function definitions, it requires all the translation units depending on this header to be recompiled. Hence I think having them in the cpp file is a not benefit for development cycles.

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

38

yes these are all fixes generated by include-cleaner, I don't have any unused include findings.

277

do you mean initializer for CommentFID? if so CommentFID is actually used in more places below

277

thanks missed that. unified all the usages for getUniqueID, they all should be the same.

you're absolutely right about checking for FE, in practice it's almost never null as it's the file in which we've found the IWYU pragma. But SourceManager actually allows working with virtual files, in which case there might not be a phyiscal file entry.
code in rest of this logic was also relying on FE being non-null without checking for it. Moved code around a little bit and used filename inside the SourceManager instead to make sure export/keep logic can work with virtual files too.

for now not adding support to always_keep and private pragmas for virtual files, as it requires deeper changes in the way we store uniqueids (we can no longer depend on them) and they're not common in practice (usually those virtual files are provided as predefines from the compiler and they likely won't contain IWYU pragmas).

[kadircet]

• Rebase
• Properly check for null-ness of FileEntry.

[VitaNuo]

Thanks!

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

walkUsed(MainFileDecls, RecordedPreprocessor.MacroReferences, &RecordedPI,
                Missing.push_back({Ref, Providers.front()});
            });
 
   std::vector<const include_cleaner::Include *> Unused;
   for (const include_cleaner::Include &I :
        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.
     if (auto PHeader = RecordedPI.getPublic(*I.Resolved); !PHeader.empty()) {
       PHeader = PHeader.trim("<>\"");
       // Since most private -> public mappings happen in a verbatim way, we
       // check textually here. This might go wrong in presence of symlinks or
       // header mappings. But that's not different than rest of the places.

clang-tools-extra/clangd/IncludeCleaner.cpp

bool isIgnored(llvm::StringRef HeaderPath, HeaderFilter IgnoreHeaders) {
   llvm::sys::path::native(NormalizedPath, llvm::sys::path::Style::posix);
   for (auto &Filter : IgnoreHeaders) {
     if (Filter(NormalizedPath))
       return true;
   }
   return false;
 }
 
-bool mayConsiderUnused(
-    const Inclusion &Inc, ParsedAST &AST,
+bool mayConsiderUnused(const Inclusion &Inc, ParsedAST &AST,
                        const include_cleaner::PragmaIncludes *PI) {
-  if (PI && PI->shouldKeep(Inc.HashLine + 1))
-      return false;
+  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)))
   // 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();
-  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) {
     // Check if main file is the public interface for a private header. If so we
     // shouldn't diagnose it as unused.
     if (auto PHeader = PI->getPublic(*FE); !PHeader.empty()) {
       PHeader = PHeader.trim("<>\"");
       // Since most private -> public mappings happen in a verbatim way, we
       // check textually here. This might go wrong in presence of symlinks or
       // header mappings. But that's not different than rest of the places.

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

public:
   void record(const CompilerInstance &CI);
 
   /// Installs an analysing PPCallback and CommentHandler and populates results
   /// to the structure.
   void record(Preprocessor &P);
 
   /// 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.
   llvm::StringRef getPublic(const FileEntry *File) const;
 
   /// Returns all direct exporter headers for the given header file.
   /// Returns empty if there is none.
   llvm::SmallVector<const FileEntry *> getExporters(const FileEntry *File,
llvm::DenseMap<llvm::sys::fs::UniqueID,
                  llvm::SmallVector</*FileEntry::getName()*/ llvm::StringRef>>
       IWYUExportBy;
   llvm::DenseMap<tooling::stdlib::Header,
                  llvm::SmallVector</*FileEntry::getName()*/ llvm::StringRef>>
       StdIWYUExportBy;
 
   /// 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;
 
   // FIXME: add support for clang use_instead pragma
 };
 
 /// Recorded main-file parser events relevant to include-cleaner.

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

walkUsed(ASTRoots, MacroRefs, PI, SM,
            });
 
   AnalysisResults Results;
   for (const Include &I : Inc.all()) {
     if (Used.contains(&I) || !I.Resolved ||
         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.
       if (auto PHeader = PI->getPublic(*I.Resolved); !PHeader.empty()) {
         PHeader = PHeader.trim("<>\"");
         // Since most private -> public mappings happen in a verbatim way, we
         // check textually here. This might go wrong in presence of symlinks or
         // header mappings. But that's not different than rest of the places.

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

 //===--- Record.cpp - Record compiler events ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "clang-include-cleaner/Record.h"
 #include "clang-include-cleaner/Types.h"
 #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"
 #include "clang/Lex/DirectoryLookup.h"
 #include "clang/Lex/MacroInfo.h"
 #include "clang/Lex/PPCallbacks.h"
 #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>
 
 namespace clang::include_cleaner {
 namespace {
 
 class PPRecorder : public PPCallbacks {
 public:
   PPRecorder(RecordedPP &Recorded, const Preprocessor &PP)
       : Recorded(Recorded), PP(PP), SM(PP.getSourceManager()) {
public:
 
   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")) {
-      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.
       if (!ExportStack.empty()) {
         assert(ExportStack.back().Block);
         ExportStack.pop_back();
       }
     }
 
     if (InMainFile) {
       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();
       }
     }
+
+    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;
   }
 
 private:
   StringRef save(llvm::StringRef S) { return UniqueStrings.save(S); }
 
   bool InMainFile = false;
   const SourceManager &SM;

 bool PragmaIncludes::isSelfContained(const FileEntry *FE) const {
   return !NonSelfContainedFiles.contains(FE->getUniqueID());
 }
 
 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))
     return TD->getTemplateSpecializationKind() == TSK_ImplicitInstantiation;
   return false;
 }
 } // namespace
 

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

TEST_F(PragmaIncludeTest, SelfContained) {
   )cpp";
   Inputs.ExtraFiles["unguarded.h"] = "";
   TestAST Processed = build();
   auto &FM = Processed.fileManager();
   EXPECT_TRUE(PI.isSelfContained(FM.getFile("guarded.h").get()));
   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