[include-fixer] Create a mode in vim integration to show multiple potential headers.

[hokein]

Summary

Some changes in the patch:

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

[hokein]

Oh, sorry, I miss two separate commits here. This patch should not be ready for review. I need to rebase it after commit D20581.

[hokein]

Rebase

[hokein]

Fix a nit.

[ioeric]

include-fixer/IncludeFixer.cpp

239–243

This doesn't seem to be the right name. It does more than minimizing headers.

include-fixer/tool/ClangIncludeFixer.cpp

120

These comments seem redundant since it is already clear what those parameters are from variables' names.

128

This should only be for vim-mode or STDINMode I think. For normal mode, we should apply replacements on the rewriter.

207

Same here.

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

30

I think we can simply have one mode where a header is inserted if there is only one result, and a dialog is shown only if there are multiple headers for user to choose from.

30

Also remove semicolons here and below...

94

I think we should also update the buffer with the diff method below. It doesn't make sense to clear the buffer if we are just inserting one line.

[hokein]

Update and address comments.

[bkramer]

Can you add some lit tests for the various command line modes clang-include-fixer has now. We can't reasonably test the vim integration but we can tests the bits it's composed of.

[bkramer]

include-fixer/IncludeFixer.h

52–53

maybe 'about the symbol being queried'?

73

'Header' looks like it should be a StringRef too.

include-fixer/IncludeFixerContext.h

20

I think some doxygen comments would be in order for this struct and its members.

include-fixer/tool/ClangIncludeFixer.cpp

61

typo: Ouput

[ioeric]

include-fixer/IncludeFixer.cpp

241

I think function name should start with lower case in llvm.

include-fixer/IncludeFixer.h

71

Function name should start with lower case.

include-fixer/tool/ClangIncludeFixer.cpp

221

I'd put this before creating replacements and right after you get the context since those replacements are not used here. And the "Added #include..." message above doesn't make sense in this mode.

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

48

Maybe handle cases where to_eval is not a valid index?

88

If there is only one candidate, it doesn't make sense to ask users to choose it IMO. We can simply insert the header in this case.

And I think we should make users pick the correct header if there are multiple candidates by default. What do you think, Ben? @bkramer

[bkramer]

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

88

I agree, we should ask by default when there is more than one suggested header.

[hokein]

Address comments.

[hokein]

In D20621#439447, @bkramer wrote:

Can you add some lit tests for the various command line modes clang-include-fixer has now. We can't reasonably test the vim integration but we can tests the bits it's composed of.

Done.

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

48

vim can handle this for us.

[hokein]

Remove unneeded headers.

[klimek]

include-fixer/IncludeFixer.cpp

397

Do we want UINT_MAX? I find the wording in the standard to be too involved for me to easily understand that this is both well-defined and what we want and portable.

include-fixer/IncludeFixer.h

63–73

1. This only needs one style, so I think we should pass it in instead of FilePath and FallbackStyle
2. From the docs it seems like FirstIncludeOffset should always be INT_MAX? Can we just leave out that parameter and make the function do the right thing?
3. Generally, I'd order parameters by "importance", that is, the ones that are core to the functionality go first (C++ kinda implies this by only allowing later parameters being defaultable); so basically, I'd make this:
   
   .. insertHeader(StringRef Code, StringRef Header, const FormatStyle& Style);

[hokein]

Refactor createReplacementsForHeaders.

[ioeric]

include-fixer/IncludeFixer.h

80

I don't see why we'd want Style to be optional.

include-fixer/tool/ClangIncludeFixer.cpp

110

We might still want to use clang::format::getStyle("file", FilePath, FallbackStyle);, which will use format style defined in clang-format config file in the path, if such config file exists. Style option is just a fallback style.

[hokein]

Use format::getStyle to get clang-format style.

[hokein]

include-fixer/IncludeFixer.h

80

Using a default argument in Style can simplify the code in some cases, for example in the Unittest. @klimek might have idea?

[bkramer]

include-fixer/IncludeFixer.h

71

I think you can drop the part about 'uninitialization' from this comment.

79

This doesn't look clang-formatted :p

[hokein]

Fix code style.

[bkramer]

LG. Can't wait to use it myself :)

[hokein]

Update a out-of-date comment.

[klimek]

include-fixer/IncludeFixer.h

74

This is still missing documentation on what the -1U special case means.

[hokein]

Add -1U comment back.

[klimek]

lg

[hokein]

In D20621#444050, @bkramer wrote:

LG. Can't wait to use it myself :)

Currently, the header is only inserted at the first line of the file because we don't output the FirstIncludeOffset to py script. A follow-up patch will come soon.

include-fixer/IncludeFixer.h

 //===-- IncludeFixer.h - Include inserter -----------------------*- 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_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>
 #include <vector>
 
 namespace clang {
 
 class CompilerInvocation;
 class DiagnosticConsumer;
 class FileManager;
 class PCHContainerOperations;
 
 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 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;
 
   bool
   runInvocation(clang::CompilerInvocation *Invocation,
                 clang::FileManager *Files,
                 std::shared_ptr<clang::PCHContainerOperations> PCHContainerOps,
                 clang::DiagnosticConsumer *Diagnostics) override;
 
 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;
 
   /// The fallback format style for formatting after insertion if no
   /// clang-format config file was found.
   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.
+/// -1U(UINT_MAX) indicates uninitialization, which will insert the header at
+/// first line. If there is no #include block in the Code, 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
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_INCLUDE_FIXER_INCLUDEFIXER_H

include-fixer/IncludeFixer.cpp

void NextInclude(clang::SourceManager *SourceManager,
                    clang::SourceLocation EndLocation) {
     unsigned Offset = SourceManager->getFileOffset(HashLocation);
     if (FirstIncludeOffset == -1U)
       FirstIncludeOffset = Offset;
   }
 
   /// 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;
 
     // Get the FileEntry for the include.
     StringRef StrippedInclude = Include.trim("\"<>");
     const FileEntry *Entry =
         SourceManager.getFileManager().getFile(StrippedInclude);
 
     // If the file doesn't exist return the path from the database.
     // FIXME: This should never happen.
     if (!Entry)
       return Include;
 
     bool IsSystem;
     std::string Suggestion =
         HeaderSearch.suggestPathToFileForDiagnostics(Entry, &IsSystem);
 
     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));
-
-    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;
+    FixerContext.SymbolIdentifer = QuerySymbol;
+    FixerContext.FirstIncludeOffset = FirstIncludeOffset;
+    for (const auto &Header : SymbolQueryResults)
+      FixerContext.Headers.push_back(
+          minimizeInclude(Header, SourceManager, HeaderSearch));
+
+    return FixerContext;
   }
 
   /// Sets the location at the very top of the file.
   void setFileBegin(clang::SourceLocation Location) { FileBegin = Location; }
 
 private:
   /// Query the database for a given identifier.
   bool query(StringRef Query, SourceLocation Loc) {
     assert(!Query.empty() && "Empty query!");
 
     // Skip other identifers once we have discovered an identfier successfully.
     if (!SymbolQueryResults.empty())
       return false;
 
     DEBUG(llvm::dbgs() << "Looking up '" << Query << "' at ");
     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();
   }
 
   /// The client to use to find cross-references.
   SymbolIndexManager &SymbolIndexMgr;
 
   /// The absolute path to the file being processed.
   std::string Filename;
 
   /// The location of the beginning of the tracked file.
   clang::SourceLocation FileBegin;
 
   /// The offset of the last include in the original source file. This will
   /// be used as the insertion point for new include directives.
   unsigned FirstIncludeOffset = -1U;
 
   /// The fallback format style for formatting after insertion if there is no
   /// 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;
 
   /// Whether we should use the smallest possible include path.
   bool MinimizeIncludePaths = true;
 };
 
if (!IncludeFile || TrackedFile != SourceFile)
     return;
   EnclosingPass->NextInclude(SourceManager, IncludeFile->getName(), FileName,
                              IsAngled, HashLocation, FileNameRange.getEnd());
 }
 
 } // 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;
 
 bool IncludeFixerActionFactory::runInvocation(
     clang::CompilerInvocation *Invocation, clang::FileManager *Files,
     std::shared_ptr<clang::PCHContainerOperations> PCHContainerOps,
     clang::DiagnosticConsumer *Diagnostics) {
   assert(Invocation->getFrontendOpts().Inputs.size() == 1);
bool IncludeFixerActionFactory::runInvocation(
   // fatal. A missing #include can cause thousands of errors.
   Compiler.getDiagnostics().setErrorLimit(0);
 
   // Run the parser, gather missing includes.
   auto ScopedToolAction = llvm::make_unique<Action>(
       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

include-fixer/IncludeFixerContext.h

+//===-- 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

include-fixer/tool/ClangIncludeFixer.cpp

 //===-- ClangIncludeFixer.cpp - Standalone include fixer ------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace clang;
 using namespace llvm;
cl::opt<bool>
     STDINMode("stdin",
               cl::desc("Override source file's content (in the overlaying\n"
                        "virtual file system) with input from <stdin> and run\n"
                        "the tool on the new content with the compilation\n"
                        "options of the source file. This mode is currently\n"
                        "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 "
                    "headers if there is no clang-format config file found."),
           cl::init("llvm"), cl::cat(IncludeFixerCategory));
 
 int includeFixerMain(int argc, const char **argv) {
   tooling::CommonOptionsParser options(argc, argv, IncludeFixerCategory);
if (STDINMode) {
     }
     Code = std::move(CodeOrErr.get());
     if (Code->getBufferSize() == 0)
       return 0;  // Skip empty files.
 
     tool.mapVirtualFile(options.getSourcePathList().front(), Code->getBuffer());
   }
 
+  format::FormatStyle InsertStyle;
+  format::getPredefinedStyle(Style, format::FormatStyle::LK_Cpp, &InsertStyle);
+
+  if (!InsertHeader.empty()) {
+    if (!STDINMode) {
+      errs() << "Should be running in STDIN mode\n";
+      return 1;
+    }
+
+    StringRef FilePath = options.getSourcePathList().front();
+
+    // 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) {
   case fixed: {
     // Parse input and fill the database with it.
     // <symbol>=<header><, header...>
     // Multiple symbols can be given, separated by semicolons.
     std::map<std::string, std::vector<std::string>> SymbolsMap;
case yaml: {
     }
 
     SymbolIndexMgr->addSymbolIndex(std::move(*DB));
     break;
   }
   }
 
   // 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;
+
+  StringRef FilePath = options.getSourcePathList().front();
+  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);
   DiagnosticsEngine Diagnostics(new DiagnosticIDs, &*DiagOpts);
   TextDiagnosticPrinter DiagnosticPrinter(outs(), &*DiagOpts);
   SourceManager SM(Diagnostics, tool.getFiles());
   Diagnostics.setClient(&DiagnosticPrinter, false);
 
   if (STDINMode) {
     tooling::Replacements Replaces(Replacements.begin(), Replacements.end());

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

 # be run on the current buffer.
 #
 # It operates on the current, potentially unsaved buffer and does not create
 # or save any files. To revert a fix, just undo.
 
 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
 # on the path.
 # Change this to the full path if clang-include-fixer is not on the path.
 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')
   parser.add_argument('-db', default='yaml',
                       help='clang-include-fixer input format.')
   parser.add_argument('-input', default='',
                       help='String to initialize the database.')
   args = parser.parse_args()
 
   # Get the current text.
   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)
-
-  # If successful, replace buffer contents.
-  if stderr:
-    print stderr
-
-  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]]
+  stdout, stderr = execute(command, text)
+  if len(lines) < 2:
+    print "No header is included.\n"
+    return
+
+  # 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 __name__ == '__main__':
   main()

test/include-fixer/commandline_options.cpp

+// 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;

unittests/include-fixer/IncludeFixerTest.cpp

std::vector<SymbolInfo> Symbols = {
                  1, {{SymbolInfo::ContextType::EnumDecl, "Color"},
                      {SymbolInfo::ContextType::Namespace, "b"},
                      {SymbolInfo::ContextType::Namespace, "a"}}),
   };
   auto SymbolIndexMgr = llvm::make_unique<include_fixer::SymbolIndexManager>();
   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);
   return Context.getRewrittenText(ID);
 }
 
 TEST(IncludeFixer, Typo) {
   EXPECT_EQ("#include <string>\n\nstd::string foo;\n",