[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

241

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

include-fixer/tool/ClangIncludeFixer.cpp

121

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

129

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

208

Same here.

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

31

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.

31

Also remove semicolons here and below...

78

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

51–52

maybe 'about the symbol being queried'?

72

'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

62

typo: Ouput

[ioeric]

include-fixer/IncludeFixer.cpp

243

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

include-fixer/IncludeFixer.h

70

Function name should start with lower case.

include-fixer/tool/ClangIncludeFixer.cpp

217

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

49

Maybe handle cases where to_eval is not a valid index?

72

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

72

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

49

vim can handle this for us.

[hokein]

Remove unneeded headers.

[klimek]

include-fixer/IncludeFixer.cpp

398

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

62–72

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

79

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

include-fixer/tool/ClangIncludeFixer.cpp

111

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

79

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

70

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

78

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

73

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/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 &Headers, 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 queried symbol.
+  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;
 };
 
+/// Insert a header before the first #include in \p Code and run
+/// clang-format to sort all headers.
+/// \param Code The source code.
+/// \param FilePath The source file path.
+/// \param StyleName Fallback style for reformatting.
+/// \param FirstIncludeOffset The insertion point for new include directives.
+/// \param Header The header being inserted.
+/// \return Replacements for inserting and sorting headers.
+std::vector<clang::tooling::Replacement> CreateReplacementsForHeader(
+    StringRef Code, StringRef FilePath, StringRef FallbackStyle,
+    unsigned FirstIncludeOffset, const std::string &Header);
+
 } // namespace include_fixer
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_INCLUDE_FIXER_INCLUDEFIXER_H

include-fixer/IncludeFixer.cpp

std::string minimizeInclude(StringRef 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,
+  bool MinimizeAllIncludeHeaders(clang::SourceManager &SourceManager,
                                  clang::HeaderSearch &HeaderSearch,
-               std::set<std::string> &Headers,
-               std::vector<clang::tooling::Replacement> &Replacements) {
+                                 IncludeFixerContext &Context) {
     if (SymbolQueryResults.empty())
       return false;
 
-    // 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.
+    Context.SymbolIdentifer = QuerySymbol;
+    Context.FirstIncludeOffset = FirstIncludeOffset;
+    for (const auto &Header : SymbolQueryResults)
+      Context.Headers.push_back(
+          minimizeInclude(Header, SourceManager, HeaderSearch));
+
     return true;
   }
 
   /// 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(
   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);
+  ScopedToolAction->MinimizeAllIncludeHeaders(
+      Compiler.getSourceManager(),
+      Compiler.getPreprocessor().getHeaderSearchInfo(), Context);
 
   // 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> CreateReplacementsForHeader(
+    StringRef Code, StringRef FilePath, StringRef FallbackStyle,
+    unsigned FirstIncludeOffset, const std::string &Header) {
+  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 + "\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::format::FormatStyle Style =
+      clang::format::getStyle("file", FilePath, 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;
+}
+
 } // 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 {
+
+struct IncludeFixerContext {
+  std::string SymbolIdentifer;
+  std::vector<std::string> Headers;
+  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/Frontend/TextDiagnosticPrinter.h"
 #include "clang/Rewrite/Core/Rewriter.h"
 #include "clang/Tooling/CommonOptionsParser.h"
+#include "clang/Tooling/ReplacementsYaml.h"
 #include "clang/Tooling/Tooling.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/YAMLTraits.h"
 
 using namespace clang;
 using namespace llvm;
 
 namespace {
 cl::OptionCategory IncludeFixerCategory("Tool options");
 
 enum DatabaseFormatTy {
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> OuputHeaders(
+    "output-headers",
+    cl::desc("Output the queried symbol and 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 to a file.\n"
+                          " clang-include-fixer -insert-header=\"foo.h\""
+                          " path/to/source.cc"),
+                 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());
   }
 
+  if (!InsertHeader.empty()) {
+    StringRef FilePath = options.getSourcePathList().front();
+    auto Buffer = llvm::MemoryBuffer::getFile(FilePath);
+    if (!Buffer) {
+      errs() << "Couldn't open file: " << FilePath;
+      return 1;
+    }
+
+    // FIXME: Insert the header in the FirstIncludeOffset.
+    std::vector<tooling::Replacement> Replacements =
+        clang::include_fixer::CreateReplacementsForHeader(
+            /*Code=*/Buffer.get()->getBuffer(),
+            /*FilePath=*/FilePath,
+            /*FallbackStyle=*/Style,
+            /*FirstIncludeOffset=*/-1U,
+            /*Header=*/InsertHeader);
+    tooling::Replacements Replaces(Replacements.begin(), Replacements.end());
+    std::string ChangedCode =
+        tooling::applyAllReplacements(Buffer.get()->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 (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::CreateReplacementsForHeader(
+          /*Code=*/Buffer.get()->getBuffer(),
+          /*FilePath=*/FilePath,
+          /*FallbackStyle=*/Style,
+          /*FirstIncludeOffset=*/Context.FirstIncludeOffset,
+          /*Header=*/Context.Headers.front());
+
   if (!Quiet)
-    for (const auto &Header : Headers)
-      llvm::errs() << "Added #include " << Header;
+    llvm::errs() << "Added #include" << Context.Headers.front();
+
+  if (OuputHeaders) {
+    // FIXME: Output all the fields in Context.
+    llvm::outs() << Context.SymbolIdentifer << "\n";
+    for (const auto &Header : Context.Headers)
+      llvm::outs() << Header << "\n";
+    return 0;
+  }
 
   // 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);
 

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

 
 # 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')
 
+choosing_mode = False;
+if vim.eval('exists("g:clang_include_fixer_choosing_mode")') == "1":
+  choosing_mode = vim.eval('g:clang_include_fixer_choosing_mode')
+
+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 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)
 
+  if choosing_mode:
+    # Run command to get all headers.
+    command = [binary, "-output-headers", "-db="+args.db, "-input="+args.input,
+               vim.current.buffer.name]
+    stdout, stderr = execute(command, text)
+    lines = stdout.splitlines()
+    if len(lines) == 0:
+      return
+    symbol = lines[0]
+    choices_message = ""
+    index = 1;
+    for line in lines[1:]:
+      choices_message += "&" + str(index) + line + "\n"
+      index += 1
+    select = ShowDialog("choose a header file for {0}.".format(symbol), choices_message)
+    # Insert a selected header.
+    command = [binary, "-insert-header="+lines[select], vim.current.buffer.name]
+    stdout, stderr = execute(command, text)
+    vim.current.buffer[:] = None;
+    vim.current.buffer.append(stdout.splitlines())
+    print "Added #include {0}.".format(lines[select])
+    return;
+
+
   # Call clang-include-fixer.
   command = [binary, "-stdin", "-db="+args.db, "-input="+args.input,
              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)
 
   # 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]]
 
 if __name__ == '__main__':
   main()

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::CreateReplacementsForHeader(
+        Code, "input.cc", "llvm", FixerContext.FirstIncludeOffset,
+        FixerContext.Headers.front());
+  }
   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",