[include-fixer] Create a mode in vim integration to show multiple potential 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 @@ public:
 
   /// 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 @@ public:
     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 @@ private:
     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 @@ private:
   /// 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,10 +377,9 @@ 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
@@ -431,5 +387,41 @@ bool IncludeFixerActionFactory::runInvocation(
   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 @@ private:
   /// 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 @@ private:
   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,10 +179,9 @@ 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()
@@ -150,9 +189,31 @@ int includeFixerMain(int argc, const char **argv) {
     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);

clang-tools-extra/include-fixer/tool/clang-include-fixer.py

@@ -18,7 +18,6 @@
 import argparse
 import difflib
 import subprocess
-import sys
 import vim
 
 # set g:clang_include_fixer_path to the path to clang-include-fixer if it is not
@@ -28,6 +27,39 @@ binary = 'clang-include-fixer'
 if vim.eval('exists("g:clang_include_fixer_path")') == "1":
   binary = vim.eval('g:clang_include_fixer_path')
 
+maximum_suggested_headers=3
+if vim.eval('exists("g:clang_include_fixer_maximum_suggested_headers")') == "1":
+  maximum_suggested_headers = max(
+      1,
+      vim.eval('g:clang_include_fixer_maximum_suggested_headers'))
+
+
+def ShowDialog(message, choices, default_choice_index=0):
+  to_eval = "confirm('{0}', '{1}', '{2}')".format(message,
+                                                  choices,
+                                                  default_choice_index)
+  return int(vim.eval(to_eval));
+
+
+def execute(command, text):
+  p = subprocess.Popen(command,
+                       stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+                       stdin=subprocess.PIPE)
+  return p.communicate(input=text)
+
+
+def InsertHeaderToVimBuffer(header, text):
+  command = [binary, "-stdin", "-insert-header="+header,
+             vim.current.buffer.name]
+  stdout, stderr = execute(command, text)
+  if stdout:
+    lines = stdout.splitlines()
+    sequence = difflib.SequenceMatcher(None, vim.current.buffer, lines)
+    for op in reversed(sequence.get_opcodes()):
+      if op[0] is not 'equal':
+        vim.current.buffer[op[1]:op[2]] = lines[op[3]:op[4]]
+
+
 def main():
   parser = argparse.ArgumentParser(
       description='Vim integration for clang-include-fixer')
@@ -41,24 +73,36 @@ def main():
   buf = vim.current.buffer
   text = '\n'.join(buf)
 
-  # Call clang-include-fixer.
-  command = [binary, "-stdin", "-db="+args.db, "-input="+args.input,
+  # Run command to get all headers.
+  command = [binary, "-stdin", "-output-headers", "-db="+args.db, "-input="+args.input, "-debug",
              vim.current.buffer.name]
-  p = subprocess.Popen(command,
-                       stdout=subprocess.PIPE, stderr=subprocess.PIPE,
-                       stdin=subprocess.PIPE)
-  stdout, stderr = p.communicate(input=text)
+  stdout, stderr = execute(command, text)
+  lines = stdout.splitlines()
+  if len(lines) < 2:
+    print "No header is included.\n"
+    return
 
-  # If successful, replace buffer contents.
-  if stderr:
-    print stderr
+  # The first line is the symbol name.
+  symbol = lines[0]
+  # If there is only one suggested header, insert it directly.
+  if len(lines) == 2 or maximum_suggested_headers == 1:
+    InsertHeaderToVimBuffer(lines[1], text)
+    print "Added #include {0} for {1}.\n".format(lines[1], symbol)
+    return
+
+  choices_message = ""
+  index = 1;
+  for header in lines[1:1+maximum_suggested_headers]:
+    choices_message += "&" + str(index) + header + "\n"
+    index += 1
+
+  select = ShowDialog("choose a header file for {0}.".format(symbol),
+                      choices_message)
+  # Insert a selected header.
+  InsertHeaderToVimBuffer(lines[select], text)
+  print "Added #include {0} for {1}.\n".format(lines[select], symbol)
+  return;
 
-  if stdout:
-    lines = stdout.splitlines()
-    sequence = difflib.SequenceMatcher(None, vim.current.buffer, lines)
-    for op in reversed(sequence.get_opcodes()):
-      if op[0] is not 'equal':
-        vim.current.buffer[op[1]:op[2]] = lines[op[3]:op[4]]
 
 if __name__ == '__main__':
   main()

clang-tools-extra/test/include-fixer/commandline_options.cpp

@@ -0,0 +1,12 @@
+// REQUIRES: shell
+// RUN: sed -e 's#//.*$##' %s > %t.cpp
+// RUN: clang-include-fixer -db=fixed -input='foo= "foo.h","bar.h"' -output-headers %t.cpp -- | FileCheck %s -check-prefix=CHECK-HEADERS
+// RUN: cat %t.cpp | clang-include-fixer -stdin -insert-header='"foo.h"' %t.cpp | FileCheck %s -check-prefix=CHECK
+//
+// CHECK-HEADERS: "foo.h"
+// CHECK-HEADERS: "bar.h"
+//
+// CHECK: #include "foo.h"
+// CHECK: foo f;
+
+foo f;

clang-tools-extra/unittests/include-fixer/IncludeFixerTest.cpp

@@ -70,11 +70,16 @@ static std::string runIncludeFixer(
   SymbolIndexMgr->addSymbolIndex(
       llvm::make_unique<include_fixer::InMemorySymbolIndex>(Symbols));
 
-  std::set<std::string> Headers;
-  std::vector<clang::tooling::Replacement> Replacements;
-  IncludeFixerActionFactory Factory(*SymbolIndexMgr, Headers, Replacements,
-                                    "llvm");
+  IncludeFixerContext FixerContext;
+  IncludeFixerActionFactory Factory(*SymbolIndexMgr, FixerContext, "llvm");
+
   runOnCode(&Factory, Code, "input.cc", ExtraArgs);
+  std::vector<clang::tooling::Replacement> Replacements;
+  if (!FixerContext.Headers.empty()) {
+    Replacements = clang::include_fixer::createInsertHeaderReplacements(
+        Code, "input.cc", FixerContext.Headers.front(),
+        FixerContext.FirstIncludeOffset);
+  }
   clang::RewriterTestContext Context;
   clang::FileID ID = Context.createInMemoryFile("input.cc", Code);
   clang::tooling::applyAllReplacements(Replacements, Context.Rewrite);