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

[kadircet]

Summary

Depends on D156122

[VitaNuo]

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

248

If I understand correctly, you should be able to extract IncludedFile from IncludedHeader->physical() and then you don't need the extra parameter.

Also, technically this and the below code used to work for standard and verbatim headers before this change. Now it probably won't, since there will be no corresponding IncludedFile. Is this actually something to worry about?

262

nit: remove braces.

419

Consider inlining this function.

[kadircet]

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

248

If I understand correctly, you should be able to extract IncludedFile from IncludedHeader->physical() and then you don't need the extra parameter.

Not really, if included file is recognized as a system include we won't have a physical file entry.

Also, technically this and the below code used to work for standard and verbatim headers before this change. Now it probably won't, since there will be no corresponding IncludedFile. Is this actually something to worry about?

It still does (and I am adding tests for those). All of them were still recorded based on their line number into this structure, now instead we're recording them through their fileentry (even if their include_cleaner::Header representation is non-physical, we're talking about an include here, so it must have a logical file it resolves to, when it exists). So this still implies some behavior changes of course:

• We no longer respect IWYU pragmas assoaciated with non-existent/virtual files, but all of the users already suppress unused diagnostics for unresolved files and moreover the new shouldKeep API forces the caller to pass in a file-entry. Hence going forward we give explicit info to the callers that they should deal with unresolved headers on their own way.
• If there are multiple include directives resolving to the same physical file, we'll bind their fate. We already assume in other places that a file will be included at most once, so I don't think that's also a worth concern.

WDYT?

262

because the condition is spanning multiple lines, i'd rather keep it

419

same concerns about inlining as in the previous patch

[kadircet]

• Add tests for pragmas on stdlib headers

[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) || 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.
     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 mayConsiderUnused(const Inclusion &Inc, ParsedAST &AST,
   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) {
-    if (PI->shouldKeep(Inc.HashLine + 1) || PI->shouldKeep(*FE))
+    if (PI->shouldKeep(*FE))
       return false;
     // 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/clangd/unittests/IncludeCleanerTests.cpp

 #include "Diagnostics.h"
 #include "IncludeCleaner.h"
 #include "ParsedAST.h"
 #include "SourceCode.h"
 #include "TestFS.h"
 #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"
-#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/StringRef.h"
 #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 {
 namespace clangd {
 namespace {
 
 using ::testing::AllOf;
 using ::testing::ElementsAre;

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;
   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,
                                                     FileManager &FM) const;
   llvm::SmallVector<const FileEntry *> getExporters(tooling::stdlib::Header,
                                                     FileManager &FM) const;
 
   /// Returns true if the given file is a self-contained file.
   bool isSelfContained(const FileEntry *File) const;
 
   /// Returns true if the given file is marked with the IWYU private pragma.
   bool isPrivate(const FileEntry *File) const;
 
 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.
   //
   // !!NOTE: instead of using a FileEntry* to identify the physical file, we
   // deliberately use the UniqueID to ensure the result is stable across
   // FileManagers (for clangd's preamble and main-file builds).
   llvm::DenseMap<llvm::sys::fs::UniqueID, llvm::StringRef /*VerbatimSpelling*/>
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;
+  // 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;
 
   // 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->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.
       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

void InclusionDirective(SourceLocation HashLoc, const Token &IncludeTok,
     std::optional<Header> IncludedHeader;
     if (IsAngled)
       if (auto StandardHeader =
               tooling::stdlib::Header::named("<" + FileName.str() + ">")) {
         IncludedHeader = *StandardHeader;
       }
     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();
     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) {
         switch (IncludedHeader->kind()) {
         case Header::Physical:
           Out->IWYUExportBy[IncludedHeader->physical()->getUniqueID()]
               .push_back(Top.Path);
           break;
         case Header::Standard:
           Out->StdIWYUExportBy[IncludedHeader->standard()].push_back(Top.Path);
           break;
         case Header::Verbatim:
           assert(false && "unexpected Verbatim header");
           break;
         }
       }
       // 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.
   }
 
   bool HandleComment(Preprocessor &PP, SourceRange Range) override {
     auto &SM = PP.getSourceManager();
     auto Pragma =
if (Pragma->consume_front("private")) {
         PublicHeader = save(Pragma->startswith("<") || Pragma->startswith("\"")
                                 ? (*Pragma)
                                 : ("\"" + *Pragma + "\"").str());
       }
       Out->IWYUPublic.insert({FE->getLastRef().getUniqueID(), PublicHeader});
       return false;
     }
     if (Pragma->consume_front("always_keep")) {
-      Out->AlwaysKeep.insert(FE->getUniqueID());
+      Out->ShouldKeep.insert(FE->getUniqueID());
       return false;
     }
     int CommentLine = SM.getLineNumber(SM.getFileID(Range.getBegin()),
                                        SM.getFileOffset(Range.getBegin()));
     // Record export pragma.
     if (Pragma->startswith("export")) {
       ExportStack.push_back({CommentLine, CommentFID,
                              save(SM.getFileEntryForID(CommentFID)->getName()),

 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());
+  return ShouldKeep.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;
 }

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

 //===-- RecordTest.cpp ----------------------------------------------------===//
 //
 // 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/Basic/SourceLocation.h"
 #include "clang/Frontend/FrontendAction.h"
 #include "clang/Frontend/FrontendActions.h"
 #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 {
 using testing::ElementsAreArray;
 using testing::IsEmpty;
 
 // Matches a Decl* if it is a NamedDecl with the given name.
 MATCHER_P(named, N, "") {
protected:
 
   void createEmptyFiles(llvm::ArrayRef<StringRef> FileNames) {
     for (llvm::StringRef File : FileNames)
       Inputs.ExtraFiles[File] = "";
   }
 };
 
 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
-
-    $normal^#include "normal.h"
-
-    $begin_keep^// IWYU pragma: begin_keep 
-    $keep3^#include "keep3.h"
-    $end_keep^// IWYU pragma: end_keep
+    #include "export1.h" // IWYU pragma: export
+    // IWYU pragma: begin_exports
+    #include "export2.h"
+    #include "export3.h"
+    // IWYU pragma: end_exports
+
+    #include "normal.h"
+
+    // IWYU pragma: begin_keep
+    #include "keep3.h"
+    // IWYU pragma: end_keep
 
-    // IWYU pragma: begin_keep 
-    $keep4^#include "keep4.h"
+    // IWYU pragma: begin_keep
+    #include "keep4.h"
     // IWYU pragma: begin_keep
-    $keep5^#include "keep5.h"
+    #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();
+  )cpp";
   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));
+  auto &FM = Processed.fileManager();
 
-  // Keep
-  EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("keep1"))));
-  EXPECT_TRUE(PI.shouldKeep(OffsetToLineNum(MainFile.point("keep2"))));
+  EXPECT_FALSE(PI.shouldKeep(FM.getFile("normal.h").get()));
 
-  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"))));
+  // 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()));
 
   // 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()));
 }
 
 TEST_F(PragmaIncludeTest, IWYUPrivate) {
   Inputs.Code = R"cpp(
     #include "public.h"
   )cpp";
   Inputs.ExtraFiles["public.h"] = R"cpp(
     #include "private.h"