[find-all-symbols] Add IWYU private pragma support.

Reviewers: djasper, klimek

Subscribers: kimgr, cfe-commits, bkramer, ioeric

Differential Revision: http://reviews.llvm.org/D19816

llvm-svn: 269779

clang-tools-extra/include-fixer/find-all-symbols/CMakeLists.txt

@@ -4,6 +4,7 @@ set(LLVM_LINK_COMPONENTS
 
 add_clang_library(findAllSymbols
   FindAllSymbols.cpp
+  PragmaCommentHandler.cpp
   SymbolInfo.cpp
 
   LINK_LIBS

clang-tools-extra/include-fixer/find-all-symbols/FindAllSymbols.cpp

@@ -8,6 +8,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "FindAllSymbols.h"
+#include "HeaderMapCollector.h"
 #include "SymbolInfo.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclCXX.h"
@@ -56,8 +57,9 @@ std::vector<SymbolInfo::Context> GetContexts(const NamedDecl *ND) {
   return Contexts;
 }
 
-llvm::Optional<SymbolInfo> CreateSymbolInfo(const NamedDecl *ND,
-                                            const SourceManager &SM) {
+llvm::Optional<SymbolInfo>
+CreateSymbolInfo(const NamedDecl *ND, const SourceManager &SM,
+                 const HeaderMapCollector::HeaderMap &HeaderMappingTable) {
   SymbolInfo::SymbolKind Type;
   if (llvm::isa<VarDecl>(ND)) {
     Type = SymbolInfo::SymbolKind::Variable;
@@ -94,6 +96,11 @@ llvm::Optional<SymbolInfo> CreateSymbolInfo(const NamedDecl *ND,
   if (FilePath.empty())
     return llvm::None;
 
+  // Check pragma remapping header.
+  auto Iter = HeaderMappingTable.find(FilePath);
+  if (Iter != HeaderMappingTable.end())
+    FilePath = Iter->second;
+
   return SymbolInfo(ND->getNameAsString(), Type, FilePath.str(),
                     SM.getExpansionLineNumber(Loc), GetContexts(ND));
 }
@@ -207,7 +214,8 @@ void FindAllSymbols::run(const MatchFinder::MatchResult &Result) {
   assert(ND && "Matched declaration must be a NamedDecl!");
   const SourceManager *SM = Result.SourceManager;
 
-  llvm::Optional<SymbolInfo> Symbol = CreateSymbolInfo(ND, *SM);
+  llvm::Optional<SymbolInfo> Symbol =
+      CreateSymbolInfo(ND, *SM, Collector->getHeaderMappingTable());
   if (Symbol)
     Reporter->reportResult(
         SM->getFileEntryForID(SM->getMainFileID())->getName(), *Symbol);

clang-tools-extra/include-fixer/find-all-symbols/FindAllSymbols.h

@@ -17,6 +17,8 @@
 namespace clang {
 namespace find_all_symbols {
 
+class HeaderMapCollector;
+
 /// \brief FindAllSymbols collects all classes, free standing functions and
 /// global variables with some extra information such as the path of the header
 /// file, the namespaces they are contained in, the type of variables and the
@@ -39,15 +41,21 @@ class FindAllSymbols : public clang::ast_matchers::MatchFinder::MatchCallback {
                               const SymbolInfo &Symbol) = 0;
   };
 
-  explicit FindAllSymbols(ResultReporter *Reporter) : Reporter(Reporter) {}
+  explicit FindAllSymbols(ResultReporter *Reporter,
+                          HeaderMapCollector *Collector)
+      : Reporter(Reporter), Collector(Collector) {}
 
   void registerMatchers(clang::ast_matchers::MatchFinder *MatchFinder);
 
   void
   run(const clang::ast_matchers::MatchFinder::MatchResult &result) override;
 
 private:
+  // Reporter for SymbolInfo.
   ResultReporter *const Reporter;
+  // A remapping header file collector allowing clients include a different
+  // header.
+  HeaderMapCollector *const Collector;
 };
 
 } // namespace find_all_symbols

clang-tools-extra/include-fixer/find-all-symbols/HeaderMapCollector.h

@@ -0,0 +1,38 @@
+//===-- HeaderMapCoolector.h - find all symbols------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TOOLS_EXTRA_FIND_ALL_SYMBOLS_HEADER_MAP_COLLECTOR_H
+#define LLVM_CLANG_TOOLS_EXTRA_FIND_ALL_SYMBOLS_HEADER_MAP_COLLECTOR_H
+
+#include "llvm/ADT/StringMap.h"
+#include <string>
+
+namespace clang {
+namespace find_all_symbols {
+
+/// \brief HeaderMappCollector collects all remapping header files.
+class HeaderMapCollector {
+public:
+  typedef llvm::StringMap<std::string> HeaderMap;
+
+  void addHeaderMapping(llvm::StringRef OrignalHeaderPath,
+                        llvm::StringRef MappingHeaderPath) {
+    HeaderMappingTable[OrignalHeaderPath] = MappingHeaderPath;
+  };
+  const HeaderMap &getHeaderMappingTable() { return HeaderMappingTable; };
+
+private:
+  /// A string-to-string map saving the mapping relationship.
+  HeaderMap HeaderMappingTable;
+};
+
+} // namespace find_all_symbols
+} // namespace clang
+
+#endif // LLVM_CLANG_TOOLS_EXTRA_FIND_ALL_SYMBOLS_HEADER_MAP_COLLECTOR_H

clang-tools-extra/include-fixer/find-all-symbols/PragmaCommentHandler.cpp

@@ -0,0 +1,37 @@
+//===-- PragmaCommentHandler.cpp - find all symbols -----------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PragmaCommentHandler.h"
+#include "FindAllSymbols.h"
+#include "HeaderMapCollector.h"
+#include "clang/Lex/Preprocessor.h"
+#include "llvm/Support/Regex.h"
+
+namespace clang {
+namespace find_all_symbols {
+namespace {
+const char IWYUPragma[] = "// IWYU pragma: private, include ";
+} // namespace
+
+bool PragmaCommentHandler::HandleComment(Preprocessor &PP, SourceRange Range) {
+  StringRef Text =
+      Lexer::getSourceText(CharSourceRange::getCharRange(Range),
+                           PP.getSourceManager(), PP.getLangOpts());
+  size_t Pos = Text.find(IWYUPragma);
+  if (Pos == StringRef::npos)
+    return false;
+  StringRef RemappingFilePath = Text.substr(Pos + std::strlen(IWYUPragma));
+  Collector->addHeaderMapping(
+      PP.getSourceManager().getFilename(Range.getBegin()),
+      RemappingFilePath.trim("\"<>"));
+  return false;
+}
+
+} // namespace find_all_symbols
+} // namespace clang

clang-tools-extra/include-fixer/find-all-symbols/PragmaCommentHandler.h

@@ -0,0 +1,41 @@
+//===-- PragmaCommentHandler.h - find all symbols----------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TOOLS_EXTRA_FIND_ALL_SYMBOLS_PRAGMA_COMMENT_HANDLER_H
+#define LLVM_CLANG_TOOLS_EXTRA_FIND_ALL_SYMBOLS_PRAGMA_COMMENT_HANDLER_H
+
+#include "clang/Basic/SourceLocation.h"
+#include "clang/Lex/Preprocessor.h"
+#include <map>
+
+namespace clang {
+namespace find_all_symbols {
+
+class HeaderMapCollector;
+
+/// \brief PragmaCommentHandler parses pragma comment on include files to
+/// determine when we should include a different header from the header that
+/// directly defines a symbol.
+///
+/// Currently it only supports IWYU private pragma:
+/// https://github.com/include-what-you-use/include-what-you-use/blob/master/docs/IWYUPragmas.md#iwyu-pragma-private
+class PragmaCommentHandler : public clang::CommentHandler {
+public:
+  PragmaCommentHandler(HeaderMapCollector *Collector) : Collector(Collector) {}
+
+  bool HandleComment(Preprocessor &PP, SourceRange Range) override;
+
+private:
+  HeaderMapCollector *const Collector;
+};
+
+} // namespace find_all_symbols
+} // namespace clang
+
+#endif // LLVM_CLANG_TOOLS_EXTRA_FIND_ALL_SYMBOLS_PRAGMA_COMMENT_HANDLER_H

clang-tools-extra/include-fixer/find-all-symbols/tool/FindAllSymbolsMain.cpp

@@ -8,8 +8,14 @@
 //===----------------------------------------------------------------------===//
 
 #include "FindAllSymbols.h"
+#include "HeaderMapCollector.h"
+#include "PragmaCommentHandler.h"
 #include "SymbolInfo.h"
 #include "clang/ASTMatchers/ASTMatchFinder.h"
+#include "clang/ASTMatchers/ASTMatchers.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "clang/Frontend/FrontendActions.h"
+#include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/CommonOptionsParser.h"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -82,6 +88,31 @@ class YamlReporter
   std::map<std::string, std::set<SymbolInfo>> Symbols;
 };
 
+class FindAllSymbolsAction : public clang::ASTFrontendAction {
+public:
+  FindAllSymbolsAction()
+      : Reporter(), MatchFinder(), Collector(), Handler(&Collector),
+        Matcher(&Reporter, &Collector) {
+    Matcher.registerMatchers(&MatchFinder);
+  }
+
+  std::unique_ptr<clang::ASTConsumer>
+  CreateASTConsumer(clang::CompilerInstance &Compiler,
+                    StringRef InFile) override {
+    Compiler.getPreprocessor().addCommentHandler(&Handler);
+    return MatchFinder.newASTConsumer();
+  }
+
+  void EndSourceFileAction() override { Reporter.Write(OutputDir); }
+
+private:
+  YamlReporter Reporter;
+  clang::ast_matchers::MatchFinder MatchFinder;
+  HeaderMapCollector Collector;
+  PragmaCommentHandler Handler;
+  FindAllSymbols Matcher;
+};
+
 bool Merge(llvm::StringRef MergeDir, llvm::StringRef OutputFile) {
   std::error_code EC;
   std::set<SymbolInfo> UniqueSymbols;
@@ -141,12 +172,8 @@ int main(int argc, const char **argv) {
     clang::find_all_symbols::Merge(MergeDir, sources[0]);
     return 0;
   }
-
-  clang::find_all_symbols::YamlReporter Reporter;
-  clang::find_all_symbols::FindAllSymbols Matcher(&Reporter);
-  clang::ast_matchers::MatchFinder MatchFinder;
-  Matcher.registerMatchers(&MatchFinder);
-  Tool.run(newFrontendActionFactory(&MatchFinder).get());
-  Reporter.Write(OutputDir);
+  Tool.run(
+      newFrontendActionFactory<clang::find_all_symbols::FindAllSymbolsAction>()
+          .get());
   return 0;
 }

clang-tools-extra/unittests/include-fixer/find-all-symbols/FindAllSymbolsTests.cpp

@@ -8,11 +8,14 @@
 //===----------------------------------------------------------------------===//
 
 #include "FindAllSymbols.h"
+#include "HeaderMapCollector.h"
+#include "PragmaCommentHandler.h"
 #include "SymbolInfo.h"
 #include "clang/ASTMatchers/ASTMatchFinder.h"
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/FileSystemOptions.h"
 #include "clang/Basic/VirtualFileSystem.h"
+#include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/PCHContainerOperations.h"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
@@ -50,25 +53,58 @@ class MockReporter
   std::vector<SymbolInfo> Symbols;
 };
 
+class TestFindAllSymbolsAction : public clang::ASTFrontendAction {
+public:
+  TestFindAllSymbolsAction(FindAllSymbols::ResultReporter *Reporter)
+      : MatchFinder(), Collector(), Handler(&Collector),
+        Matcher(Reporter, &Collector) {
+    Matcher.registerMatchers(&MatchFinder);
+  }
+
+  std::unique_ptr<clang::ASTConsumer>
+  CreateASTConsumer(clang::CompilerInstance &Compiler,
+                    StringRef InFile) override {
+    Compiler.getPreprocessor().addCommentHandler(&Handler);
+    return MatchFinder.newASTConsumer();
+  }
+
+private:
+  ast_matchers::MatchFinder MatchFinder;
+  HeaderMapCollector Collector;
+  PragmaCommentHandler Handler;
+  FindAllSymbols Matcher;
+};
+
+class TestFindAllSymbolsActionFactory
+    : public clang::tooling::FrontendActionFactory {
+public:
+  TestFindAllSymbolsActionFactory(MockReporter *Reporter)
+      : Reporter(Reporter) {}
+  clang::FrontendAction *create() override {
+    return new TestFindAllSymbolsAction(Reporter);
+  }
+
+private:
+  MockReporter *const Reporter;
+};
+
 class FindAllSymbolsTest : public ::testing::Test {
 public:
   bool hasSymbol(const SymbolInfo &Symbol) {
     return Reporter.hasSymbol(Symbol);
   }
 
   bool runFindAllSymbols(StringRef Code) {
-    FindAllSymbols matcher(&Reporter);
-    clang::ast_matchers::MatchFinder MatchFinder;
-    matcher.registerMatchers(&MatchFinder);
-
     llvm::IntrusiveRefCntPtr<vfs::InMemoryFileSystem> InMemoryFileSystem(
         new vfs::InMemoryFileSystem);
     llvm::IntrusiveRefCntPtr<FileManager> Files(
         new FileManager(FileSystemOptions(), InMemoryFileSystem));
 
     std::string FileName = "symbol.cc";
-    std::unique_ptr<clang::tooling::FrontendActionFactory> Factory =
-        clang::tooling::newFrontendActionFactory(&MatchFinder);
+
+    std::unique_ptr<clang::tooling::FrontendActionFactory> Factory(
+        new TestFindAllSymbolsActionFactory(&Reporter));
+
     tooling::ToolInvocation Invocation(
         {std::string("find_all_symbols"), std::string("-fsyntax-only"),
          std::string("-std=c++11"), FileName},
@@ -329,5 +365,18 @@ TEST_F(FindAllSymbolsTest, EnumTest) {
   EXPECT_FALSE(hasSymbol(Symbol));
 }
 
+TEST_F(FindAllSymbolsTest, IWYUPrivatePragmaTest) {
+  static const char Code[] = R"(
+    // IWYU pragma: private, include "bar.h"
+    struct Bar {
+    };
+  )";
+  runFindAllSymbols(Code);
+
+  SymbolInfo Symbol =
+      SymbolInfo("Bar", SymbolInfo::SymbolKind::Class, "bar.h", 3, {});
+  EXPECT_TRUE(hasSymbol(Symbol));
+}
+
 } // namespace find_all_symbols
 } // namespace clang