[include-cleaner] Capture private headers in PragmaIncludes.

[VitaNuo]

Summary

Save file IDs of IWYU private headers and report them as private.

[VitaNuo]

Improve the comment.

[VitaNuo]

Fix formatting.

[hokein]

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

76

I think this function should return true if the given file is marked with the IWYU private mapping pragma, because conceptually these file are also private.

95

It should collect files with mapping IWYU private pragmas.

and I think the last sentence can be removed (the intention of using UniqueID is clear from other comments).

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

269

I'd suggest moving this to the above line 236 where we handle the private mapping.

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

350

no extra ; token needed for the #include.

371

add EXPECT_TRUE(PI.isPrivate(PrivateFE.get())).

[kadircet]

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

97

i'd actually merge this with the previous map, and store empty verbatim spellings. semantics of getPublic is to return empty path when there are no mappings anyway.

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

361

nit: no need for the declaration of Private here (nor in private.h). it's actually introducing a double definition error into TU now.

[VitaNuo]

Address review comments.

[VitaNuo]

The formatting has introduced a bit of unrelated diff. It's not much, so I hope that's fine.

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

95

I've followed Kadir's suggestion below now and merged the private pragmas with the IWYUPublic map.

97

Sounds reasonable. I wasn't sure which of the options to pick in the first place.

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

361

Out of curiosity: what's TU?

[VitaNuo]

Fix build.

[VitaNuo]

Fix comment.

[kadircet]

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

236

nit: you can drop the braces here. LLVM convention is to not have braces when a condition/loop body has only a single statement.

239

nit: you can also rewrite this as:

StringRef PublicHeader;
if (Pragma->consume_front(", include ")) {
  PublicHeader = save(Pragma->startswith("<") || Pragma->startswith("\"")
                      ? (*Pragma)
                      : ("\"" + *Pragma + "\"").str());
}
Out->IWYUPublic.insert({FE->getLastRef().getUniqueID(), PublicHeader});

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

361

it means translation unit, it's the source file provided to the compiler after all preprocessor directives are processed.

[VitaNuo]

Small refactoring.

[VitaNuo]

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

236

Ah ok, thank you. It's a bit funny, my whole life until now I've been rigorously taught to never do this ;-)

239

Yeah this is good, thanks.

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

361

Thanks.

[VitaNuo]

Fix build.

[VitaNuo]

Fix history.

[VitaNuo]

Apply the local changes.

[hokein]

Thanks, looks good from my side, a few nits on the test.

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

364

nit: IIUC, we're using a blank line to separate the tests per FileEntry, this EXPECT_EQ statement should belong to the above group (which is PrivateFE), right?

367–371

nit: add EXPECT_FALSE(PI.isPrivate(PublicFE.get()));

[VitaNuo]

Address the remaining comments.

[VitaNuo]

Rebase.

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

 // These are the simplest way to connect include-cleaner logic to the parser,
 // but other ways are possible (for example clangd records includes separately).
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef CLANG_INCLUDE_CLEANER_RECORD_H
 #define CLANG_INCLUDE_CLEANER_RECORD_H
 
+#include "clang-include-cleaner/Types.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/FileSystem/UniqueID.h"
-#include "clang-include-cleaner/Types.h"
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/StringMap.h"
 #include <memory>
 #include <vector>
 
 namespace clang {
 class ASTConsumer;
 class ASTContext;
 class CompilerInstance;
 class Decl;
public:
   /// 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;
 
   /// 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.
+  /// 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*/>
       IWYUPublic;
 
   /// A reverse map from the underlying header to its exporter headers.
   //
   //  There's no way to get a FileEntry from a UniqueID, especially when it
   //  hasn't been opened before. So store the full path and convert it to a
   //  FileEntry by opening the file again through a FileManager.
   llvm::DenseMap<llvm::sys::fs::UniqueID,
                  llvm::SmallVector</*RealPathNames*/ llvm::StringRef>>
       IWYUExportBy;
 
   /// Contains all non self-contained files detected during the parsing.
   llvm::DenseSet<llvm::sys::fs::UniqueID> NonSelfContainedFiles;
 

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

void InclusionDirective(SourceLocation Hash, const Token &IncludeTok,
     I.Spelled = SpelledFilename;
     Recorded.Includes.add(I);
   }
 
   void MacroExpands(const Token &MacroName, const MacroDefinition &MD,
                     SourceRange Range, const MacroArgs *Args) override {
     if (!Active)
       return;
-    recordMacroRef(MacroName, *MD.getMacroInfo());
+    recordMacroRef(MacroName, *MD.getMacroInfo(), RefType::Explicit);
   }
 
   void MacroDefined(const Token &MacroName, const MacroDirective *MD) override {
     if (!Active)
       return;
 
     const auto *MI = MD->getMacroInfo();
     // The tokens of a macro definition could refer to a macro.
     // Formally this reference isn't resolved until this macro is expanded,
     // but we want to treat it as a reference anyway.
     for (const auto &Tok : MI->tokens()) {
       auto *II = Tok.getIdentifierInfo();
       // Could this token be a reference to a macro? (Not param to this macro).
       if (!II || !II->hadMacroDefinition() ||
           llvm::is_contained(MI->params(), II))
         continue;
       if (const MacroInfo *MI = PP.getMacroInfo(II))
-        recordMacroRef(Tok, *MI);
+        recordMacroRef(Tok, *MI, RefType::Explicit);
     }
   }
 
   void MacroUndefined(const Token &MacroName, const MacroDefinition &MD,
                       const MacroDirective *) override {
     if (!Active)
       return;
     if (const auto *MI = MD.getMacroInfo())
-      recordMacroRef(MacroName, *MI);
+      recordMacroRef(MacroName, *MI, RefType::Explicit);
   }
 
   void Ifdef(SourceLocation Loc, const Token &MacroNameTok,
              const MacroDefinition &MD) override {
     if (!Active)
       return;
     if (const auto *MI = MD.getMacroInfo())
       recordMacroRef(MacroNameTok, *MI, RefType::Ambiguous);
void Defined(const Token &MacroNameTok, const MacroDefinition &MD,
                SourceRange Range) override {
     if (!Active)
       return;
     if (const auto *MI = MD.getMacroInfo())
       recordMacroRef(MacroNameTok, *MI, RefType::Ambiguous);
   }
 
 private:
-  void recordMacroRef(const Token &Tok, const MacroInfo &MI,
-                      RefType RT = RefType::Explicit) {
+  void recordMacroRef(const Token &Tok, const MacroInfo &MI, RefType RT) {
     if (MI.isBuiltinMacro())
       return; // __FILE__ is not a reference.
     Recorded.MacroReferences.push_back(SymbolReference{
         Tok.getLocation(),
         Macro{Tok.getIdentifierInfo(), MI.getDefinitionLoc()}, RT});
   }
 
   bool Active = false;
public:
   }
 
   bool HandleComment(Preprocessor &PP, SourceRange Range) override {
     auto &SM = PP.getSourceManager();
     auto Pragma = parseIWYUPragma(SM.getCharacterData(Range.getBegin()));
     if (!Pragma)
       return false;
 
-    if (Pragma->consume_front("private, include ")) {
+    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.
-      if (auto *FE = SM.getFileEntryForID(SM.getFileID(Range.getBegin())))
-        Out->IWYUPublic.insert(
-            {FE->getLastRef().getUniqueID(),
-             save(Pragma->startswith("<") || Pragma->startswith("\"")
+        PublicHeader = save(Pragma->startswith("<") || Pragma->startswith("\"")
                                 ? (*Pragma)
-                      : ("\"" + *Pragma + "\"").str())});
+                                : ("\"" + *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()));
     // Record export pragma.
     if (Pragma->startswith("export")) {
       ExportStack.push_back(
           {CommentLine, CommentFID,
            save(SM.getFileEntryForID(CommentFID)->tryGetRealPathName()),
            false});
     } else if (Pragma->startswith("begin_exports")) {
       ExportStack.push_back(
           {CommentLine, CommentFID,
            save(SM.getFileEntryForID(CommentFID)->tryGetRealPathName()), 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"))
         return false;
       // Given:
       //
       // #include "foo.h"
       // #include "bar.h" // IWYU pragma: keep
PragmaIncludes::getExporters(const FileEntry *File, FileManager &FM) const {
   }
   return Results;
 }
 
 bool PragmaIncludes::isSelfContained(const FileEntry *FE) const {
   return !NonSelfContainedFiles.contains(FE->getUniqueID());
 }
 
+bool PragmaIncludes::isPrivate(const FileEntry *FE) const {
+  return IWYUPublic.find(FE->getUniqueID()) != IWYUPublic.end();
+}
+
 std::unique_ptr<ASTConsumer> RecordedAST::record() {
   class Recorder : public ASTConsumer {
     RecordedAST *Out;
 
   public:
     Recorder(RecordedAST *Out) : Out(Out) {}
     void Initialize(ASTContext &Ctx) override { Out->Ctx = &Ctx; }
     bool HandleTopLevelDecl(DeclGroupRef DG) override {

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

TEST_F(PragmaIncludeTest, IWYUKeep) {
 
   EXPECT_FALSE(PI.shouldKeep(11));
 }
 
 TEST_F(PragmaIncludeTest, IWYUPrivate) {
   Inputs.Code = R"cpp(
     #include "public.h"
   )cpp";
-  Inputs.ExtraFiles["public.h"] = "#include \"private.h\"";
+  Inputs.ExtraFiles["public.h"] = R"cpp(
+    #include "private.h"
+    #include "private2.h"
+  )cpp";
   Inputs.ExtraFiles["private.h"] = R"cpp(
     // IWYU pragma: private, include "public2.h"
-    class Private {};
+  )cpp";
+  Inputs.ExtraFiles["private2.h"] = R"cpp(
+    // IWYU pragma: private
   )cpp";
   TestAST Processed = build();
   auto PrivateFE = Processed.fileManager().getFile("private.h");
   assert(PrivateFE);
+  EXPECT_TRUE(PI.isPrivate(PrivateFE.get()));
+
   EXPECT_EQ(PI.getPublic(PrivateFE.get()), "\"public2.h\"");
   auto PublicFE = Processed.fileManager().getFile("public.h");
   assert(PublicFE);
   EXPECT_EQ(PI.getPublic(PublicFE.get()), ""); // no mapping.
+
+  auto Private2FE = Processed.fileManager().getFile("private2.h");
+  assert(Private2FE);
+  EXPECT_TRUE(PI.isPrivate(Private2FE.get()));
 }
 
 TEST_F(PragmaIncludeTest, IWYUExport) {
   Inputs.Code = R"cpp(// Line 1
     #include "export1.h"
     #include "export2.h"
   )cpp";
   Inputs.ExtraFiles["export1.h"] = R"cpp(