[clang-tools-extra] 10d183b - [include-cleaner] Capture private headers in PragmaIncludes.

[Haojian Wu via cfe-commits]

Author: Viktoriia Bakalova
Date: 2022-11-28T10:02:02+01:00
New Revision: 10d183b889daab4512d476c1645d24d4e8946e8c

URL: https://github.com/llvm/llvm-project/commit/10d183b889daab4512d476c1645d24d4e8946e8c
DIFF: https://github.com/llvm/llvm-project/commit/10d183b889daab4512d476c1645d24d4e8946e8c.diff

LOG: [include-cleaner] Capture private headers in PragmaIncludes.

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

Reviewed By: hokein

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

Added: 
    

Modified: 
    clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h
    clang-tools-extra/include-cleaner/lib/Record.cpp
    clang-tools-extra/include-cleaner/unittests/RecordTest.cpp

Removed: 
    


################################################################################
diff  --git a/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h b/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h
index a569f219939d..66fd0c7915fa 100644
--- a/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h
+++ b/clang-tools-extra/include-cleaner/include/clang-include-cleaner/Record.h
@@ -17,16 +17,15 @@
 #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>
 
@@ -73,6 +72,9 @@ class PragmaIncludes {
   /// 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
@@ -80,7 +82,8 @@ class PragmaIncludes {
   /// 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

diff  --git a/clang-tools-extra/include-cleaner/lib/Record.cpp b/clang-tools-extra/include-cleaner/lib/Record.cpp
index eb9f5c3e390b..e93bc14d4b46 100644
--- a/clang-tools-extra/include-cleaner/lib/Record.cpp
+++ b/clang-tools-extra/include-cleaner/lib/Record.cpp
@@ -233,14 +233,18 @@ class PragmaIncludes::RecordPragma : public PPCallbacks, public CommentHandler {
     if (!Pragma)
       return false;
 
-    if (Pragma->consume_front("private, 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("\"")
-                      ? (*Pragma)
-                      : ("\"" + *Pragma + "\"").str())});
+    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());
@@ -346,6 +350,10 @@ 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;

diff  --git a/clang-tools-extra/include-cleaner/unittests/RecordTest.cpp b/clang-tools-extra/include-cleaner/unittests/RecordTest.cpp
index fb4e0632fd24..0cc163751fd7 100644
--- a/clang-tools-extra/include-cleaner/unittests/RecordTest.cpp
+++ b/clang-tools-extra/include-cleaner/unittests/RecordTest.cpp
@@ -346,18 +346,30 @@ 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.
+  EXPECT_FALSE(PI.isPrivate(PublicFE.get()));
+
+  auto Private2FE = Processed.fileManager().getFile("private2.h");
+  assert(Private2FE);
+  EXPECT_TRUE(PI.isPrivate(Private2FE.get()));
 }
 
 TEST_F(PragmaIncludeTest, IWYUExport) {