[include-fixer] Create a mode in vim integration to show multiple pot…

…ential headers.

Summary:
Some changes in the patch:

* Add two commandline flags in clang-include-fixer.
* Introduce a IncludeFixerContext for the queried symbol.
* Pull out CreateReplacementsForHeader.

Reviewers: bkramer

Subscribers: klimek, cfe-commits, ioeric

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

llvm-svn: 271258

clang-tools-extra/include-fixer/IncludeFixer.cpp

@@ -214,7 +214,7 @@ class Action : public clang::ASTFrontendAction,
 
   /// Get the minimal include for a given path.
   std::string minimizeInclude(StringRef Include,
-                              clang::SourceManager &SourceManager,
+                              const clang::SourceManager &SourceManager,
                               clang::HeaderSearch &HeaderSearch) {
     if (!MinimizeIncludePaths)
       return Include;
@@ -236,66 +236,21 @@ class Action : public clang::ASTFrontendAction,
     return IsSystem ? '<' + Suggestion + '>' : '"' + Suggestion + '"';
   }
 
-  /// Insert all headers before the first #include in \p Code and run
-  /// clang-format to sort all headers.
-  /// \return Replacements for inserting and sorting headers.
-  std::vector<clang::tooling::Replacement>
-  CreateReplacementsForHeaders(StringRef Code,
-                               const std::set<std::string> &Headers) {
-    // Create replacements for new headers.
-    clang::tooling::Replacements Insertions;
-    if (FirstIncludeOffset == -1U) {
-      // FIXME: skip header guards.
-      FirstIncludeOffset = 0;
-      // If there is no existing #include, then insert an empty line after new
-      // header block.
-      if (Code.front() != '\n')
-        Insertions.insert(
-            clang::tooling::Replacement(Filename, FirstIncludeOffset, 0, "\n"));
-    }
-    // Keep inserting new headers before the first header.
-    for (StringRef Header : Headers) {
-      std::string Text = "#include " + Header.str() + "\n";
-      Insertions.insert(
-          clang::tooling::Replacement(Filename, FirstIncludeOffset, 0, Text));
-    }
-    DEBUG({
-      llvm::dbgs() << "Header insertions:\n";
-      for (const auto &R : Insertions)
-        llvm::dbgs() << R.toString() << '\n';
-    });
-
-    clang::format::FormatStyle Style =
-        clang::format::getStyle("file", Filename, FallbackStyle);
-    clang::tooling::Replacements Replaces =
-        formatReplacements(Code, Insertions, Style);
-    // FIXME: remove this when `clang::tooling::Replacements` is implemented as
-    // `std::vector<clang::tooling::Replacement>`.
-    std::vector<clang::tooling::Replacement> Results;
-    std::copy(Replaces.begin(), Replaces.end(), std::back_inserter(Results));
-    return Results;
-  }
-
-  /// Generate replacements for the suggested includes.
-  /// \return true if changes will be made, false otherwise.
-  bool Rewrite(clang::SourceManager &SourceManager,
-               clang::HeaderSearch &HeaderSearch,
-               std::set<std::string> &Headers,
-               std::vector<clang::tooling::Replacement> &Replacements) {
+  /// Get the include fixer context for the queried symbol.
+  IncludeFixerContext
+  getIncludeFixerContext(const clang::SourceManager &SourceManager,
+                         clang::HeaderSearch &HeaderSearch) {
+    IncludeFixerContext FixerContext;
     if (SymbolQueryResults.empty())
-      return false;
+      return FixerContext;
 
-    // FIXME: Rank the results and pick the best one instead of the first one.
-    const auto &ToTry = SymbolQueryResults.front();
-    Headers.insert(minimizeInclude(ToTry, SourceManager, HeaderSearch));
+    FixerContext.SymbolIdentifer = QuerySymbol;
+    FixerContext.FirstIncludeOffset = FirstIncludeOffset;
+    for (const auto &Header : SymbolQueryResults)
+      FixerContext.Headers.push_back(
+          minimizeInclude(Header, SourceManager, HeaderSearch));
 
-    StringRef Code = SourceManager.getBufferData(SourceManager.getMainFileID());
-    Replacements = CreateReplacementsForHeaders(Code, Headers);
-
-    // We currently abort after the first inserted include. The more
-    // includes we have the less safe this becomes due to error recovery
-    // changing the results.
-    return true;
+    return FixerContext;
   }
 
   /// Sets the location at the very top of the file.
@@ -314,6 +269,7 @@ class Action : public clang::ASTFrontendAction,
     DEBUG(Loc.print(llvm::dbgs(), getCompilerInstance().getSourceManager()));
     DEBUG(llvm::dbgs() << " ...");
 
+    QuerySymbol = Query.str();
     SymbolQueryResults = SymbolIndexMgr.search(Query);
     DEBUG(llvm::dbgs() << SymbolQueryResults.size() << " replies\n");
     return !SymbolQueryResults.empty();
@@ -336,6 +292,9 @@ class Action : public clang::ASTFrontendAction,
   /// clang-format config file found.
   std::string FallbackStyle;
 
+  /// The symbol being queried.
+  std::string QuerySymbol;
+
   /// The query results of an identifier. We only include the first discovered
   /// identifier to avoid getting caught in results from error recovery.
   std::vector<std::string> SymbolQueryResults;
@@ -385,12 +344,10 @@ void PreprocessorHooks::InclusionDirective(
 } // namespace
 
 IncludeFixerActionFactory::IncludeFixerActionFactory(
-    SymbolIndexManager &SymbolIndexMgr, std::set<std::string> &Headers,
-    std::vector<clang::tooling::Replacement> &Replacements, StringRef StyleName,
-    bool MinimizeIncludePaths)
-    : SymbolIndexMgr(SymbolIndexMgr), Headers(Headers),
-      Replacements(Replacements), MinimizeIncludePaths(MinimizeIncludePaths),
-      FallbackStyle(StyleName) {}
+    SymbolIndexManager &SymbolIndexMgr, IncludeFixerContext &Context,
+    StringRef StyleName, bool MinimizeIncludePaths)
+    : SymbolIndexMgr(SymbolIndexMgr), Context(Context),
+      MinimizeIncludePaths(MinimizeIncludePaths), FallbackStyle(StyleName) {}
 
 IncludeFixerActionFactory::~IncludeFixerActionFactory() = default;
 
@@ -420,16 +377,51 @@ bool IncludeFixerActionFactory::runInvocation(
       SymbolIndexMgr, FallbackStyle, MinimizeIncludePaths);
   Compiler.ExecuteAction(*ScopedToolAction);
 
-  // Generate replacements.
-  ScopedToolAction->Rewrite(Compiler.getSourceManager(),
-                            Compiler.getPreprocessor().getHeaderSearchInfo(),
-                            Headers, Replacements);
+  Context = ScopedToolAction->getIncludeFixerContext(
+      Compiler.getSourceManager(),
+      Compiler.getPreprocessor().getHeaderSearchInfo());
 
   // Technically this should only return true if we're sure that we have a
   // parseable file. We don't know that though. Only inform users of fatal
   // errors.
   return !Compiler.getDiagnostics().hasFatalErrorOccurred();
 }
 
+std::vector<clang::tooling::Replacement>
+createInsertHeaderReplacements(StringRef Code, StringRef FilePath,
+                               StringRef Header, unsigned FirstIncludeOffset,
+                               const clang::format::FormatStyle &Style) {
+  if (Header.empty())
+    return {};
+  // Create replacements for new headers.
+  clang::tooling::Replacements Insertions;
+  if (FirstIncludeOffset == -1U) {
+    // FIXME: skip header guards.
+    FirstIncludeOffset = 0;
+    // If there is no existing #include, then insert an empty line after new
+    // header block.
+    if (Code.front() != '\n')
+      Insertions.insert(
+          clang::tooling::Replacement(FilePath, FirstIncludeOffset, 0, "\n"));
+  }
+  // Keep inserting new headers before the first header.
+  std::string Text = "#include " + Header.str() + "\n";
+  Insertions.insert(
+      clang::tooling::Replacement(FilePath, FirstIncludeOffset, 0, Text));
+  DEBUG({
+    llvm::dbgs() << "Header insertions:\n";
+    for (const auto &R : Insertions)
+      llvm::dbgs() << R.toString() << '\n';
+  });
+
+  clang::tooling::Replacements Replaces =
+      formatReplacements(Code, Insertions, Style);
+  // FIXME: remove this when `clang::tooling::Replacements` is implemented as
+  // `std::vector<clang::tooling::Replacement>`.
+  std::vector<clang::tooling::Replacement> Results;
+  std::copy(Replaces.begin(), Replaces.end(), std::back_inserter(Results));
+  return Results;
+}
+
 } // namespace include_fixer
 } // namespace clang

clang-tools-extra/include-fixer/IncludeFixer.h

@@ -10,7 +10,9 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_INCLUDE_FIXER_INCLUDEFIXER_H
 #define LLVM_CLANG_TOOLS_EXTRA_INCLUDE_FIXER_INCLUDEFIXER_H
 
+#include "IncludeFixerContext.h"
 #include "SymbolIndexManager.h"
+#include "clang/Format/Format.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "clang/Tooling/Tooling.h"
 #include <memory>
@@ -28,13 +30,12 @@ namespace include_fixer {
 class IncludeFixerActionFactory : public clang::tooling::ToolAction {
 public:
   /// \param SymbolIndexMgr A source for matching symbols to header files.
-  /// \param Replacements Storage for the output of the fixer.
+  /// \param Context A context for the symbol being queried.
   /// \param StyleName Fallback style for reformatting.
   /// \param MinimizeIncludePaths whether inserted include paths are optimized.
-  IncludeFixerActionFactory(
-      SymbolIndexManager &SymbolIndexMgr, std::set<std::string> &Headers,
-      std::vector<clang::tooling::Replacement> &Replacements,
-      StringRef StyleName, bool MinimizeIncludePaths = true);
+  IncludeFixerActionFactory(SymbolIndexManager &SymbolIndexMgr,
+                            IncludeFixerContext &Context, StringRef StyleName,
+                            bool MinimizeIncludePaths = true);
 
   ~IncludeFixerActionFactory() override;
 
@@ -48,11 +49,8 @@ class IncludeFixerActionFactory : public clang::tooling::ToolAction {
   /// The client to use to find cross-references.
   SymbolIndexManager &SymbolIndexMgr;
 
-  /// Headers to be added.
-  std::set<std::string> &Headers;
-
-  /// Replacements are written here.
-  std::vector<clang::tooling::Replacement> &Replacements;
+  /// The context that contains all information about the symbol being queried.
+  IncludeFixerContext &Context;
 
   /// Whether inserted include paths should be optimized.
   bool MinimizeIncludePaths;
@@ -62,6 +60,25 @@ class IncludeFixerActionFactory : public clang::tooling::ToolAction {
   std::string FallbackStyle;
 };
 
+/// Create replacements for the header being inserted. The replacements will
+/// insert a header before the first #include in \p Code, and sort all headers
+/// with the given clang-format style.
+///
+/// \param Code The source code.
+/// \param FilePath The source file path.
+/// \param Header The header being inserted.
+/// \param FirstIncludeOffset The insertion point for new include directives.
+/// The default value -1U means inserting the header at the first line, and if
+/// there is no #include block, it will create a header block by inserting a
+/// newline.
+/// \param Style clang-format style being used.
+///
+/// \return Replacements for inserting and sorting headers.
+std::vector<clang::tooling::Replacement> createInsertHeaderReplacements(
+    StringRef Code, StringRef FilePath, StringRef Header,
+    unsigned FirstIncludeOffset = -1U,
+    const clang::format::FormatStyle &Style = clang::format::getLLVMStyle());
+
 } // namespace include_fixer
 } // namespace clang
 

clang-tools-extra/include-fixer/IncludeFixerContext.h

@@ -0,0 +1,32 @@
+//===-- IncludeFixerContext.h - Include fixer context -----------*- 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_INCLUDE_FIXER_INCLUDEFIXERCONTEXT_H
+#define LLVM_CLANG_TOOLS_EXTRA_INCLUDE_FIXER_INCLUDEFIXERCONTEXT_H
+
+#include <string>
+#include <vector>
+
+namespace clang {
+namespace include_fixer {
+
+/// \brief A context for the symbol being queried.
+struct IncludeFixerContext {
+  /// \brief The symbol name.
+  std::string SymbolIdentifer;
+  /// \brief The headers which have SymbolIdentifier definitions.
+  std::vector<std::string> Headers;
+  /// \brief The insertion point for new include header.
+  unsigned FirstIncludeOffset;
+};
+
+} // namespace include_fixer
+} // namespace clang
+
+#endif // LLVM_CLANG_TOOLS_EXTRA_INCLUDE_FIXER_INCLUDEFIXERCONTEXT_H

clang-tools-extra/include-fixer/tool/ClangIncludeFixer.cpp

@@ -9,8 +9,10 @@
 
 #include "InMemorySymbolIndex.h"
 #include "IncludeFixer.h"
+#include "IncludeFixerContext.h"
 #include "SymbolIndexManager.h"
 #include "YamlSymbolIndex.h"
+#include "clang/Format/Format.h"
 #include "clang/Frontend/TextDiagnosticPrinter.h"
 #include "clang/Rewrite/Core/Rewriter.h"
 #include "clang/Tooling/CommonOptionsParser.h"
@@ -56,6 +58,23 @@ cl::opt<bool>
                        "used for editor integration."),
               cl::init(false), cl::cat(IncludeFixerCategory));
 
+cl::opt<bool> OutputHeaders(
+    "output-headers",
+    cl::desc("Output the symbol being quired and all its relevant headers.\n"
+             "The first line is the symbol name; The other lines\n"
+             "are the headers: \n"
+             "   b::foo\n"
+             "   path/to/foo_a.h\n"
+             "   path/to/foo_b.h\n"),
+    cl::init(false), cl::cat(IncludeFixerCategory));
+
+cl::opt<std::string> InsertHeader(
+    "insert-header",
+    cl::desc("Insert a specific header. This should run with STDIN mode.\n"
+             "The result is written to stdout. It is currently used for\n"
+             "editor integration."),
+    cl::init(""), cl::cat(IncludeFixerCategory));
+
 cl::opt<std::string>
     Style("style",
           cl::desc("Fallback style for reformatting after inserting new "
@@ -87,6 +106,27 @@ int includeFixerMain(int argc, const char **argv) {
     tool.mapVirtualFile(options.getSourcePathList().front(), Code->getBuffer());
   }
 
+  StringRef FilePath = options.getSourcePathList().front();
+  format::FormatStyle InsertStyle = format::getStyle("file", FilePath, Style);
+
+  if (!InsertHeader.empty()) {
+    if (!STDINMode) {
+      errs() << "Should be running in STDIN mode\n";
+      return 1;
+    }
+
+    // FIXME: Insert the header in right FirstIncludeOffset.
+    std::vector<tooling::Replacement> Replacements =
+        clang::include_fixer::createInsertHeaderReplacements(
+            Code->getBuffer(), FilePath, InsertHeader,
+            /*FirstIncludeOffset=*/0, InsertStyle);
+    tooling::Replacements Replaces(Replacements.begin(), Replacements.end());
+    std::string ChangedCode =
+        tooling::applyAllReplacements(Code->getBuffer(), Replaces);
+    llvm::outs() << ChangedCode;
+    return 0;
+  }
+
   // Set up data source.
   auto SymbolIndexMgr = llvm::make_unique<include_fixer::SymbolIndexManager>();
   switch (DatabaseFormat) {
@@ -139,20 +179,41 @@ int includeFixerMain(int argc, const char **argv) {
   }
 
   // Now run our tool.
-  std::set<std::string> Headers;  // Headers to be added.
-  std::vector<tooling::Replacement> Replacements;
-  include_fixer::IncludeFixerActionFactory Factory(
-      *SymbolIndexMgr, Headers, Replacements, Style, MinimizeIncludePaths);
+  include_fixer::IncludeFixerContext Context;
+  include_fixer::IncludeFixerActionFactory Factory(*SymbolIndexMgr, Context,
+                                                   Style, MinimizeIncludePaths);
 
   if (tool.run(&Factory) != 0) {
     llvm::errs()
         << "Clang died with a fatal error! (incorrect include paths?)\n";
     return 1;
   }
 
+  if (OutputHeaders) {
+    // FIXME: Output IncludeFixerContext as YAML.
+    llvm::outs() << Context.SymbolIdentifer << "\n";
+    for (const auto &Header : Context.Headers)
+      llvm::outs() << Header << "\n";
+    return 0;
+  }
+
+  if (Context.Headers.empty())
+    return 0;
+
+  auto Buffer = llvm::MemoryBuffer::getFile(FilePath);
+  if (!Buffer) {
+    errs() << "Couldn't open file: " << FilePath;
+    return 1;
+  }
+
+  // FIXME: Rank the results and pick the best one instead of the first one.
+  std::vector<tooling::Replacement> Replacements =
+      clang::include_fixer::createInsertHeaderReplacements(
+          /*Code=*/Buffer.get()->getBuffer(), FilePath, Context.Headers.front(),
+          Context.FirstIncludeOffset, InsertStyle);
+
   if (!Quiet)
-    for (const auto &Header : Headers)
-      llvm::errs() << "Added #include " << Header;
+    llvm::errs() << "Added #include" << Context.Headers.front();
 
   // Set up a new source manager for applying the resulting replacements.
   IntrusiveRefCntPtr<DiagnosticOptions> DiagOpts(new DiagnosticOptions);