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 --- 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 #include @@ -73,6 +72,9 @@ /// 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 @@ /// export` right after). llvm::DenseSet 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 --- a/clang-tools-extra/include-cleaner/lib/Record.cpp +++ b/clang-tools-extra/include-cleaner/lib/Record.cpp @@ -233,14 +233,18 @@ 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 @@ return !NonSelfContainedFiles.contains(FE->getUniqueID()); } +bool PragmaIncludes::isPrivate(const FileEntry *FE) const { + return IWYUPublic.find(FE->getUniqueID()) != IWYUPublic.end(); +} + std::unique_ptr 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 --- a/clang-tools-extra/include-cleaner/unittests/RecordTest.cpp +++ b/clang-tools-extra/include-cleaner/unittests/RecordTest.cpp @@ -346,18 +346,30 @@ 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) {