[include-fixer] Add a prototype for a new include fixing tool.

[bkramer]

Summary

The goal of this tool is fairly simple, look up unknown identifiers in a
global database and add the corresponding #include line. It accomplishes
this by hooking into Sema as an ExternalSemaSource and responding to typo
correction callbacks. This means we can see the unknown identifier before
it's being munged by error recovery.

This doesn't work perfectly yet as some typo corrections don't emit
callbacks (delayed typos), but I think this is fixable. We also handle
only one include at a time as this is meant to be run directly from
the editing environment eventually. Adding multiple includes at the same
time is tricky because of error recovery.

This version only has a a dummy database, so all you can do is fixing
missing includes of <string>, but the indexer to build a database will
follow soon.

[klimek]

include-fixer/FixedXrefsDB.h

18

Add '.', remove "intended for testing". I don't think telling people to not use it outside of testing helps a lot. For example, in the compilation database we found that it helps to allow people to pass a fixed-compilation-db at the command line; similar things might make sense here.

20

I'd make it constructible from a map<string, vector<string>>. Especially for tests, I'd want to see the test setup.

include-fixer/IncludeFixer.cpp

117

specification, *,*

132

Can we sort this so the public interface comes first? Also, why is the public interface so large?

179

This is not actually rewriting, but constructing the replacements, right?

184–190

Can we reuse functionality from clang-format here?

include-fixer/IncludeFixer.h

26

Why are we passing ownership?

36

That also seems weird in the interface here.

[bkramer]

• FixedXrefsDB -> FakeXrefsDB
• FakeXrefsDB is now configurable
• Cosmetic & comment fixes.
• Moved query to private methods.
• Always insert new includes at the top, let clang-format sort out the rest (not yet implemented in the tool, pending fixReplacements)
• Don't take ownership of XrefsDB

[bkramer]

From now on you shall be known as InMemoryXrefsDB.

[bkramer]

include-fixer/IncludeFixer.cpp

132

Sorted. Most of it is dealing with clang callbacks. I also don't consider this to be a public interface as it's all in the cpp file.

184–190

Yes we should. I changed it to insert new includes before the first include, clang-format will clean up the rest. clang-format isn't wired up to the tool yet, I'd like to use the new applyAllReplacements with formatting tools for that, should be ready soon.

include-fixer/IncludeFixer.h

26

Taking a reference now.

36

Removed.

[klimek]

Cool, lg.

include-fixer/IncludeFixer.cpp

134

Ah, right, didn't see that it's all in the cpp.

CMakeLists.txt

 add_subdirectory(clang-apply-replacements)
 add_subdirectory(clang-rename)
 add_subdirectory(modularize)
 if(CLANG_ENABLE_STATIC_ANALYZER)
 add_subdirectory(clang-tidy)
 endif()
 
 add_subdirectory(clang-query)
+add_subdirectory(include-fixer)
 add_subdirectory(pp-trace)
 add_subdirectory(tool-template)
 
 # Add the common testsuite after all the tools.
 # TODO: Support tests with more granularity when features are off?
 if(CLANG_ENABLE_STATIC_ANALYZER AND CLANG_INCLUDE_TESTS)
 add_subdirectory(test)
 add_subdirectory(unittests)
 endif()
 
 option(CLANG_TOOLS_EXTRA_INCLUDE_DOCS "Generate build targets for the Clang Extra Tools docs."
   ${LLVM_INCLUDE_DOCS})
 if( CLANG_TOOLS_EXTRA_INCLUDE_DOCS )
   add_subdirectory(docs)
 endif()
 

include-fixer/CMakeLists.txt

+set(LLVM_LINK_COMPONENTS
+  support
+  )
+
+add_clang_library(clangIncludeFixer
+  IncludeFixer.cpp
+  FixedXrefsDB.cpp
+
+  LINK_LIBS
+  clangAST
+  clangBasic
+  clangFrontend
+  clangLex
+  clangParse
+  clangSema
+  clangTooling
+  clangToolingCore
+  )
+
+add_subdirectory(tool)

include-fixer/FixedXrefsDB.h

+//===-- FixedXrefsDB.h - Xrefs DB for testing -------------------*- 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_FIXED_XREFS_DB_H
+#define LLVM_CLANG_TOOLS_EXTRA_INCLUDE_FIXER_FIXED_XREFS_DB_H
+
+#include "XrefsDB.h"
+
+namespace clang {
+namespace include_fixer {
+
+/// Xref database with fixed content, intended for testing
+// FIXME: Move to unittest once ClangIncludeFixer doesn't depend on it anymore.
+class FixedXrefsDB : public XrefsDB {
+public:
+  std::vector<std::string> search(llvm::StringRef Identifier) override;
+};
+
+} // namespace include_fixer
+} // namespace clang
+
+#endif // LLVM_CLANG_TOOLS_EXTRA_INCLUDE_FIXER_FIXED_XREFS_DB_H

include-fixer/FixedXrefsDB.cpp

+//===-- FixedXrefsDB.cpp - Xrefs DB for testing ---------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "FixedXrefsDB.h"
+
+namespace clang {
+namespace include_fixer {
+
+std::vector<std::string> FixedXrefsDB::search(llvm::StringRef Identifier) {
+  std::vector<std::string> Result;
+  if (Identifier.startswith("std::string") ||
+      Identifier.startswith("::std::string"))
+    Result.push_back("<string>");
+  return Result;
+}
+}
+}

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_INCLUDE_FIXER_H
+#define LLVM_CLANG_TOOLS_EXTRA_INCLUDE_FIXER_INCLUDE_FIXER_H
+
+#include "XrefsDB.h"
+#include "clang/Lex/PreprocessorOptions.h"
+#include "clang/Tooling/Core/Replacement.h"
+#include "clang/Tooling/Tooling.h"
+
+namespace clang {
+namespace include_fixer {
+
+class IncludeFixerActionFactory : public clang::tooling::ToolAction {
+public:
+  /// \param Xrefs A source for matching symbols to header files.
+  /// \param Replacements Storage for the output of the fixer.
+  IncludeFixerActionFactory(
+      std::unique_ptr<XrefsDB> Xrefs,
+      std::vector<clang::tooling::Replacement> &Replacements);
+  ~IncludeFixerActionFactory();
+
+  bool
+  runInvocation(clang::CompilerInvocation *Invocation,
+                clang::FileManager *Files,
+                std::shared_ptr<clang::PCHContainerOperations> PCHContainerOps,
+                clang::DiagnosticConsumer *Diagnostics) override;
+
+  XrefsDB *getXrefsDB() const { return Xrefs.get(); }
+
+private:
+  /// The client to use to find cross-references.
+  std::unique_ptr<XrefsDB> Xrefs;
+
+  /// Replacements are written here.
+  std::vector<clang::tooling::Replacement> &Replacements;
+};
+
+} // namespace include_fixer
+} // namespace clang
+
+#endif // LLVM_CLANG_TOOLS_EXTRA_INCLUDE_FIXER_INCLUDE_FIXER_H

include-fixer/IncludeFixer.cpp

+//===-- IncludeFixer.cpp - Include inserter based on sema callbacks -------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "IncludeFixer.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "clang/Frontend/TextDiagnosticPrinter.h"
+#include "clang/Lex/Preprocessor.h"
+#include "clang/Parse/ParseAST.h"
+#include "clang/Rewrite/Core/Rewriter.h"
+#include "clang/Sema/ExternalSemaSource.h"
+#include "clang/Sema/Sema.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
+
+#define DEBUG_TYPE "include-fixer"
+
+using namespace clang;
+
+namespace clang {
+namespace include_fixer {
+namespace {
+
+class Action;
+
+class PreprocessorHooks : public clang::PPCallbacks {
+public:
+  explicit PreprocessorHooks(Action *enclosing_pass)
+      : EnclosingPass(enclosing_pass), TrackedFile(nullptr) {}
+
+  void FileChanged(clang::SourceLocation loc,
+                   clang::PPCallbacks::FileChangeReason reason,
+                   clang::SrcMgr::CharacteristicKind file_type,
+                   clang::FileID prev_fid) override;
+
+  void InclusionDirective(clang::SourceLocation HashLocation,
+                          const clang::Token &IncludeToken,
+                          llvm::StringRef FileName, bool IsAngled,
+                          clang::CharSourceRange FileNameRange,
+                          const clang::FileEntry *IncludeFile,
+                          llvm::StringRef SearchPath,
+                          llvm::StringRef relative_path,
+                          const clang::Module *imported) override;
+
+private:
+  /// The current Action.
+  Action *EnclosingPass;
+
+  /// The current FileEntry.
+  const clang::FileEntry *TrackedFile;
+};
+
+/// \brief Manages a full parse and any subsequent reparses for a single file.
+class Action : public clang::ASTFrontendAction,
+               public clang::ExternalSemaSource {
+public:
+  explicit Action(IncludeFixerActionFactory &Factory) : Factory(Factory) {}
+
+  std::unique_ptr<clang::ASTConsumer>
+  CreateASTConsumer(clang::CompilerInstance &Compiler,
+                    StringRef InFile) override {
+    Filename = InFile;
+    Compiler.getPreprocessor().addPPCallbacks(
+        llvm::make_unique<PreprocessorHooks>(this));
+    return llvm::make_unique<clang::ASTConsumer>();
+  }
+
+  void ExecuteAction() override {
+    clang::CompilerInstance *Compiler = &getCompilerInstance();
+    assert(!Compiler->hasSema() && "CI already has Sema");
+
+    // Set up our hooks into sema and parse the AST.
+    if (hasCodeCompletionSupport() &&
+        !Compiler->getFrontendOpts().CodeCompletionAt.FileName.empty())
+      Compiler->createCodeCompletionConsumer();
+
+    clang::CodeCompleteConsumer *CompletionConsumer = nullptr;
+    if (Compiler->hasCodeCompletionConsumer())
+      CompletionConsumer = &Compiler->getCodeCompletionConsumer();
+
+    Compiler->createSema(getTranslationUnitKind(), CompletionConsumer);
+    Compiler->getSema().addExternalSource(this);
+
+    clang::ParseAST(Compiler->getSema(), Compiler->getFrontendOpts().ShowStats,
+                    Compiler->getFrontendOpts().SkipFunctionBodies);
+  }
+
+  /// Callback for incomplete types. If we encounter a forward declaration we
+  /// have the fully qualified name ready. Just query that.
+  bool MaybeDiagnoseMissingCompleteType(clang::SourceLocation Loc,
+                                        clang::QualType T) override {
+    clang::ASTContext &context = getCompilerInstance().getASTContext();
+    query(T.getUnqualifiedType().getAsString(context.getPrintingPolicy()));
+    return false;
+  }
+
+  /// Callback for unknown identifiers. Try to piece together as much
+  /// qualification as we can get and do a query.
+  clang::TypoCorrection CorrectTypo(const DeclarationNameInfo &Typo,
+                                    int LookupKind, Scope *S, CXXScopeSpec *SS,
+                                    CorrectionCandidateCallback &CCC,
+                                    DeclContext *MemberContext,
+                                    bool EnteringContext,
+                                    const ObjCObjectPointerType *OPT) override {
+    // We don't want to look up inner parts of nested name specifies. Looking up
+    // the header where a namespace is defined in is rarely useful.
+    if (LookupKind == clang::Sema::LookupNestedNameSpecifierName) {
+      DEBUG(llvm::dbgs() << "ignoring " << Typo.getAsString() << "\n");
+      return clang::TypoCorrection();
+    }
+
+    /// If we have a scope specification use that to get more precise results.
+    std::string QueryString;
+    if (SS && SS->getRange().isValid()) {
+      auto Range = CharSourceRange::getTokenRange(SS->getRange().getBegin(),
+                                                  Typo.getLoc());
+      QueryString =
+          Lexer::getSourceText(Range, getCompilerInstance().getSourceManager(),
+                               getCompilerInstance().getLangOpts());
+    } else {
+      QueryString = Typo.getAsString();
+    }
+
+    return query(QueryString);
+  }
+
+private:
+  /// Query the database for a given identifier.
+  clang::TypoCorrection query(StringRef Query) {
+    assert(!Query.empty() && "Empty query!");
+
+    // Save database lookups by not looking up identifiers multiple times.
+    if (!SeenQueries.insert(Query).second)
+      return clang::TypoCorrection();
+
+    DEBUG(llvm::dbgs() << "Looking up " << Query << " ... ");
+
+    std::string error_text;
+    auto SearchReply = Factory.getXrefsDB()->search(Query);
+    DEBUG(llvm::dbgs() << SearchReply.size() << " replies\n");
+    if (SearchReply.empty())
+      return clang::TypoCorrection();
+
+    // Add those files to the set of includes to try out.
+    // FIXME: Rank the results and pick the best one instead of the first one.
+    TryInclude(Query, SearchReply[0]);
+
+    // FIXME: We should just return the name we got as input here and prevent
+    // clang from trying to correct the typo by itself. That may change the
+    // identifier to something that's not wanted by the user.
+    return clang::TypoCorrection();
+  }
+
+public:
+  StringRef filename() const { return Filename; }
+
+  /// Called for each include file we discover is in the file.
+  /// \param SourceManager the active SourceManager
+  /// \param canonical_path the canonical path to the include file
+  /// \param uttered_path the path as it appeared in the program
+  /// \param IsAngled whether angle brackets were used
+  /// \param HashLocation the source location of the include's \#
+  /// \param EndLocation the source location following the include
+  void NextInclude(clang::SourceManager *SourceManager,
+                   llvm::StringRef canonical_path, llvm::StringRef uttered_path,
+                   bool IsAngled, clang::SourceLocation HashLocation,
+                   clang::SourceLocation EndLocation) {
+    unsigned Offset = SourceManager->getFileOffset(EndLocation);
+    if (Offset > LastIncludeOffset) {
+      LastIncludeOffset = Offset;
+    }
+  }
+
+  /// \brief Rewrite the associated source file with our tentative suggestions.
+  /// \param rewriter a valid Rewriter.
+  /// \return true if changes will be made, false otherwise.
+  bool Rewrite(clang::SourceManager &SourceManager,
+               std::vector<clang::tooling::Replacement> &replacements) {
+    for (const auto &ToTry : Untried) {
+      DEBUG(llvm::dbgs() << "Adding include " << ToTry << "\n");
+      std::string ToAdd = "\n#include " + ToTry;
+      // If this is the only include in the file, add the newline after it, not
+      // before.
+      if (LastIncludeOffset == 0)
+        std::rotate(ToAdd.begin(), ToAdd.begin() + 1, ToAdd.end());
+
+      replacements.push_back(clang::tooling::Replacement(
+          SourceManager, FileBegin.getLocWithOffset(LastIncludeOffset), 0,
+          ToAdd));
+
+      // 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.
+      // FIXME: Handle multiple includes at once.
+      return true;
+    }
+    return false;
+  }
+
+  /// \brief Gets the location at the very top of the file (in this pass).
+  clang::SourceLocation file_begin() const { return FileBegin; }
+
+  /// \brief Sets the location at the very top of the file (in this pass).
+  void setFileBegin(clang::SourceLocation Location) { FileBegin = Location; }
+
+  /// \brief Add an include to the set of includes to try.
+  /// \param include_path The include path to try (as a quoted include).
+  void TryInclude(const std::string &query, const std::string &include_path) {
+    Untried.insert(include_path);
+  }
+
+private:
+  /// The `ActionFactory` orchestrating this run.
+  IncludeFixerActionFactory &Factory;
+
+  // Remeber things we looked up to avoid querying things twice.
+  llvm::StringSet<> SeenQueries;
+
+  /// The absolute path to the file being processed.
+  std::string Filename;
+
+  /// The location of the beginning of the tracked file. This changes after
+  /// each pass.
+  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 LastIncludeOffset = 0;
+
+  /// Includes we have left to try.
+  std::set<std::string> Untried;
+};
+
+void PreprocessorHooks::FileChanged(clang::SourceLocation Loc,
+                                    clang::PPCallbacks::FileChangeReason Reason,
+                                    clang::SrcMgr::CharacteristicKind FileType,
+                                    clang::FileID PrevFID) {
+  // Remember where the main file starts.
+  if (Reason == clang::PPCallbacks::EnterFile) {
+    clang::SourceManager *SourceManager =
+        &EnclosingPass->getCompilerInstance().getSourceManager();
+    clang::FileID loc_id = SourceManager->getFileID(Loc);
+    if (const clang::FileEntry *FileEntry =
+            SourceManager->getFileEntryForID(loc_id)) {
+      if (FileEntry->getName() == EnclosingPass->filename()) {
+        EnclosingPass->setFileBegin(Loc);
+        TrackedFile = FileEntry;
+      }
+    }
+  }
+}
+
+void PreprocessorHooks::InclusionDirective(
+    clang::SourceLocation HashLocation, const clang::Token &IncludeToken,
+    llvm::StringRef FileName, bool IsAngled,
+    clang::CharSourceRange FileNameRange, const clang::FileEntry *IncludeFile,
+    llvm::StringRef SearchPath, llvm::StringRef relative_path,
+    const clang::Module *imported) {
+  // Remember include locations so we can insert our new include at the end of
+  // the include block.
+  clang::SourceManager *SourceManager =
+      &EnclosingPass->getCompilerInstance().getSourceManager();
+  auto IDPosition = SourceManager->getDecomposedExpansionLoc(HashLocation);
+  const FileEntry *SourceFile =
+      SourceManager->getFileEntryForID(IDPosition.first);
+  if (!IncludeFile || TrackedFile != SourceFile)
+    return;
+  EnclosingPass->NextInclude(SourceManager, IncludeFile->getName(), FileName,
+                             IsAngled, HashLocation, FileNameRange.getEnd());
+}
+
+} // namespace
+
+IncludeFixerActionFactory::IncludeFixerActionFactory(
+    std::unique_ptr<XrefsDB> Xrefs,
+    std::vector<clang::tooling::Replacement> &Replacements)
+    : Xrefs(std::move(Xrefs)), Replacements(Replacements) {}
+
+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);
+
+  // Set up Clang.
+  clang::CompilerInstance Compiler(PCHContainerOps);
+  Compiler.setInvocation(Invocation);
+  Compiler.setFileManager(Files);
+
+  // Create the compiler's actual diagnostics engine. We want to drop all
+  // diagnostics here.
+  Compiler.createDiagnostics(new clang::IgnoringDiagConsumer,
+                             /*ShouldOwnClient=*/true);
+  Compiler.createSourceManager(*Files);
+
+  // Run the parser, gather missing includes.
+  auto ScopedToolAction = llvm::make_unique<Action>(*this);
+  Compiler.ExecuteAction(*ScopedToolAction);
+
+  // Generate replacements.
+  ScopedToolAction->Rewrite(Compiler.getSourceManager(), Replacements);
+
+  // Technically this should only return true if we're sure that we have a
+  // parseable file. We don't know that though.
+  return true;
+}
+
+} // namespace include_fixer
+} // namespace clang

include-fixer/XrefsDB.h

+//===-- XrefsDB.h - Interface for symbol-header matching --------*- 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_XREFS_DB_H
+#define LLVM_CLANG_TOOLS_EXTRA_INCLUDE_FIXER_XREFS_DB_H
+
+#include "llvm/ADT/StringRef.h"
+#include <vector>
+
+namespace clang {
+namespace include_fixer {
+
+/// This class provides an interface for finding the header files corresponding
+/// to an indentifier in the source code.
+class XrefsDB {
+public:
+  virtual ~XrefsDB() = default;
+
+  /// Search for header files to be included for an identifier.
+  /// \param Identifier The identifier being searched for. May or may not be
+  ///                   fully qualified.
+  /// \returns A list of inclusion candidates, in a format ready for being
+  ///          pasted after an #include token.
+  // FIXME: Expose the type name so we can also insert using declarations (or
+  // fix the usage)
+  virtual std::vector<std::string> search(llvm::StringRef Identifier) = 0;
+};
+
+} // namespace include_fixer
+} // namespace clang
+
+#endif // LLVM_CLANG_TOOLS_EXTRA_INCLUDE_FIXER_XREFS_DB_H

include-fixer/tool/CMakeLists.txt

+include_directories(${CMAKE_CURRENT_SOURCE_DIR}/..)
+
+add_clang_executable(clang-include-fixer ClangIncludeFixer.cpp)
+target_link_libraries(clang-include-fixer
+  clangBasic
+  clangFrontend
+  clangIncludeFixer
+  clangRewrite
+  clangTooling
+  clangToolingCore
+  )

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 "FixedXrefsDB.h"
+#include "IncludeFixer.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;
+
+static llvm::cl::OptionCategory tool_options("Tool options");
+
+int main(int argc, char **argv) {
+  clang::tooling::CommonOptionsParser options(argc, (const char **)argv,
+                                              tool_options);
+  clang::tooling::ClangTool tool(options.getCompilations(),
+                                 options.getSourcePathList());
+  // Set up the data source.
+  auto xrefs_db = llvm::make_unique<include_fixer::FixedXrefsDB>();
+
+  // Now run our tool.
+  std::vector<clang::tooling::Replacement> replacements;
+  include_fixer::IncludeFixerActionFactory factory(std::move(xrefs_db),
+                                                   replacements);
+
+  tool.run(&factory); // Always succeeds.
+
+  // Set up a new source manager for applying the resulting replacements.
+  llvm::IntrusiveRefCntPtr<clang::DiagnosticOptions> DiagOpts(
+      new clang::DiagnosticOptions);
+  clang::DiagnosticsEngine Diagnostics(new clang::DiagnosticIDs, &*DiagOpts);
+  clang::TextDiagnosticPrinter DiagnosticPrinter(llvm::outs(), &*DiagOpts);
+  clang::SourceManager source_manager(Diagnostics, tool.getFiles());
+  Diagnostics.setClient(&DiagnosticPrinter, false);
+
+  // Write replacements to disk.
+  clang::Rewriter Rewrites(source_manager, clang::LangOptions());
+  clang::tooling::applyAllReplacements(replacements, Rewrites);
+  return Rewrites.overwriteChangedFiles();
+}

unittests/CMakeLists.txt

 add_custom_target(ExtraToolsUnitTests)
 set_target_properties(ExtraToolsUnitTests PROPERTIES FOLDER "Extra Tools Unit Tests")
 
 function(add_extra_unittest test_dirname)
   add_unittest(ExtraToolsUnitTests ${test_dirname} ${ARGN})
 endfunction()
 
 add_subdirectory(clang-apply-replacements)
 add_subdirectory(clang-rename)
 add_subdirectory(clang-query)
 add_subdirectory(clang-tidy)
+add_subdirectory(include-fixer)

unittests/include-fixer/CMakeLists.txt

+set(LLVM_LINK_COMPONENTS
+  support
+  )
+
+get_filename_component(INCLUDE_FIXER_SOURCE_DIR
+  ${CMAKE_CURRENT_SOURCE_DIR}/../../include-fixer REALPATH)
+include_directories(
+  ${INCLUDE_FIXER_SOURCE_DIR}
+  )
+
+add_extra_unittest(IncludeFixerTests
+  IncludeFixerTest.cpp
+  )
+
+target_link_libraries(IncludeFixerTests
+  clangBasic
+  clangFrontend
+  clangIncludeFixer
+  clangRewrite
+  clangTooling
+  clangToolingCore
+  )

unittests/include-fixer/IncludeFixerTest.cpp

+//===-- IncludeFixerTest.cpp - Include fixer unit tests -------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "../../../../unittests/Tooling/RewriterTestContext.h"
+#include "FixedXrefsDB.h"
+#include "IncludeFixer.h"
+#include "clang/Tooling/Tooling.h"
+#include "gtest/gtest.h"
+using namespace clang;
+
+namespace clang {
+namespace include_fixer {
+namespace {
+
+static bool runOnCode(tooling::ToolAction *ToolAction, StringRef Code,
+                      StringRef FileName) {
+  llvm::IntrusiveRefCntPtr<vfs::InMemoryFileSystem> InMemoryFileSystem(
+      new vfs::InMemoryFileSystem);
+  llvm::IntrusiveRefCntPtr<FileManager> Files(
+      new FileManager(FileSystemOptions(), InMemoryFileSystem));
+  tooling::ToolInvocation Invocation(
+      {std::string("include_fixer"), std::string("-fsyntax-only"),
+       FileName.str()},
+      ToolAction, Files.get(), std::make_shared<PCHContainerOperations>());
+
+  InMemoryFileSystem->addFile(FileName, 0,
+                              llvm::MemoryBuffer::getMemBuffer(Code));
+
+  InMemoryFileSystem->addFile("foo.h", 0,
+                              llvm::MemoryBuffer::getMemBuffer("\n"));
+  InMemoryFileSystem->addFile("bar.h", 0,
+                              llvm::MemoryBuffer::getMemBuffer("\n"));
+  return Invocation.run();
+}
+
+static std::string runIncludeFixer(StringRef Code) {
+  auto xrefs_db = llvm::make_unique<FixedXrefsDB>();
+  std::vector<clang::tooling::Replacement> replacements;
+  IncludeFixerActionFactory factory(std::move(xrefs_db), replacements);
+  runOnCode(&factory, Code, "input.cc");
+  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>\nstd::string foo;\n",
+            runIncludeFixer("std::string foo;\n"));
+
+  // FIXME: Don't add includes after code.
+  EXPECT_EQ("#include \"foo.h\"\nstd::string foo;\n"
+            "#include \"bar.h\"\n#include <string>\n",
+            runIncludeFixer(
+                "#include \"foo.h\"\nstd::string foo;\n#include \"bar.h\"\n"));
+
+  EXPECT_EQ("#include \"foo.h\"\n#include <string>\nstd::string foo;\n",
+            runIncludeFixer("#include \"foo.h\"\nstd::string foo;\n"));
+
+  EXPECT_EQ(
+      "#include \"foo.h\"\n#include <string>\nstd::string::size_type foo;\n",
+      runIncludeFixer("#include \"foo.h\"\nstd::string::size_type foo;\n"));
+
+  // The fixed xrefs db doesn't know how to handle string without std::.
+  EXPECT_EQ("string foo;\n", runIncludeFixer("string foo;\n"));
+}
+
+TEST(IncludeFixer, IncompleteType) {
+  EXPECT_EQ(
+      "#include \"foo.h\"\n#include <string>\n"
+      "namespace std {\nclass string;\n}\nstring foo;\n",
+      runIncludeFixer("#include \"foo.h\"\n"
+                      "namespace std {\nclass string;\n}\nstring foo;\n"));
+}
+
+} // namespace
+} // namespace include_fixer
+} // namespace clang