Add SourceManagerForFile helper which sets up SourceManager and dependencies for a single file with code snippet

[ioeric]

Summary

This can be used to create a virtual environment (incl. VFS, source manager) for code snippets.

[ioeric]

[WIP][Tooling] Pull #include manipulation code from clangFormat into libToolingCore.

[ioeric]

• Revert last revision.

[sammccall]

include/clang/Tooling/Core/Environment.h

2

I'm not sure about this abstraction - I can't tell what it represents, and what in future should/shouldn't be added to it.

• The read API seems to be SourceManager& and FileID, and there doesn't seem to be much interaction. I think I'd usually find it easier to understand to have those two than an Environment.
• the utility to create a self-contained source manager wrapping a single virtual file is useful, but seems like that's more specialized/narrower (struct SingleSourceManager?)
• the fact that these virtual sourcemanagers are completely self-owning when passed as Environment seems maybe-convenient but also surprising and non-orthogonal (it's a maybe-owning-pointer by another name). Is this important somewhere?

[ioeric]

• Create SourceManagerForFile instead of Environment; put it in SourceManager.h which seems to be a better place.

[ioeric]

Thanks for the comment and suggestion!

I've changed this to SourceManagerForFile and move it to SourceManager.h, which seems to be a more natural fit. SourceManagerForFile sounds like a sensible name to me, but I'm guess there might be a better name...

[ioeric]

• Replaced forward decl of DiagnosticsEngine with header include.

[sammccall]

include/clang/Basic/SourceManager.h

1819 ↗

(On Diff #145386)

nit: single in-memory file

1824 ↗

(On Diff #145386)

why is this not just a constructor? it looks like it can't fail?

1827 ↗

(On Diff #145386)

this could also just be get() or operator*

lib/Format/TokenAnalyzer.h

75

add a comment that this is only present if constructed from a string?

[ioeric]

• address review comments.

include/clang/Format/Format.h

 /// Various functions to configurably format source code.
 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_FORMAT_FORMAT_H
 #define LLVM_CLANG_FORMAT_FORMAT_H
 
 #include "clang/Basic/LangOptions.h"
+#include "clang/Tooling/Core/HeaderIncludes.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/Support/Regex.h"
 #include <system_error>
 
 namespace clang {
 
 class Lexer;
 class SourceManager;
 class DiagnosticConsumer;
 
struct FormatStyle {
   /// In the .clang-format configuration file, this can be configured like:
   /// \code{.yaml}
   ///   ForEachMacros: ['RANGES_FOR', 'FOREACH']
   /// \endcode
   ///
   /// For example: BOOST_FOREACH.
   std::vector<std::string> ForEachMacros;
 
-  /// \brief Styles for sorting multiple ``#include`` blocks.
-  enum IncludeBlocksStyle {
-    /// \brief Sort each ``#include`` block separately.
-    /// \code
-    ///    #include "b.h"               into      #include "b.h"
-    ///
-    ///    #include <lib/main.h>                  #include "a.h"
-    ///    #include "a.h"                         #include <lib/main.h>
-    /// \endcode
-    IBS_Preserve,
-    /// \brief Merge multiple ``#include`` blocks together and sort as one.
-    /// \code
-    ///    #include "b.h"               into      #include "a.h"
-    ///                                           #include "b.h"
-    ///    #include <lib/main.h>                  #include <lib/main.h>
-    ///    #include "a.h"
-    /// \endcode
-    IBS_Merge,
-    /// \brief Merge multiple ``#include`` blocks together and sort as one.
-    /// Then split into groups based on category priority. See
-    /// ``IncludeCategories``.
-    /// \code
-    ///    #include "b.h"               into      #include "a.h"
-    ///                                           #include "b.h"
-    ///    #include <lib/main.h>
-    ///    #include "a.h"                         #include <lib/main.h>
-    /// \endcode
-    IBS_Regroup,
-  };
-
-  /// \brief Dependent on the value, multiple ``#include`` blocks can be sorted
-  /// as one and divided based on category.
-  IncludeBlocksStyle IncludeBlocks;
-
-  /// \brief See documentation of ``IncludeCategories``.
-  struct IncludeCategory {
-    /// \brief The regular expression that this category matches.
-    std::string Regex;
-    /// \brief The priority to assign to this category.
-    int Priority;
-    bool operator==(const IncludeCategory &Other) const {
-      return Regex == Other.Regex && Priority == Other.Priority;
-    }
-  };
-
-  /// \brief Regular expressions denoting the different ``#include`` categories
-  /// used for ordering ``#includes``.
-  ///
-  /// These regular expressions are matched against the filename of an include
-  /// (including the <> or "") in order. The value belonging to the first
-  /// matching regular expression is assigned and ``#includes`` are sorted first
-  /// according to increasing category number and then alphabetically within
-  /// each category.
-  ///
-  /// If none of the regular expressions match, INT_MAX is assigned as
-  /// category. The main header for a source file automatically gets category 0.
-  /// so that it is generally kept at the beginning of the ``#includes``
-  /// (http://llvm.org/docs/CodingStandards.html#include-style). However, you
-  /// can also assign negative priorities if you have certain headers that
-  /// always need to be first.
-  ///
-  /// To configure this in the .clang-format file, use:
-  /// \code{.yaml}
-  ///   IncludeCategories:
-  ///     - Regex:           '^"(llvm|llvm-c|clang|clang-c)/'
-  ///       Priority:        2
-  ///     - Regex:           '^(<|"(gtest|gmock|isl|json)/)'
-  ///       Priority:        3
-  ///     - Regex:           '.*'
-  ///       Priority:        1
-  /// \endcode
-  std::vector<IncludeCategory> IncludeCategories;
-
-  /// \brief Specify a regular expression of suffixes that are allowed in the
-  /// file-to-main-include mapping.
-  ///
-  /// When guessing whether a #include is the "main" include (to assign
-  /// category 0, see above), use this regex of allowed suffixes to the header
-  /// stem. A partial match is done, so that:
-  /// - "" means "arbitrary suffix"
-  /// - "$" means "no suffix"
-  ///
-  /// For example, if configured to "(_test)?$", then a header a.h would be seen
-  /// as the "main" include in both a.cc and a_test.cc.
-  std::string IncludeIsMainRegex;
+  tooling::IncludeStyle IncludeStyle;
 
   /// \brief Indent case labels one level from the switch statement.
   ///
   /// When ``false``, use the same indentation level as for the switch statement.
   /// Switch statement body is always indented one level more than case labels.
   /// \code
   ///    false:                                 true:
   ///    switch (fool) {                vs.     switch (fool) {
return AccessModifierOffset == R.AccessModifierOffset &&
            ContinuationIndentWidth == R.ContinuationIndentWidth &&
            Cpp11BracedListStyle == R.Cpp11BracedListStyle &&
            DerivePointerAlignment == R.DerivePointerAlignment &&
            DisableFormat == R.DisableFormat &&
            ExperimentalAutoDetectBinPacking ==
                R.ExperimentalAutoDetectBinPacking &&
            FixNamespaceComments == R.FixNamespaceComments &&
            ForEachMacros == R.ForEachMacros &&
-           IncludeBlocks == R.IncludeBlocks &&
-           IncludeCategories == R.IncludeCategories &&
+           IncludeStyle.IncludeBlocks == R.IncludeStyle.IncludeBlocks &&
+           IncludeStyle.IncludeCategories == R.IncludeStyle.IncludeCategories &&
            IndentCaseLabels == R.IndentCaseLabels &&
            IndentPPDirectives == R.IndentPPDirectives &&
            IndentWidth == R.IndentWidth && Language == R.Language &&
            IndentWrappedFunctionNames == R.IndentWrappedFunctionNames &&
            JavaScriptQuotes == R.JavaScriptQuotes &&
            JavaScriptWrapImports == R.JavaScriptWrapImports &&
            KeepEmptyLinesAtTheStartOfBlocks ==
                R.KeepEmptyLinesAtTheStartOfBlocks &&
            MacroBlockBegin == R.MacroBlockBegin &&
            MacroBlockEnd == R.MacroBlockEnd &&
            MaxEmptyLinesToKeep == R.MaxEmptyLinesToKeep &&
            NamespaceIndentation == R.NamespaceIndentation &&
            ObjCBinPackProtocolList == R.ObjCBinPackProtocolList &&
            ObjCBlockIndentWidth == R.ObjCBlockIndentWidth &&
            ObjCSpaceAfterProperty == R.ObjCSpaceAfterProperty &&
            ObjCSpaceBeforeProtocolList == R.ObjCSpaceBeforeProtocolList &&
-           PenaltyBreakAssignment ==
-               R.PenaltyBreakAssignment &&
+           PenaltyBreakAssignment == R.PenaltyBreakAssignment &&
            PenaltyBreakBeforeFirstCallParameter ==
                R.PenaltyBreakBeforeFirstCallParameter &&
            PenaltyBreakComment == R.PenaltyBreakComment &&
            PenaltyBreakFirstLessLess == R.PenaltyBreakFirstLessLess &&
            PenaltyBreakString == R.PenaltyBreakString &&
            PenaltyExcessCharacter == R.PenaltyExcessCharacter &&
            PenaltyReturnTypeOnItsOwnLine == R.PenaltyReturnTypeOnItsOwnLine &&
            PointerAlignment == R.PointerAlignment &&
formatReplacements(StringRef Code, const tooling::Replacements &Replaces,
                    const FormatStyle &Style);
 
 /// \brief Returns the replacements corresponding to applying \p Replaces and
 /// cleaning up the code after that on success; otherwise, return an llvm::Error
 /// carrying llvm::StringError.
 /// This also supports inserting/deleting C++ #include directives:
 /// - If a replacement has offset UINT_MAX, length 0, and a replacement text
 ///   that is an #include directive, this will insert the #include into the
-///   correct block in the \p Code. When searching for points to insert new
-///   header, this ignores #include's after the #include block(s) in the
-///   beginning of a file to avoid inserting headers into code sections where
-///   new #include's should not be added by default. These code sections
-///   include:
-///     - raw string literals (containing #include).
-///     - #if blocks.
-///     - Special #include's among declarations (e.g. functions).
+///   correct block in the \p Code.
 /// - If a replacement has offset UINT_MAX, length 1, and a replacement text
 ///   that is the name of the header to be removed, the header will be removed
 ///   from \p Code if it exists.
+/// See documentation for `tooling::HeaderIncludes` for the behavior.
 llvm::Expected<tooling::Replacements>
 cleanupAroundReplacements(StringRef Code, const tooling::Replacements &Replaces,
                           const FormatStyle &Style);
 
 /// \brief Represents the status of a formatting attempt.
 struct FormattingAttemptStatus {
   /// \brief A value of ``false`` means that any of the affected ranges were not
   /// formatted due to a non-recoverable syntax error.

include/clang/Tooling/Core/Environment.h

+//===--- Environment.h - Sourece tooling environment --*- 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_TOOLING_CORE_ENVIRONMENT_H
+#define LLVM_CLANG_TOOLING_CORE_ENVIRONMENT_H
+
+#include "clang/Basic/Diagnostic.h"
+#include "clang/Basic/FileManager.h"
+#include "clang/Basic/SourceManager.h"
+
+namespace clang {
+namespace tooling {
+
+/// An environment for handling file contents with SourceManager.
+class Environment {
+public:
+  Environment(SourceManager &SM, FileID ID) : ID(ID), SM(SM) {}
+
+  Environment(FileID ID, std::unique_ptr<FileManager> FileMgr,
+              std::unique_ptr<SourceManager> VirtualSM,
+              std::unique_ptr<DiagnosticsEngine> Diagnostics)
+      : ID(ID), SM(*VirtualSM), FileMgr(std::move(FileMgr)),
+        VirtualSM(std::move(VirtualSM)), Diagnostics(std::move(Diagnostics)) {}
+
+  // This sets up an virtual file system with file \p FileName containing the
+  // fragment \p Content.
+  static std::unique_ptr<Environment>
+  createVirtualEnvironment(StringRef Content, StringRef FileName);
+
+  FileID getFileID() const { return ID; }
+
+  const SourceManager &getSourceManager() const { return SM; }
+
+private:
+  FileID ID;
+  SourceManager &SM;
+
+  // The order of these fields are important - they should be in the same order
+  // as they are created in `createVirtualEnvironment` so that they can be
+  // deleted in the reverse order as they are created.
+  std::unique_ptr<FileManager> FileMgr;
+  std::unique_ptr<SourceManager> VirtualSM;
+  std::unique_ptr<DiagnosticsEngine> Diagnostics;
+};
+
+} // namespace tooling
+} // namespace clang
+
+#endif // LLVM_CLANG_TOOLING_CORE_ENVIRONMENT_H

include/clang/Tooling/Core/HeaderIncludes.h

+//===--- HeaderIncludes.h - Insert/Delete #includes for C++ code--*- 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_TOOLING_CORE_HEADERINCLUDES_H
+#define LLVM_CLANG_TOOLING_CORE_HEADERINCLUDES_H
+
+#include "clang/Basic/SourceManager.h"
+#include "clang/Tooling/Core/Environment.h"
+#include "clang/Tooling/Core/Replacement.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/Regex.h"
+#include <unordered_map>
+
+namespace clang {
+namespace tooling {
+
+struct IncludeStyle {
+  /// \brief Styles for sorting multiple ``#include`` blocks.
+  enum IncludeBlocksStyle {
+    /// \brief Sort each ``#include`` block separately.
+    /// \code
+    ///    #include "b.h"               into      #include "b.h"
+    ///
+    ///    #include <lib/main.h>                  #include "a.h"
+    ///    #include "a.h"                         #include <lib/main.h>
+    /// \endcode
+    IBS_Preserve,
+    /// \brief Merge multiple ``#include`` blocks together and sort as one.
+    /// \code
+    ///    #include "b.h"               into      #include "a.h"
+    ///                                           #include "b.h"
+    ///    #include <lib/main.h>                  #include <lib/main.h>
+    ///    #include "a.h"
+    /// \endcode
+    IBS_Merge,
+    /// \brief Merge multiple ``#include`` blocks together and sort as one.
+    /// Then split into groups based on category priority. See
+    /// ``IncludeCategories``.
+    /// \code
+    ///    #include "b.h"               into      #include "a.h"
+    ///                                           #include "b.h"
+    ///    #include <lib/main.h>
+    ///    #include "a.h"                         #include <lib/main.h>
+    /// \endcode
+    IBS_Regroup,
+  };
+
+  /// \brief Dependent on the value, multiple ``#include`` blocks can be sorted
+  /// as one and divided based on category.
+  IncludeBlocksStyle IncludeBlocks;
+
+  /// \brief See documentation of ``IncludeCategories``.
+  struct IncludeCategory {
+    /// \brief The regular expression that this category matches.
+    std::string Regex;
+    /// \brief The priority to assign to this category.
+    int Priority;
+    bool operator==(const IncludeCategory &Other) const {
+      return Regex == Other.Regex && Priority == Other.Priority;
+    }
+  };
+
+  /// \brief Regular expressions denoting the different ``#include`` categories
+  /// used for ordering ``#includes``.
+  ///
+  /// These regular expressions are matched against the filename of an include
+  /// (including the <> or "") in order. The value belonging to the first
+  /// matching regular expression is assigned and ``#includes`` are sorted first
+  /// according to increasing category number and then alphabetically within
+  /// each category.
+  ///
+  /// If none of the regular expressions match, INT_MAX is assigned as
+  /// category. The main header for a source file automatically gets category 0.
+  /// so that it is generally kept at the beginning of the ``#includes``
+  /// (http://llvm.org/docs/CodingStandards.html#include-style). However, you
+  /// can also assign negative priorities if you have certain headers that
+  /// always need to be first.
+  ///
+  /// To configure this in the .clang-format file, use:
+  /// \code{.yaml}
+  ///   IncludeCategories:
+  ///     - Regex:           '^"(llvm|llvm-c|clang|clang-c)/'
+  ///       Priority:        2
+  ///     - Regex:           '^(<|"(gtest|gmock|isl|json)/)'
+  ///       Priority:        3
+  ///     - Regex:           '.*'
+  ///       Priority:        1
+  /// \endcode
+  std::vector<IncludeCategory> IncludeCategories;
+
+  /// \brief Specify a regular expression of suffixes that are allowed in the
+  /// file-to-main-include mapping.
+  ///
+  /// When guessing whether a #include is the "main" include (to assign
+  /// category 0, see above), use this regex of allowed suffixes to the header
+  /// stem. A partial match is done, so that:
+  /// - "" means "arbitrary suffix"
+  /// - "$" means "no suffix"
+  ///
+  /// For example, if configured to "(_test)?$", then a header a.h would be seen
+  /// as the "main" include in both a.cc and a_test.cc.
+  std::string IncludeIsMainRegex;
+};
+
+/// This class manages priorities of C++ #include categories and calculates
+/// priorities for headers.
+/// FIXME(ioeric): move this class into implementation file when clang-format's
+/// include sorting functions are also moved here.
+class IncludeCategoryManager {
+public:
+  IncludeCategoryManager(const IncludeStyle &Style, StringRef FileName);
+
+  /// Returns the priority of the category which \p IncludeName belongs to.
+  /// If \p CheckMainHeader is true and \p IncludeName is a main header, returns
+  /// 0. Otherwise, returns the priority of the matching category or INT_MAX.
+  int getIncludePriority(StringRef IncludeName, bool CheckMainHeader) const;
+
+private:
+  bool isMainHeader(StringRef IncludeName) const;
+
+  const IncludeStyle Style;
+  bool IsMainFile;
+  StringRef FileName;
+  StringRef FileStem;
+  mutable SmallVector<llvm::Regex, 4> CategoryRegexs;
+};
+
+/// Generates replacements for inserting or deleting #include headers in a file.
+class HeaderIncludes {
+public:
+  HeaderIncludes(llvm::StringRef FileName, llvm::StringRef Code,
+                 const IncludeStyle &Style);
+
+  /// Inserts an #include directive of \p IncludeName into the code. If \p
+  /// IncludeName is quoted e.g. <...> or "...", it will be #included as is;
+  /// by default,  it is quoted with "".
+  ///
+  /// When searching for points to insert new header, this ignores #include's
+  /// after the #include block(s) in the beginning of a file to avoid inserting
+  /// headers into code sections where new #include's should not be added by
+  /// default. These code sections include:
+  ///   - raw string literals (containing #include).
+  ///   - #if blocks.
+  ///   - Special #include's among declarations (e.g. functions).
+  ///
+  /// Returns a replacement that inserts the new header into a suitable #include
+  /// block (#include order is preserved with best effort e.g. if they are
+  /// already sorted; caller should consider sorting the #includes with
+  /// clang-format after insertions). If \p IncludeName already exists, this
+  /// returns None.
+  llvm::Optional<Replacement> insert(llvm::StringRef IncludeName) const;
+
+  /// Removes all existing #includes of \p Header. If \p Header is not quoted
+  /// (e.g. without <> or ""), #include of \p Header with both quotations will
+  /// be removed.
+  Replacements remove(llvm::StringRef Header) const;
+
+private:
+  struct Include {
+    Include(StringRef Name, Range R) : Name(Name), R(R) {}
+
+    // An include header quoted with either <> or "".
+    std::string Name;
+    // The range of the whole line of include directive including any eading
+    // whitespaces and trailing comment.
+    Range R;
+  };
+
+  void addExistingInclude(Include IncludeToAdd, unsigned NextLineOffset);
+
+  StringRef FileName;
+  StringRef Code;
+
+  // Map from include name (quotation trimmed) to a list of existing includes (in case
+  // there are more than one) with the name in the current file.
+  // <x> and "x" will be treated as the same header when deleting #includes.
+  llvm::StringMap<llvm::SmallVector<Include, 1>> ExistingIncludes;
+
+  std::unordered_map<int, llvm::SmallVector<const Include *, 8>>
+      IncludesByPriority;
+
+  int FirstIncludeOffset;
+  // All new headers should be inserted after this offset.
+  unsigned MinInsertOffset;
+  // Code with the first MinInsertOffset characters trimmed.
+  StringRef TrimmedCode;
+  // Max insertion offset in the original code.
+  unsigned MaxInsertOffset;
+  IncludeCategoryManager Categories;
+  // Record the offset of the end of the last include in each category.
+  std::unordered_map<int, int> CategoryEndOffsets;
+
+  // All possible priorities.
+  std::set<int> Priorities;
+
+  // Matches a whole #include directive.
+  llvm::Regex IncludeRegex;
+};
+
+} // namespace tooling
+} // namespace clang
+
+#endif // LLVM_CLANG_TOOLING_CORE_HEADERINCLUDES_H

lib/Format/Format.cpp

 #include <algorithm>
 #include <memory>
 #include <string>
 
 #define DEBUG_TYPE "format-formatter"
 
 using clang::format::FormatStyle;
 
-LLVM_YAML_IS_SEQUENCE_VECTOR(clang::format::FormatStyle::IncludeCategory)
+LLVM_YAML_IS_SEQUENCE_VECTOR(clang::tooling::IncludeStyle::IncludeCategory)
 LLVM_YAML_IS_SEQUENCE_VECTOR(clang::format::FormatStyle::RawStringFormat)
 
 namespace llvm {
 namespace yaml {
 template <> struct ScalarEnumerationTraits<FormatStyle::LanguageKind> {
   static void enumeration(IO &IO, FormatStyle::LanguageKind &Value) {
     IO.enumCase(Value, "Cpp", FormatStyle::LK_Cpp);
     IO.enumCase(Value, "Java", FormatStyle::LK_Java);
static void mapping(IO &IO, FormatStyle &Style) {
     IO.mapOptional("ContinuationIndentWidth", Style.ContinuationIndentWidth);
     IO.mapOptional("Cpp11BracedListStyle", Style.Cpp11BracedListStyle);
     IO.mapOptional("DerivePointerAlignment", Style.DerivePointerAlignment);
     IO.mapOptional("DisableFormat", Style.DisableFormat);
     IO.mapOptional("ExperimentalAutoDetectBinPacking",
                    Style.ExperimentalAutoDetectBinPacking);
     IO.mapOptional("FixNamespaceComments", Style.FixNamespaceComments);
     IO.mapOptional("ForEachMacros", Style.ForEachMacros);
-    IO.mapOptional("IncludeBlocks", Style.IncludeBlocks);
-    IO.mapOptional("IncludeCategories", Style.IncludeCategories);
-    IO.mapOptional("IncludeIsMainRegex", Style.IncludeIsMainRegex);
+    IO.mapOptional("IncludeBlocks", Style.IncludeStyle.IncludeBlocks);
+    IO.mapOptional("IncludeCategories", Style.IncludeStyle.IncludeCategories);
+    IO.mapOptional("IncludeIsMainRegex", Style.IncludeStyle.IncludeIsMainRegex);
     IO.mapOptional("IndentCaseLabels", Style.IndentCaseLabels);
     IO.mapOptional("IndentPPDirectives", Style.IndentPPDirectives);
     IO.mapOptional("IndentWidth", Style.IndentWidth);
     IO.mapOptional("IndentWrappedFunctionNames",
                    Style.IndentWrappedFunctionNames);
     IO.mapOptional("JavaScriptQuotes", Style.JavaScriptQuotes);
     IO.mapOptional("JavaScriptWrapImports", Style.JavaScriptWrapImports);
     IO.mapOptional("KeepEmptyLinesAtTheStartOfBlocks",
static void mapping(IO &IO, FormatStyle::BraceWrappingFlags &Wrapping) {
     IO.mapOptional("BeforeElse", Wrapping.BeforeElse);
     IO.mapOptional("IndentBraces", Wrapping.IndentBraces);
     IO.mapOptional("SplitEmptyFunction", Wrapping.SplitEmptyFunction);
     IO.mapOptional("SplitEmptyRecord", Wrapping.SplitEmptyRecord);
     IO.mapOptional("SplitEmptyNamespace", Wrapping.SplitEmptyNamespace);
   }
 };
 
-template <> struct MappingTraits<FormatStyle::IncludeCategory> {
-  static void mapping(IO &IO, FormatStyle::IncludeCategory &Category) {
+using clang::tooling::IncludeStyle;
+
+template <> struct MappingTraits<IncludeStyle::IncludeCategory> {
+  static void mapping(IO &IO, IncludeStyle::IncludeCategory &Category) {
     IO.mapOptional("Regex", Category.Regex);
     IO.mapOptional("Priority", Category.Priority);
   }
 };
 
-template <> struct ScalarEnumerationTraits<FormatStyle::IncludeBlocksStyle> {
-  static void enumeration(IO &IO, FormatStyle::IncludeBlocksStyle &Value) {
-    IO.enumCase(Value, "Preserve", FormatStyle::IBS_Preserve);
-    IO.enumCase(Value, "Merge", FormatStyle::IBS_Merge);
-    IO.enumCase(Value, "Regroup", FormatStyle::IBS_Regroup);
+template <> struct ScalarEnumerationTraits<IncludeStyle::IncludeBlocksStyle> {
+  static void enumeration(IO &IO, IncludeStyle::IncludeBlocksStyle &Value) {
+    IO.enumCase(Value, "Preserve", IncludeStyle::IBS_Preserve);
+    IO.enumCase(Value, "Merge", IncludeStyle::IBS_Merge);
+    IO.enumCase(Value, "Regroup", IncludeStyle::IBS_Regroup);
   }
 };
 
 template <> struct MappingTraits<FormatStyle::RawStringFormat> {
   static void mapping(IO &IO, FormatStyle::RawStringFormat &Format) {
     IO.mapOptional("Language", Format.Language);
     IO.mapOptional("Delimiters", Format.Delimiters);
     IO.mapOptional("EnclosingFunctions", Format.EnclosingFunctions);
FormatStyle getLLVMStyle() {
   LLVMStyle.ContinuationIndentWidth = 4;
   LLVMStyle.Cpp11BracedListStyle = true;
   LLVMStyle.DerivePointerAlignment = false;
   LLVMStyle.ExperimentalAutoDetectBinPacking = false;
   LLVMStyle.FixNamespaceComments = true;
   LLVMStyle.ForEachMacros.push_back("foreach");
   LLVMStyle.ForEachMacros.push_back("Q_FOREACH");
   LLVMStyle.ForEachMacros.push_back("BOOST_FOREACH");
-  LLVMStyle.IncludeCategories = {{"^\"(llvm|llvm-c|clang|clang-c)/", 2},
-                                 {"^(<|\"(gtest|gmock|isl|json)/)", 3},
-                                 {".*", 1}};
-  LLVMStyle.IncludeIsMainRegex = "(Test)?$";
-  LLVMStyle.IncludeBlocks = FormatStyle::IBS_Preserve;
+  LLVMStyle.IncludeStyle.IncludeCategories = {
+      {"^\"(llvm|llvm-c|clang|clang-c)/", 2},
+  LLVMStyle.IncludeStyle.IncludeIsMainRegex = "(Test)?$";
+  LLVMStyle.IncludeStyle.IncludeBlocks = tooling::IncludeStyle::IBS_Preserve;
   LLVMStyle.IndentCaseLabels = false;
   LLVMStyle.IndentPPDirectives = FormatStyle::PPDIS_None;
   LLVMStyle.IndentWrappedFunctionNames = false;
   LLVMStyle.IndentWidth = 2;
   LLVMStyle.JavaScriptQuotes = FormatStyle::JSQS_Leave;
   LLVMStyle.JavaScriptWrapImports = true;
   LLVMStyle.TabWidth = 8;
   LLVMStyle.MaxEmptyLinesToKeep = 1;
FormatStyle getGoogleStyle(FormatStyle::LanguageKind Language) {
   GoogleStyle.AccessModifierOffset = -1;
   GoogleStyle.AlignEscapedNewlines = FormatStyle::ENAS_Left;
   GoogleStyle.AllowShortIfStatementsOnASingleLine = true;
   GoogleStyle.AllowShortLoopsOnASingleLine = true;
   GoogleStyle.AlwaysBreakBeforeMultilineStrings = true;
   GoogleStyle.AlwaysBreakTemplateDeclarations = true;
   GoogleStyle.ConstructorInitializerAllOnOneLineOrOnePerLine = true;
   GoogleStyle.DerivePointerAlignment = true;
-  GoogleStyle.IncludeCategories = {
+  GoogleStyle.IncludeStyle.IncludeCategories = {
       {"^<ext/.*\\.h>", 2}, {"^<.*\\.h>", 1}, {"^<.*", 2}, {".*", 3}};
-  GoogleStyle.IncludeIsMainRegex = "([-_](test|unittest))?$";
+  GoogleStyle.IncludeStyle.IncludeIsMainRegex = "([-_](test|unittest))?$";
   GoogleStyle.IndentCaseLabels = true;
   GoogleStyle.KeepEmptyLinesAtTheStartOfBlocks = false;
   GoogleStyle.ObjCBinPackProtocolList = FormatStyle::BPS_Never;
   GoogleStyle.ObjCSpaceAfterProperty = false;
   GoogleStyle.ObjCSpaceBeforeProtocolList = true;
   GoogleStyle.PointerAlignment = FormatStyle::PAS_Left;
   GoogleStyle.RawStringFormats = {{
       FormatStyle::LK_TextProto,
Indices.erase(std::unique(Indices.begin(), Indices.end(),
                 Indices.end());
 
   int CurrentCategory = Includes.front().Category;
 
   // If the #includes are out of order, we generate a single replacement fixing
   // the entire block. Otherwise, no replacement is generated.
   if (Indices.size() == Includes.size() &&
       std::is_sorted(Indices.begin(), Indices.end()) &&
-      Style.IncludeBlocks == FormatStyle::IBS_Preserve)
+      Style.IncludeStyle.IncludeBlocks == tooling::IncludeStyle::IBS_Preserve)
     return;
 
   std::string result;
   for (unsigned Index : Indices) {
     if (!result.empty()) {
       result += "\n";
-      if (Style.IncludeBlocks == FormatStyle::IBS_Regroup &&
+      if (Style.IncludeStyle.IncludeBlocks ==
+              tooling::IncludeStyle::IBS_Regroup &&
           CurrentCategory != Includes[Index].Category)
         result += "\n";
     }
     result += Includes[Index].Text;
     if (Cursor && CursorIndex == Index)
       *Cursor = IncludesBeginOffset + result.size() - CursorToEOLOffset;
     CurrentCategory = Includes[Index].Category;
   }
 
   auto Err = Replaces.add(tooling::Replacement(
       FileName, Includes.front().Offset, IncludesBlockSize, result));
   // FIXME: better error handling. For now, just skip the replacement for the
   // release version.
   if (Err) {
     llvm::errs() << llvm::toString(std::move(Err)) << "\n";
     assert(false);
   }
 }
 
 namespace {
 
-// This class manages priorities of #include categories and calculates
-// priorities for headers.
-class IncludeCategoryManager {
-public:
-  IncludeCategoryManager(const FormatStyle &Style, StringRef FileName)
-      : Style(Style), FileName(FileName) {
-    FileStem = llvm::sys::path::stem(FileName);
-    for (const auto &Category : Style.IncludeCategories)
-      CategoryRegexs.emplace_back(Category.Regex, llvm::Regex::IgnoreCase);
-    IsMainFile = FileName.endswith(".c") || FileName.endswith(".cc") ||
-                 FileName.endswith(".cpp") || FileName.endswith(".c++") ||
-                 FileName.endswith(".cxx") || FileName.endswith(".m") ||
-                 FileName.endswith(".mm");
-  }
-
-  // Returns the priority of the category which \p IncludeName belongs to.
-  // If \p CheckMainHeader is true and \p IncludeName is a main header, returns
-  // 0. Otherwise, returns the priority of the matching category or INT_MAX.
-  int getIncludePriority(StringRef IncludeName, bool CheckMainHeader) {
-    int Ret = INT_MAX;
-    for (unsigned i = 0, e = CategoryRegexs.size(); i != e; ++i)
-      if (CategoryRegexs[i].match(IncludeName)) {
-        Ret = Style.IncludeCategories[i].Priority;
-        break;
-      }
-    if (CheckMainHeader && IsMainFile && Ret > 0 && isMainHeader(IncludeName))
-      Ret = 0;
-    return Ret;
-  }
-
-private:
-  bool isMainHeader(StringRef IncludeName) const {
-    if (!IncludeName.startswith("\""))
-      return false;
-    StringRef HeaderStem =
-        llvm::sys::path::stem(IncludeName.drop_front(1).drop_back(1));
-    if (FileStem.startswith(HeaderStem) ||
-        FileStem.startswith_lower(HeaderStem)) {
-      llvm::Regex MainIncludeRegex(
-          (HeaderStem + Style.IncludeIsMainRegex).str(),
-          llvm::Regex::IgnoreCase);
-      if (MainIncludeRegex.match(FileStem))
-        return true;
-    }
-    return false;
-  }
-
-  const FormatStyle &Style;
-  bool IsMainFile;
-  StringRef FileName;
-  StringRef FileStem;
-  SmallVector<llvm::Regex, 4> CategoryRegexs;
-};
-
 const char IncludeRegexPattern[] =
     R"(^[\t\ ]*#[\t\ ]*(import|include)[^"<]*(["<][^">]*[">]))";
 
 } // anonymous namespace
 
 tooling::Replacements sortCppIncludes(const FormatStyle &Style, StringRef Code,
                                       ArrayRef<tooling::Range> Ranges,
                                       StringRef FileName,
                                       tooling::Replacements &Replaces,
                                       unsigned *Cursor) {
   unsigned Prev = 0;
   unsigned SearchFrom = 0;
   llvm::Regex IncludeRegex(IncludeRegexPattern);
   SmallVector<StringRef, 4> Matches;
   SmallVector<IncludeDirective, 16> IncludesInBlock;
 
   // In compiled files, consider the first #include to be the main #include of
   // the file if it is not a system #include. This ensures that the header
   // doesn't have hidden dependencies
   // (http://llvm.org/docs/CodingStandards.html#include-style).
   //
   // FIXME: Do some sanity checking, e.g. edit distance of the base name, to fix
   // cases where the first #include is unlikely to be the main header.
-  IncludeCategoryManager Categories(Style, FileName);
+  tooling::IncludeCategoryManager Categories(Style.IncludeStyle, FileName);
   bool FirstIncludeBlock = true;
   bool MainIncludeFound = false;
   bool FormattingOff = false;
 
   for (;;) {
     auto Pos = Code.find('\n', SearchFrom);
     StringRef Line =
         Code.substr(Prev, (Pos != StringRef::npos ? Pos : Code.size()) - Prev);
 
     StringRef Trimmed = Line.trim();
     if (Trimmed == "// clang-format off")
       FormattingOff = true;
     else if (Trimmed == "// clang-format on")
       FormattingOff = false;
 
     const bool EmptyLineSkipped =
-        Trimmed.empty() && (Style.IncludeBlocks == FormatStyle::IBS_Merge ||
-                            Style.IncludeBlocks == FormatStyle::IBS_Regroup);
+        Trimmed.empty() &&
+        (Style.IncludeStyle.IncludeBlocks == tooling::IncludeStyle::IBS_Merge ||
+         Style.IncludeStyle.IncludeBlocks ==
+             tooling::IncludeStyle::IBS_Regroup);
 
     if (!FormattingOff && !Line.endswith("\\")) {
       if (IncludeRegex.match(Line, &Matches)) {
         StringRef IncludeName = Matches[2];
         int Category = Categories.getIncludePriority(
             IncludeName,
             /*CheckMainHeader=*/!MainIncludeFound && FirstIncludeBlock);
         if (Category == 0)
inline bool isHeaderInsertion(const tooling::Replacement &Replace) {
   return Replace.getOffset() == UINT_MAX && Replace.getLength() == 0 &&
          llvm::Regex(IncludeRegexPattern).match(Replace.getReplacementText());
 }
 
 inline bool isHeaderDeletion(const tooling::Replacement &Replace) {
   return Replace.getOffset() == UINT_MAX && Replace.getLength() == 1;
 }
 
-// Returns the offset after skipping a sequence of tokens, matched by \p
-// GetOffsetAfterSequence, from the start of the code.
-// \p GetOffsetAfterSequence should be a function that matches a sequence of
-// tokens and returns an offset after the sequence.
-unsigned getOffsetAfterTokenSequence(
-    StringRef FileName, StringRef Code, const FormatStyle &Style,
-    llvm::function_ref<unsigned(const SourceManager &, Lexer &, Token &)>
-        GetOffsetAfterSequence) {
-  std::unique_ptr<Environment> Env =
-      Environment::CreateVirtualEnvironment(Code, FileName, /*Ranges=*/{});
-  const SourceManager &SourceMgr = Env->getSourceManager();
-  Lexer Lex(Env->getFileID(), SourceMgr.getBuffer(Env->getFileID()), SourceMgr,
-            getFormattingLangOpts(Style));
-  Token Tok;
-  // Get the first token.
-  Lex.LexFromRawLexer(Tok);
-  return GetOffsetAfterSequence(SourceMgr, Lex, Tok);
-}
-
-// Check if a sequence of tokens is like "#<Name> <raw_identifier>". If it is,
-// \p Tok will be the token after this directive; otherwise, it can be any token
-// after the given \p Tok (including \p Tok).
-bool checkAndConsumeDirectiveWithName(Lexer &Lex, StringRef Name, Token &Tok) {
-  bool Matched = Tok.is(tok::hash) && !Lex.LexFromRawLexer(Tok) &&
-                 Tok.is(tok::raw_identifier) &&
-                 Tok.getRawIdentifier() == Name && !Lex.LexFromRawLexer(Tok) &&
-                 Tok.is(tok::raw_identifier);
-  if (Matched)
-    Lex.LexFromRawLexer(Tok);
-  return Matched;
-}
-
-void skipComments(Lexer &Lex, Token &Tok) {
-  while (Tok.is(tok::comment))
-    if (Lex.LexFromRawLexer(Tok))
-      return;
-}
-
-// Returns the offset after header guard directives and any comments
-// before/after header guards. If no header guard presents in the code, this
-// will returns the offset after skipping all comments from the start of the
-// code.
-unsigned getOffsetAfterHeaderGuardsAndComments(StringRef FileName,
-                                               StringRef Code,
-                                               const FormatStyle &Style) {
-  return getOffsetAfterTokenSequence(
-      FileName, Code, Style,
-      [](const SourceManager &SM, Lexer &Lex, Token Tok) {
-        skipComments(Lex, Tok);
-        unsigned InitialOffset = SM.getFileOffset(Tok.getLocation());
-        if (checkAndConsumeDirectiveWithName(Lex, "ifndef", Tok)) {
-          skipComments(Lex, Tok);
-          if (checkAndConsumeDirectiveWithName(Lex, "define", Tok))
-            return SM.getFileOffset(Tok.getLocation());
-        }
-        return InitialOffset;
-      });
-}
-
-// Check if a sequence of tokens is like
-//    "#include ("header.h" | <header.h>)".
-// If it is, \p Tok will be the token after this directive; otherwise, it can be
-// any token after the given \p Tok (including \p Tok).
-bool checkAndConsumeInclusiveDirective(Lexer &Lex, Token &Tok) {
-  auto Matched = [&]() {
-    Lex.LexFromRawLexer(Tok);
-    return true;
-  };
-  if (Tok.is(tok::hash) && !Lex.LexFromRawLexer(Tok) &&
-      Tok.is(tok::raw_identifier) && Tok.getRawIdentifier() == "include") {
-    if (Lex.LexFromRawLexer(Tok))
-      return false;
-    if (Tok.is(tok::string_literal))
-      return Matched();
-    if (Tok.is(tok::less)) {
-      while (!Lex.LexFromRawLexer(Tok) && Tok.isNot(tok::greater)) {
-      }
-      if (Tok.is(tok::greater))
-        return Matched();
-    }
-  }
-  return false;
-}
-
-// Returns the offset of the last #include directive after which a new
-// #include can be inserted. This ignores #include's after the #include block(s)
-// in the beginning of a file to avoid inserting headers into code sections
-// where new #include's should not be added by default.
-// These code sections include:
-//      - raw string literals (containing #include).
-//      - #if blocks.
-//      - Special #include's among declarations (e.g. functions).
-//
-// If no #include after which a new #include can be inserted, this returns the
-// offset after skipping all comments from the start of the code.
-// Inserting after an #include is not allowed if it comes after code that is not
-// #include (e.g. pre-processing directive that is not #include, declarations).
-unsigned getMaxHeaderInsertionOffset(StringRef FileName, StringRef Code,
-                                     const FormatStyle &Style) {
-  return getOffsetAfterTokenSequence(
-      FileName, Code, Style,
-      [](const SourceManager &SM, Lexer &Lex, Token Tok) {
-        skipComments(Lex, Tok);
-        unsigned MaxOffset = SM.getFileOffset(Tok.getLocation());
-        while (checkAndConsumeInclusiveDirective(Lex, Tok))
-          MaxOffset = SM.getFileOffset(Tok.getLocation());
-        return MaxOffset;
-      });
-}
-
-bool isDeletedHeader(llvm::StringRef HeaderName,
-                     const std::set<llvm::StringRef> &HeadersToDelete) {
-  return HeadersToDelete.count(HeaderName) ||
-         HeadersToDelete.count(HeaderName.trim("\"<>"));
-}
-
 // FIXME: insert empty lines between newly created blocks.
 tooling::Replacements
 fixCppIncludeInsertions(StringRef Code, const tooling::Replacements &Replaces,
                         const FormatStyle &Style) {
   if (!Style.isCpp())
     return Replaces;
 
   tooling::Replacements HeaderInsertions;
if (isHeaderInsertion(R)) {
                    << R.getReplacementText() << "\n";
     } else {
       llvm::consumeError(Result.add(R));
     }
   }
   if (HeaderInsertions.empty() && HeadersToDelete.empty())
     return Replaces;
 
-  llvm::Regex IncludeRegex(IncludeRegexPattern);
-  llvm::Regex DefineRegex(R"(^[\t\ ]*#[\t\ ]*define[\t\ ]*[^\\]*$)");
-  SmallVector<StringRef, 4> Matches;
 
   StringRef FileName = Replaces.begin()->getFilePath();
-  IncludeCategoryManager Categories(Style, FileName);
+  tooling::HeaderIncludes Includes(FileName, Code, Style.IncludeStyle);
 
-  // Record the offset of the end of the last include in each category.
-  std::map<int, int> CategoryEndOffsets;
-  // All possible priorities.
-  // Add 0 for main header and INT_MAX for headers that are not in any category.
-  std::set<int> Priorities = {0, INT_MAX};
-  for (const auto &Category : Style.IncludeCategories)
-    Priorities.insert(Category.Priority);
-  int FirstIncludeOffset = -1;
-  // All new headers should be inserted after this offset.
-  unsigned MinInsertOffset =
-      getOffsetAfterHeaderGuardsAndComments(FileName, Code, Style);
-  StringRef TrimmedCode = Code.drop_front(MinInsertOffset);
-  // Max insertion offset in the original code.
-  unsigned MaxInsertOffset =
-      MinInsertOffset +
-      getMaxHeaderInsertionOffset(FileName, TrimmedCode, Style);
-  SmallVector<StringRef, 32> Lines;
-  TrimmedCode.split(Lines, '\n');
-  unsigned Offset = MinInsertOffset;
-  unsigned NextLineOffset;
-  std::set<StringRef> ExistingIncludes;
-  for (auto Line : Lines) {
-    NextLineOffset = std::min(Code.size(), Offset + Line.size() + 1);
-    if (IncludeRegex.match(Line, &Matches)) {
-      // The header name with quotes or angle brackets.
-      StringRef IncludeName = Matches[2];
-      ExistingIncludes.insert(IncludeName);
-      // Only record the offset of current #include if we can insert after it.
-      if (Offset <= MaxInsertOffset) {
-        int Category = Categories.getIncludePriority(
-            IncludeName, /*CheckMainHeader=*/FirstIncludeOffset < 0);
-        CategoryEndOffsets[Category] = NextLineOffset;
-        if (FirstIncludeOffset < 0)
-          FirstIncludeOffset = Offset;
-      }
-      if (isDeletedHeader(IncludeName, HeadersToDelete)) {
-        // If this is the last line without trailing newline, we need to make
-        // sure we don't delete across the file boundary.
-        unsigned Length = std::min(Line.size() + 1, Code.size() - Offset);
-        llvm::Error Err =
-            Result.add(tooling::Replacement(FileName, Offset, Length, ""));
-        if (Err) {
-          // Ignore the deletion on conflict.
-          llvm::errs() << "Failed to add header deletion replacement for "
-                       << IncludeName << ": " << llvm::toString(std::move(Err))
-                       << "\n";
-        }
-      }
-    }
-    Offset = NextLineOffset;
-  }
-
-  // Populate CategoryEndOfssets:
-  // - Ensure that CategoryEndOffset[Highest] is always populated.
-  // - If CategoryEndOffset[Priority] isn't set, use the next higher value that
-  //   is set, up to CategoryEndOffset[Highest].
-  auto Highest = Priorities.begin();
-  if (CategoryEndOffsets.find(*Highest) == CategoryEndOffsets.end()) {
-    if (FirstIncludeOffset >= 0)
-      CategoryEndOffsets[*Highest] = FirstIncludeOffset;
-    else
-      CategoryEndOffsets[*Highest] = MinInsertOffset;
-  }
-  // By this point, CategoryEndOffset[Highest] is always set appropriately:
-  //  - to an appropriate location before/after existing #includes, or
-  //  - to right after the header guard, or
-  //  - to the beginning of the file.
-  for (auto I = ++Priorities.begin(), E = Priorities.end(); I != E; ++I)
-    if (CategoryEndOffsets.find(*I) == CategoryEndOffsets.end())
-      CategoryEndOffsets[*I] = CategoryEndOffsets[*std::prev(I)];
+  for (const auto &Header : HeadersToDelete) {
+    tooling::Replacements Replaces = Includes.remove(Header);
+    for (const auto &R : Replaces) {
+      auto Err = Result.add(R);
+                     << Header << ": " << llvm::toString(std::move(Err))
 
-  bool NeedNewLineAtEnd = !Code.empty() && Code.back() != '\n';
+  llvm::Regex IncludeRegex = llvm::Regex(IncludeRegexPattern);
+  llvm::SmallVector<StringRef, 4> Matches;
   for (const auto &R : HeaderInsertions) {
     auto IncludeDirective = R.getReplacementText();
     bool Matched = IncludeRegex.match(IncludeDirective, &Matches);
     assert(Matched && "Header insertion replacement must have replacement text "
                       "'#include ...'");
     (void)Matched;
     auto IncludeName = Matches[2];
-    if (ExistingIncludes.find(IncludeName) != ExistingIncludes.end()) {
-      DEBUG(llvm::dbgs() << "Skip adding existing include : " << IncludeName
-                         << "\n");
-      continue;
-    }
-    int Category =
-        Categories.getIncludePriority(IncludeName, /*CheckMainHeader=*/true);
-    Offset = CategoryEndOffsets[Category];
-    std::string NewInclude = !IncludeDirective.endswith("\n")
-                                 ? (IncludeDirective + "\n").str()
-                                 : IncludeDirective.str();
-    // When inserting headers at end of the code, also append '\n' to the code
-    // if it does not end with '\n'.
-    if (NeedNewLineAtEnd && Offset == Code.size()) {
-      NewInclude = "\n" + NewInclude;
-      NeedNewLineAtEnd = false;
-    }
-    auto NewReplace = tooling::Replacement(FileName, Offset, 0, NewInclude);
-    auto Err = Result.add(NewReplace);
+    auto Replace = Includes.insert(IncludeName);
+    if (Replace) {
+      auto Err = Result.add(*Replace);
       if (Err) {
         llvm::consumeError(std::move(Err));
-      unsigned NewOffset = Result.getShiftedCodePosition(Offset);
-      NewReplace = tooling::Replacement(FileName, NewOffset, 0, NewInclude);
-      Result = Result.merge(tooling::Replacements(NewReplace));
+        unsigned NewOffset = Result.getShiftedCodePosition(Replace->getOffset());
+        auto Shifted = tooling::Replacement(FileName, NewOffset, 0,
+                                            Replace->getReplacementText());
+        Result = Result.merge(tooling::Replacements(Shifted));
+      }
     }
   }
   return Result;
 }
 
 } // anonymous namespace
 
 llvm::Expected<tooling::Replacements>
if (Style.Language == FormatStyle::LK_JavaScript &&
     Passes.emplace_back([&](const Environment &Env) {
       return JavaScriptRequoter(Env, Expanded).process();
     });
 
   Passes.emplace_back([&](const Environment &Env) {
     return Formatter(Env, Expanded, Status).process();
   });
 
-  std::unique_ptr<Environment> Env = Environment::CreateVirtualEnvironment(
-      Code, FileName, Ranges, FirstStartColumn, NextStartColumn,
+  Environment Env(Code, FileName, Ranges, FirstStartColumn, NextStartColumn,
                   LastStartColumn);
   llvm::Optional<std::string> CurrentCode = None;
   tooling::Replacements Fixes;
   unsigned Penalty = 0;
   for (size_t I = 0, E = Passes.size(); I < E; ++I) {
-    std::pair<tooling::Replacements, unsigned> PassFixes = Passes[I](*Env);
+    std::pair<tooling::Replacements, unsigned> PassFixes = Passes[I](Env);
     auto NewCode = applyAllReplacements(
         CurrentCode ? StringRef(*CurrentCode) : Code, PassFixes.first);
     if (NewCode) {
       Fixes = Fixes.merge(PassFixes.first);
       Penalty += PassFixes.second;
       if (I + 1 < E) {
         CurrentCode = std::move(*NewCode);
-        Env = Environment::CreateVirtualEnvironment(
+        Env = Environment(
             *CurrentCode, FileName,
             tooling::calculateRangesAfterReplacements(Fixes, Ranges),
             FirstStartColumn, NextStartColumn, LastStartColumn);
       }
     }
   }
 
   return {Fixes, Penalty};

 }
 
 tooling::Replacements cleanup(const FormatStyle &Style, StringRef Code,
                               ArrayRef<tooling::Range> Ranges,
                               StringRef FileName) {
   // cleanups only apply to C++ (they mostly concern ctor commas etc.)
   if (Style.Language != FormatStyle::LK_Cpp)
     return tooling::Replacements();
-  std::unique_ptr<Environment> Env =
-      Environment::CreateVirtualEnvironment(Code, FileName, Ranges);
-  Cleaner Clean(*Env, Style);
-  return Clean.process().first;
+  return Cleaner(Environment(Code, FileName, Ranges), Style).process().first;
 }
 
 tooling::Replacements reformat(const FormatStyle &Style, StringRef Code,
                                ArrayRef<tooling::Range> Ranges,
                                StringRef FileName, bool *IncompleteFormat) {
   FormattingAttemptStatus Status;
   auto Result = reformat(Style, Code, Ranges, FileName, &Status);
   if (!Status.FormatComplete)
     *IncompleteFormat = true;
   return Result;
 }
 
 tooling::Replacements fixNamespaceEndComments(const FormatStyle &Style,
                                               StringRef Code,
                                               ArrayRef<tooling::Range> Ranges,
                                               StringRef FileName) {
-  std::unique_ptr<Environment> Env =
-      Environment::CreateVirtualEnvironment(Code, FileName, Ranges);
-  NamespaceEndCommentsFixer Fix(*Env, Style);
-  return Fix.process().first;
+  return NamespaceEndCommentsFixer(Environment(Code, FileName, Ranges), Style)
+      .process()
+      .first;
 }
 
 tooling::Replacements sortUsingDeclarations(const FormatStyle &Style,
                                             StringRef Code,
                                             ArrayRef<tooling::Range> Ranges,
                                             StringRef FileName) {
-  std::unique_ptr<Environment> Env =
-      Environment::CreateVirtualEnvironment(Code, FileName, Ranges);
-  UsingDeclarationsSorter Sorter(*Env, Style);
-  return Sorter.process().first;
+  return UsingDeclarationsSorter(Environment(Code, FileName, Ranges), Style)
+      .process()
+      .first;
 }
 
 LangOptions getFormattingLangOpts(const FormatStyle &Style) {
   LangOptions LangOpts;
   LangOpts.CPlusPlus = 1;
   LangOpts.CPlusPlus11 = Style.Standard == FormatStyle::LS_Cpp03 ? 0 : 1;
   LangOpts.CPlusPlus14 = Style.Standard == FormatStyle::LS_Cpp03 ? 0 : 1;
   LangOpts.CPlusPlus17 = Style.Standard == FormatStyle::LS_Cpp03 ? 0 : 1;

 FormatStyle::LanguageKind guessLanguage(StringRef FileName, StringRef Code) {
   const auto GuessedLanguage = getLanguageByFileName(FileName);
   if (GuessedLanguage == FormatStyle::LK_Cpp) {
     auto Extension = llvm::sys::path::extension(FileName);
     // If there's no file extension (or it's .h), we need to check the contents
     // of the code to see if it contains Objective-C.
     if (Extension.empty() || Extension == ".h") {
       auto NonEmptyFileName = FileName.empty() ? "guess.h" : FileName;
-      std::unique_ptr<Environment> Env =
-          Environment::CreateVirtualEnvironment(Code, NonEmptyFileName, /*Ranges=*/{});
-      ObjCHeaderStyleGuesser Guesser(*Env, getLLVMStyle());
+      Environment Env(Code, NonEmptyFileName, /*Ranges=*/{});
+      ObjCHeaderStyleGuesser Guesser(Env, getLLVMStyle());
       Guesser.process();
       if (Guesser.isObjC())
         return FormatStyle::LK_ObjC;
     }
   }
   return GuessedLanguage;
 }
 

lib/Format/SortJavaScriptImports.cpp

private:
   }
 };
 
 tooling::Replacements sortJavaScriptImports(const FormatStyle &Style,
                                             StringRef Code,
                                             ArrayRef<tooling::Range> Ranges,
                                             StringRef FileName) {
   // FIXME: Cursor support.
-  std::unique_ptr<Environment> Env =
-      Environment::CreateVirtualEnvironment(Code, FileName, Ranges);
-  JavaScriptImportSorter Sorter(*Env, Style);
-  return Sorter.process().first;
+  return JavaScriptImportSorter(Environment(Code, FileName, Ranges), Style)
+      .process()
+      .first;
 }
 
 } // end namespace format
 } // end namespace clang

lib/Format/TokenAnalyzer.h

 #include "FormatTokenLexer.h"
 #include "TokenAnnotator.h"
 #include "UnwrappedLineParser.h"
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/DiagnosticOptions.h"
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Format/Format.h"
+#include "clang/Tooling/Core/Environment.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
 
 namespace clang {
 namespace format {
 
 class Environment {
 public:
   Environment(SourceManager &SM, FileID ID, ArrayRef<CharSourceRange> Ranges)
-      : ID(ID), CharRanges(Ranges.begin(), Ranges.end()), SM(SM),
-      FirstStartColumn(0),
-      NextStartColumn(0),
-      LastStartColumn(0) {}
-
-  Environment(FileID ID, std::unique_ptr<FileManager> FileMgr,
-              std::unique_ptr<SourceManager> VirtualSM,
-              std::unique_ptr<DiagnosticsEngine> Diagnostics,
-              const std::vector<CharSourceRange> &CharRanges,
-              unsigned FirstStartColumn,
-              unsigned NextStartColumn,
-              unsigned LastStartColumn)
-      : ID(ID), CharRanges(CharRanges.begin(), CharRanges.end()),
-        SM(*VirtualSM), 
-        FirstStartColumn(FirstStartColumn),
-        NextStartColumn(NextStartColumn),
-        LastStartColumn(LastStartColumn),
-        FileMgr(std::move(FileMgr)),
-        VirtualSM(std::move(VirtualSM)), Diagnostics(std::move(Diagnostics)) {}
+      : ToolEnv(new tooling::Environment(SM, ID)),
+        CharRanges(Ranges.begin(), Ranges.end()), FirstStartColumn(0),
+        NextStartColumn(0), LastStartColumn(0) {}
 
   // This sets up an virtual file system with file \p FileName containing the
   // fragment \p Code. Assumes that \p Code starts at \p FirstStartColumn,
   // that the next lines of \p Code should start at \p NextStartColumn, and
   // that \p Code should end at \p LastStartColumn if it ends in newline.
   // See also the documentation of clang::format::internal::reformat.
-  static std::unique_ptr<Environment>
-  CreateVirtualEnvironment(StringRef Code, StringRef FileName,
-                           ArrayRef<tooling::Range> Ranges,
-                           unsigned FirstStartColumn = 0,
-                           unsigned NextStartColumn = 0,
-                           unsigned LastStartColumn = 0);
-
-  FileID getFileID() const { return ID; }
+  Environment(StringRef Code, StringRef FileName,
+              ArrayRef<tooling::Range> Ranges, unsigned FirstStartColumn = 0,
+              unsigned NextStartColumn = 0, unsigned LastStartColumn = 0);
+
+  FileID getFileID() const { return ToolEnv->getFileID(); }
+
+  const SourceManager &getSourceManager() const {
+    return ToolEnv->getSourceManager();
+  }
 
   ArrayRef<CharSourceRange> getCharRanges() const { return CharRanges; }
 
-  const SourceManager &getSourceManager() const { return SM; }
-
   // Returns the column at which the fragment of code managed by this
   // environment starts.
   unsigned getFirstStartColumn() const { return FirstStartColumn; }
 
   // Returns the column at which subsequent lines of the fragment of code
   // managed by this environment should start.
   unsigned getNextStartColumn() const { return NextStartColumn; }
 
   // Returns the column at which the fragment of code managed by this
   // environment should end if it ends in a newline.
   unsigned getLastStartColumn() const { return LastStartColumn; }
 
 private:
-  FileID ID;
+  std::unique_ptr<tooling::Environment> ToolEnv;
   SmallVector<CharSourceRange, 8> CharRanges;
-  SourceManager &SM;
   unsigned FirstStartColumn;
   unsigned NextStartColumn;
   unsigned LastStartColumn;
-
-  // The order of these fields are important - they should be in the same order
-  // as they are created in `CreateVirtualEnvironment` so that they can be
-  // deleted in the reverse order as they are created.
-  std::unique_ptr<FileManager> FileMgr;
-  std::unique_ptr<SourceManager> VirtualSM;
-  std::unique_ptr<DiagnosticsEngine> Diagnostics;
 };
 
 class TokenAnalyzer : public UnwrappedLineConsumer {
 public:
   TokenAnalyzer(const Environment &Env, const FormatStyle &Style);
 
   std::pair<tooling::Replacements, unsigned> process();
 

lib/Format/TokenAnalyzer.cpp

 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "format-formatter"
 
 namespace clang {
 namespace format {
 
-// This sets up an virtual file system with file \p FileName containing \p
-// Code.
-std::unique_ptr<Environment>
-Environment::CreateVirtualEnvironment(StringRef Code, StringRef FileName,
-                                      ArrayRef<tooling::Range> Ranges,
-                                      unsigned FirstStartColumn,
-                                      unsigned NextStartColumn,
-                                      unsigned LastStartColumn) {
-  // This is referenced by `FileMgr` and will be released by `FileMgr` when it
-  // is deleted.
-  IntrusiveRefCntPtr<vfs::InMemoryFileSystem> InMemoryFileSystem(
-      new vfs::InMemoryFileSystem);
-  // This is passed to `SM` as reference, so the pointer has to be referenced
-  // in `Environment` so that `FileMgr` can out-live this function scope.
-  std::unique_ptr<FileManager> FileMgr(
-      new FileManager(FileSystemOptions(), InMemoryFileSystem));
-  // This is passed to `SM` as reference, so the pointer has to be referenced
-  // by `Environment` due to the same reason above.
-  std::unique_ptr<DiagnosticsEngine> Diagnostics(new DiagnosticsEngine(
-      IntrusiveRefCntPtr<DiagnosticIDs>(new DiagnosticIDs),
-      new DiagnosticOptions));
-  // This will be stored as reference, so the pointer has to be stored in
-  // due to the same reason above.
-  std::unique_ptr<SourceManager> VirtualSM(
-      new SourceManager(*Diagnostics, *FileMgr));
-  InMemoryFileSystem->addFile(
-      FileName, 0,
-      llvm::MemoryBuffer::getMemBuffer(Code, FileName,
-                                       /*RequiresNullTerminator=*/false));
-  FileID ID = VirtualSM->createFileID(FileMgr->getFile(FileName),
-                                      SourceLocation(), clang::SrcMgr::C_User);
-  assert(ID.isValid());
-  SourceLocation StartOfFile = VirtualSM->getLocForStartOfFile(ID);
-  std::vector<CharSourceRange> CharRanges;
+Environment::Environment(StringRef Code, StringRef FileName,
+                         unsigned FirstStartColumn, unsigned NextStartColumn,
+                         unsigned LastStartColumn)
+    : ToolEnv(tooling::Environment::createVirtualEnvironment(Code, FileName)),
+      FirstStartColumn(FirstStartColumn), NextStartColumn(NextStartColumn),
+      LastStartColumn(LastStartColumn) {
+  SourceLocation StartOfFile =
+      ToolEnv->getSourceManager().getLocForStartOfFile(ToolEnv->getFileID());
   for (const tooling::Range &Range : Ranges) {
     SourceLocation Start = StartOfFile.getLocWithOffset(Range.getOffset());
     SourceLocation End = Start.getLocWithOffset(Range.getLength());
     CharRanges.push_back(CharSourceRange::getCharRange(Start, End));
   }
-  return llvm::make_unique<Environment>(
-      ID, std::move(FileMgr), std::move(VirtualSM), std::move(Diagnostics),
-      CharRanges, FirstStartColumn, NextStartColumn, LastStartColumn);
 }
 
 TokenAnalyzer::TokenAnalyzer(const Environment &Env, const FormatStyle &Style)
     : Style(Style), Env(Env),
       AffectedRangeMgr(Env.getSourceManager(), Env.getCharRanges()),
       UnwrappedLines(1),
       Encoding(encoding::detectEncoding(
           Env.getSourceManager().getBufferData(Env.getFileID()))) {

lib/Tooling/Core/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS support)
 
 add_clang_library(clangToolingCore
+  Diagnostic.cpp
+  Environment.cpp
+  HeaderIncludes.cpp
   Lookup.cpp
   Replacement.cpp
-  Diagnostic.cpp
 
   LINK_LIBS
   clangAST
   clangBasic
   clangLex
   clangRewrite
   )

lib/Tooling/Core/Environment.cpp

+//===--- Environment.cpp - Sourece tooling environment --*- C++ -*---------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/Tooling/Core/Environment.h"
+
+namespace clang {
+namespace tooling {
+
+std::unique_ptr<Environment>
+Environment::createVirtualEnvironment(StringRef Content, StringRef FileName) {
+  // This is referenced by `FileMgr` and will be released by `FileMgr` when it
+  // is deleted.
+  IntrusiveRefCntPtr<vfs::InMemoryFileSystem> InMemoryFileSystem(
+      new vfs::InMemoryFileSystem);
+  // This is passed to `SM` as reference, so the pointer has to be referenced
+  // in `Environment` so that `FileMgr` can out-live this function scope.
+  std::unique_ptr<FileManager> FileMgr(
+      new FileManager(FileSystemOptions(), InMemoryFileSystem));
+  // This is passed to `SM` as reference, so the pointer has to be referenced
+  // by `Environment` due to the same reason above.
+  std::unique_ptr<DiagnosticsEngine> Diagnostics(new DiagnosticsEngine(
+      IntrusiveRefCntPtr<DiagnosticIDs>(new DiagnosticIDs),
+      new DiagnosticOptions));
+  // This will be stored as reference, so the pointer has to be stored in
+  // due to the same reason above.
+  std::unique_ptr<SourceManager> VirtualSM(
+      new SourceManager(*Diagnostics, *FileMgr));
+  InMemoryFileSystem->addFile(
+      FileName, 0,
+      llvm::MemoryBuffer::getMemBuffer(Content, FileName,
+                                       /*RequiresNullTerminator=*/false));
+  FileID ID = VirtualSM->createFileID(FileMgr->getFile(FileName),
+                                      SourceLocation(), clang::SrcMgr::C_User);
+  assert(ID.isValid());
+
+  return llvm::make_unique<Environment>(
+      ID, std::move(FileMgr), std::move(VirtualSM), std::move(Diagnostics));
+}
+
+} // namespace tooling
+} // namespace clang

lib/Tooling/Core/HeaderIncludes.cpp

+//===--- HeaderIncludes.cpp - Insert/Delete #includes --*- C++ -*----------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/Tooling/Core/HeaderIncludes.h"
+#include "clang/Lex/Lexer.h"
+
+namespace clang {
+namespace tooling {
+namespace {
+
+LangOptions createLangOpts() {
+  LangOptions LangOpts;
+  LangOpts.CPlusPlus = 1;
+  LangOpts.CPlusPlus11 = 1;
+  LangOpts.CPlusPlus14 = 1;
+  LangOpts.LineComment = 1;
+  LangOpts.CXXOperatorNames = 1;
+  LangOpts.Bool = 1;
+  LangOpts.ObjC1 = 1;
+  LangOpts.ObjC2 = 1;
+  LangOpts.MicrosoftExt = 1;    // To get kw___try, kw___finally.
+  LangOpts.DeclSpecKeyword = 1; // To get __declspec.
+  LangOpts.WChar = 1;           // To get wchar_t
+  return LangOpts;
+}
+
+// Returns the offset after skipping a sequence of tokens, matched by \p
+// GetOffsetAfterSequence, from the start of the code.
+// \p GetOffsetAfterSequence should be a function that matches a sequence of
+// tokens and returns an offset after the sequence.
+unsigned getOffsetAfterTokenSequence(
+    StringRef FileName, StringRef Code, const IncludeStyle &Style,
+    llvm::function_ref<unsigned(const SourceManager &, Lexer &, Token &)>
+        GetOffsetAfterSequence) {
+  std::unique_ptr<Environment> Env =
+      Environment::createVirtualEnvironment(Code, FileName);
+  const SourceManager &SourceMgr = Env->getSourceManager();
+  Lexer Lex(Env->getFileID(), SourceMgr.getBuffer(Env->getFileID()), SourceMgr,
+            createLangOpts());
+  Token Tok;
+  // Get the first token.
+  Lex.LexFromRawLexer(Tok);
+  return GetOffsetAfterSequence(SourceMgr, Lex, Tok);
+}
+
+// Check if a sequence of tokens is like "#<Name> <raw_identifier>". If it is,
+// \p Tok will be the token after this directive; otherwise, it can be any token
+// after the given \p Tok (including \p Tok).
+bool checkAndConsumeDirectiveWithName(Lexer &Lex, StringRef Name, Token &Tok) {
+  bool Matched = Tok.is(tok::hash) && !Lex.LexFromRawLexer(Tok) &&
+                 Tok.is(tok::raw_identifier) &&
+                 Tok.getRawIdentifier() == Name && !Lex.LexFromRawLexer(Tok) &&
+                 Tok.is(tok::raw_identifier);
+  if (Matched)
+    Lex.LexFromRawLexer(Tok);
+  return Matched;
+}
+
+void skipComments(Lexer &Lex, Token &Tok) {
+  while (Tok.is(tok::comment))
+    if (Lex.LexFromRawLexer(Tok))
+      return;
+}
+
+// Returns the offset after header guard directives and any comments
+// before/after header guards. If no header guard presents in the code, this
+// will returns the offset after skipping all comments from the start of the
+// code.
+unsigned getOffsetAfterHeaderGuardsAndComments(StringRef FileName,
+                                               StringRef Code,
+                                               const IncludeStyle &Style) {
+  return getOffsetAfterTokenSequence(
+      FileName, Code, Style,
+      [](const SourceManager &SM, Lexer &Lex, Token Tok) {
+        skipComments(Lex, Tok);
+        unsigned InitialOffset = SM.getFileOffset(Tok.getLocation());
+        if (checkAndConsumeDirectiveWithName(Lex, "ifndef", Tok)) {
+          skipComments(Lex, Tok);
+          if (checkAndConsumeDirectiveWithName(Lex, "define", Tok))
+            return SM.getFileOffset(Tok.getLocation());
+        }
+        return InitialOffset;
+      });
+}
+
+// Check if a sequence of tokens is like
+//    "#include ("header.h" | <header.h>)".
+// If it is, \p Tok will be the token after this directive; otherwise, it can be
+// any token after the given \p Tok (including \p Tok).
+bool checkAndConsumeInclusiveDirective(Lexer &Lex, Token &Tok) {
+  auto Matched = [&]() {
+    Lex.LexFromRawLexer(Tok);
+    return true;
+  };
+  if (Tok.is(tok::hash) && !Lex.LexFromRawLexer(Tok) &&
+      Tok.is(tok::raw_identifier) && Tok.getRawIdentifier() == "include") {
+    if (Lex.LexFromRawLexer(Tok))
+      return false;
+    if (Tok.is(tok::string_literal))
+      return Matched();
+    if (Tok.is(tok::less)) {
+      while (!Lex.LexFromRawLexer(Tok) && Tok.isNot(tok::greater)) {
+      }
+      if (Tok.is(tok::greater))
+        return Matched();
+    }
+  }
+  return false;
+}
+
+// Returns the offset of the last #include directive after which a new
+// #include can be inserted. This ignores #include's after the #include block(s)
+// in the beginning of a file to avoid inserting headers into code sections
+// where new #include's should not be added by default.
+// These code sections include:
+//      - raw string literals (containing #include).
+//      - #if blocks.
+//      - Special #include's among declarations (e.g. functions).
+//
+// If no #include after which a new #include can be inserted, this returns the
+// offset after skipping all comments from the start of the code.
+// Inserting after an #include is not allowed if it comes after code that is not
+// #include (e.g. pre-processing directive that is not #include, declarations).
+unsigned getMaxHeaderInsertionOffset(StringRef FileName, StringRef Code,
+                                     const IncludeStyle &Style) {
+  return getOffsetAfterTokenSequence(
+      FileName, Code, Style,
+      [](const SourceManager &SM, Lexer &Lex, Token Tok) {
+        skipComments(Lex, Tok);
+        unsigned MaxOffset = SM.getFileOffset(Tok.getLocation());
+        while (checkAndConsumeInclusiveDirective(Lex, Tok))
+          MaxOffset = SM.getFileOffset(Tok.getLocation());
+        return MaxOffset;
+      });
+}
+
+inline StringRef trimInclude(StringRef IncludeName) {
+  return IncludeName.trim("\"<>");
+}
+
+const char IncludeRegexPattern[] =
+    R"(^[\t\ ]*#[\t\ ]*(import|include)[^"<]*(["<][^">]*[">]))";
+
+} // anonymous namespace
+
+IncludeCategoryManager::IncludeCategoryManager(const IncludeStyle &Style,
+                                               StringRef FileName)
+    : Style(Style), FileName(FileName) {
+  FileStem = llvm::sys::path::stem(FileName);
+  for (const auto &Category : Style.IncludeCategories)
+    CategoryRegexs.emplace_back(Category.Regex, llvm::Regex::IgnoreCase);
+  IsMainFile = FileName.endswith(".c") || FileName.endswith(".cc") ||
+               FileName.endswith(".cpp") || FileName.endswith(".c++") ||
+               FileName.endswith(".cxx") || FileName.endswith(".m") ||
+               FileName.endswith(".mm");
+}
+
+int IncludeCategoryManager::getIncludePriority(StringRef IncludeName,
+                                               bool CheckMainHeader) const {
+  int Ret = INT_MAX;
+  for (unsigned i = 0, e = CategoryRegexs.size(); i != e; ++i)
+    if (CategoryRegexs[i].match(IncludeName)) {
+      Ret = Style.IncludeCategories[i].Priority;
+      break;
+    }
+  if (CheckMainHeader && IsMainFile && Ret > 0 && isMainHeader(IncludeName))
+    Ret = 0;
+  return Ret;
+}
+
+bool IncludeCategoryManager::isMainHeader(StringRef IncludeName) const {
+  if (!IncludeName.startswith("\""))
+    return false;
+  StringRef HeaderStem =
+      llvm::sys::path::stem(IncludeName.drop_front(1).drop_back(1));
+  if (FileStem.startswith(HeaderStem) ||
+      FileStem.startswith_lower(HeaderStem)) {
+    llvm::Regex MainIncludeRegex((HeaderStem + Style.IncludeIsMainRegex).str(),
+                                 llvm::Regex::IgnoreCase);
+    if (MainIncludeRegex.match(FileStem))
+      return true;
+  }
+  return false;
+}
+
+HeaderIncludes::HeaderIncludes(StringRef FileName, StringRef Code,
+                               const IncludeStyle &Style)
+    : FileName(FileName), Code(Code), FirstIncludeOffset(-1),
+      MinInsertOffset(
+          getOffsetAfterHeaderGuardsAndComments(FileName, Code, Style)),
+      TrimmedCode(Code.drop_front(MinInsertOffset)),
+      MaxInsertOffset(MinInsertOffset + getMaxHeaderInsertionOffset(
+                                            FileName, TrimmedCode, Style)),
+      Categories(Style, FileName),
+      IncludeRegex(llvm::Regex(IncludeRegexPattern)) {
+  // Add 0 for main header and INT_MAX for headers that are not in any
+  // category.
+  Priorities = {0, INT_MAX};
+  for (const auto &Category : Style.IncludeCategories)
+    Priorities.insert(Category.Priority);
+  SmallVector<StringRef, 32> Lines;
+  TrimmedCode.split(Lines, '\n');
+
+  unsigned Offset = MinInsertOffset;
+  unsigned NextLineOffset;
+  SmallVector<StringRef, 4> Matches;
+  for (auto Line : Lines) {
+    NextLineOffset = std::min(Code.size(), Offset + Line.size() + 1);
+    if (IncludeRegex.match(Line, &Matches)) {
+      // If this is the last line without trailing newline, we need to make
+      // sure we don't delete across the file boundary.
+      addExistingInclude(
+          Include(Matches[2], Range(Offset, std::min(Line.size() + 1,
+                                                     Code.size() - Offset))),
+          NextLineOffset);
+    }
+    Offset = NextLineOffset;
+  }
+
+  // Populate CategoryEndOfssets:
+  // - Ensure that CategoryEndOffset[Highest] is always populated.
+  // - If CategoryEndOffset[Priority] isn't set, use the next higher value
+  // that
+  //   is set, up to CategoryEndOffset[Highest].
+  auto Highest = Priorities.begin();
+  if (CategoryEndOffsets.find(*Highest) == CategoryEndOffsets.end()) {
+    if (FirstIncludeOffset >= 0)
+      CategoryEndOffsets[*Highest] = FirstIncludeOffset;
+    else
+      CategoryEndOffsets[*Highest] = MinInsertOffset;
+  }
+  // By this point, CategoryEndOffset[Highest] is always set appropriately:
+  //  - to an appropriate location before/after existing #includes, or
+  //  - to right after the header guard, or
+  //  - to the beginning of the file.
+  for (auto I = ++Priorities.begin(), E = Priorities.end(); I != E; ++I)
+    if (CategoryEndOffsets.find(*I) == CategoryEndOffsets.end())
+      CategoryEndOffsets[*I] = CategoryEndOffsets[*std::prev(I)];
+}
+
+// \p Offset: the start of the line following this include directive.
+void HeaderIncludes::addExistingInclude(Include IncludeToAdd,
+                                        unsigned NextLineOffset) {
+  auto Iter =
+      ExistingIncludes.try_emplace(trimInclude(IncludeToAdd.Name)).first;
+  Iter->second.push_back(std::move(IncludeToAdd));
+  auto &CurInclude = Iter->second.back();
+  // The header name with quotes or angle brackets.
+  // Only record the offset of current #include if we can insert after it.
+  if (CurInclude.R.getOffset() <= MaxInsertOffset) {
+    int Priority = Categories.getIncludePriority(
+        CurInclude.Name, /*CheckMainHeader=*/FirstIncludeOffset < 0);
+    CategoryEndOffsets[Priority] = NextLineOffset;
+    IncludesByPriority[Priority].push_back(&CurInclude);
+    if (FirstIncludeOffset < 0)
+      FirstIncludeOffset = CurInclude.R.getOffset();
+  }
+}
+
+llvm::Optional<Replacement> HeaderIncludes::insert(llvm::StringRef IncludeName) const {
+  std::string TrimmedName = trimInclude(IncludeName);
+  if (ExistingIncludes.find(TrimmedName) != ExistingIncludes.end())
+    return llvm::None;
+  std::string Quoted =
+      (IncludeName.startswith("<") || IncludeName.startswith("\""))
+          ? IncludeName.str()
+          : ("\"" + IncludeName + "\"").str();
+  StringRef QuotedName = Quoted;
+  int Priority = Categories.getIncludePriority(
+      QuotedName, /*CheckMainHeader=*/FirstIncludeOffset < 0);
+  auto CatOffset = CategoryEndOffsets.find(Priority);
+  assert(CatOffset != CategoryEndOffsets.end());
+  unsigned InsertOffset = CatOffset->second; // Fall back offset
+  auto Iter = IncludesByPriority.find(Priority);
+  if (Iter != IncludesByPriority.end()) {
+    for (const auto *Inc : Iter->second) {
+      if (QuotedName < Inc->Name) {
+        InsertOffset = Inc->R.getOffset();
+        break;
+      }
+    }
+  }
+  assert(InsertOffset <= Code.size());
+  std::string NewInclude = ("#include " + QuotedName + "\n").str();
+  // FIXME: when inserting headers at end of the code, we would also need to
+  // append '\n' to the code if it doesn't already end with '\n'.
+  return Replacement(FileName, InsertOffset, 0, NewInclude);
+}
+
+Replacements HeaderIncludes::remove(llvm::StringRef IncludeName) const {
+  StringRef TrimmedName= trimInclude(IncludeName);
+  Replacements Result;
+  auto Iter = ExistingIncludes.find(TrimmedName);
+  if (Iter == ExistingIncludes.end())
+    return Result;
+  for (const auto &Inc : Iter->second) {
+    assert(trimInclude(Inc.Name) == TrimmedName);
+    // If IncludeName is not quoted, then both quotations are deleted.
+    if (IncludeName != TrimmedName &&
+        IncludeName != Inc.Name) // Different quotations, skip.
+      continue;
+    llvm::Error Err = Result.add(tooling::Replacement(
+        FileName, Inc.R.getOffset(), Inc.R.getLength(), ""));
+    if (Err) {
+      auto ErrMsg = "Unexpected conflicts in #include deletions: " +
+                    llvm::toString(std::move(Err));
+      llvm_unreachable(ErrMsg.c_str());
+    }
+  }
+  return Result;
+}
+
+} // namespace tooling
+} // namespace clang

unittests/Format/CleanupTest.cpp

std::string Expected = "\n\nnamespace C {\n"
                          "}";
   tooling::Replacements Replaces =
       toReplacements({createReplacement(getOffset(Code, 3, 3), 6, ""),
                       createReplacement(getOffset(Code, 9, 34), 6, "")});
 
   EXPECT_EQ(Expected, formatAndApply(Code, Replaces));
 }
 
-TEST_F(CleanUpReplacementsTest, NoExistingIncludeWithoutDefine) {
-  std::string Code = "int main() {}";
-  std::string Expected = "#include \"a.h\"\n"
-                         "int main() {}";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include \"a.h\"")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, NoExistingIncludeWithDefine) {
-  std::string Code = "#ifndef A_H\n"
-                     "#define A_H\n"
-                     "class A {};\n"
-                     "#define MMM 123\n"
-                     "#endif";
-  std::string Expected = "#ifndef A_H\n"
-                         "#define A_H\n"
-                         "#include \"b.h\"\n"
-                         "class A {};\n"
-                         "#define MMM 123\n"
-                         "#endif";
-
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include \"b.h\"")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, InsertBeforeCategoryWithLowerPriority) {
-  std::string Code = "#ifndef A_H\n"
-                     "#define A_H\n"
-                     "\n"
-                     "\n"
-                     "\n"
-                     "#include <vector>\n"
-                     "class A {};\n"
-                     "#define MMM 123\n"
-                     "#endif";
-  std::string Expected = "#ifndef A_H\n"
-                         "#define A_H\n"
-                         "\n"
-                         "\n"
-                         "\n"
-                         "#include \"a.h\"\n"
-                         "#include <vector>\n"
-                         "class A {};\n"
-                         "#define MMM 123\n"
-                         "#endif";
-
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include \"a.h\"")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, InsertAfterMainHeader) {
-  std::string Code = "#include \"fix.h\"\n"
-                     "\n"
-                     "int main() {}";
-  std::string Expected = "#include \"fix.h\"\n"
-                         "#include <a>\n"
-                         "\n"
-                         "int main() {}";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <a>")});
-  Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, InsertBeforeSystemHeaderLLVM) {
-  std::string Code = "#include <memory>\n"
-                     "\n"
-                     "int main() {}";
-  std::string Expected = "#include \"z.h\"\n"
-                         "#include <memory>\n"
-                         "\n"
-                         "int main() {}";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include \"z.h\"")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, InsertAfterSystemHeaderGoogle) {
-  std::string Code = "#include <memory>\n"
-                     "\n"
-                     "int main() {}";
-  std::string Expected = "#include <memory>\n"
-                         "#include \"z.h\"\n"
-                         "\n"
-                         "int main() {}";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include \"z.h\"")});
-  Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, InsertOneIncludeLLVMStyle) {
-  std::string Code = "#include \"x/fix.h\"\n"
-                     "#include \"a.h\"\n"
-                     "#include \"b.h\"\n"
-                     "#include \"clang/Format/Format.h\"\n"
-                     "#include <memory>\n";
-  std::string Expected = "#include \"x/fix.h\"\n"
-                         "#include \"a.h\"\n"
-                         "#include \"b.h\"\n"
-                         "#include \"d.h\"\n"
-                         "#include \"clang/Format/Format.h\"\n"
-                         "#include \"llvm/x/y.h\"\n"
-                         "#include <memory>\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include \"d.h\""),
-                      createInsertion("#include \"llvm/x/y.h\"")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
 TEST_F(CleanUpReplacementsTest, InsertMultipleIncludesLLVMStyle) {
   std::string Code = "#include \"x/fix.h\"\n"
                      "#include \"a.h\"\n"
                      "#include \"b.h\"\n"
                      "#include \"clang/Format/Format.h\"\n"
                      "#include <memory>\n";
   std::string Expected = "#include \"x/fix.h\"\n"
                          "#include \"a.h\"\n"
                          "#include \"b.h\"\n"
                          "#include \"new/new.h\"\n"
                          "#include \"clang/Format/Format.h\"\n"
-                         "#include <memory>\n"
-                         "#include <list>\n";
+                         "#include <list>\n"
+                         "#include <memory>\n";
   tooling::Replacements Replaces =
       toReplacements({createInsertion("#include <list>"),
                       createInsertion("#include \"new/new.h\"")});
   EXPECT_EQ(Expected, apply(Code, Replaces));
 }
 
-TEST_F(CleanUpReplacementsTest, InsertNewSystemIncludeGoogleStyle) {
-  std::string Code = "#include \"x/fix.h\"\n"
-                     "\n"
-                     "#include \"y/a.h\"\n"
-                     "#include \"z/b.h\"\n";
-  // FIXME: inserting after the empty line following the main header might be
-  // preferred.
-  std::string Expected = "#include \"x/fix.h\"\n"
-                         "#include <vector>\n"
-                         "\n"
-                         "#include \"y/a.h\"\n"
-                         "#include \"z/b.h\"\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
 TEST_F(CleanUpReplacementsTest, InsertMultipleIncludesGoogleStyle) {
   std::string Code = "#include \"x/fix.h\"\n"
                      "\n"
+                     "#include <abc>\n"
                      "#include <vector>\n"
                      "\n"
                      "#include \"y/a.h\"\n"
                      "#include \"z/b.h\"\n";
   std::string Expected = "#include \"x/fix.h\"\n"
                          "\n"
-                         "#include <vector>\n"
+                         "#include <abc>\n"
                          "#include <list>\n"
+                         "#include <vector>\n"
                          "\n"
+                         "#include \"x/x.h\"\n"
                          "#include \"y/a.h\"\n"
-                         "#include \"z/b.h\"\n"
-                         "#include \"x/x.h\"\n";
+                         "#include \"z/b.h\"\n";
   tooling::Replacements Replaces =
       toReplacements({createInsertion("#include <list>"),
                       createInsertion("#include \"x/x.h\"")});
   Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
   EXPECT_EQ(Expected, apply(Code, Replaces));
 }
 
 TEST_F(CleanUpReplacementsTest, InsertMultipleNewHeadersAndSortLLVM) {
tooling::Replacements Replaces = toReplacements(
       {createReplacement(getOffset(Code, 4, 8), 1, "b"),
        createInsertion("#include <vector>"), createInsertion("#include <list>"),
        createInsertion("#include \"clang/x/x.h\""),
        createInsertion("#include \"y.h\""),
        createInsertion("#include \"x.h\"")});
   EXPECT_EQ(Expected, formatAndApply(Code, Replaces));
 }
 
-TEST_F(CleanUpReplacementsTest, NotConfusedByDefine) {
-  std::string Code = "void f() {}\n"
-                     "#define A \\\n"
-                     "  int i;";
-  std::string Expected = "#include <vector>\n"
-                         "void f() {}\n"
-                         "#define A \\\n"
-                         "  int i;";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, formatAndApply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, SkippedTopComment) {
-  std::string Code = "// comment\n"
-                     "\n"
-                     "   // comment\n";
-  std::string Expected = "// comment\n"
-                         "\n"
-                         "   // comment\n"
-                         "#include <vector>\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, SkippedMixedComments) {
-  std::string Code = "// comment\n"
-                     "// comment \\\n"
-                     " comment continued\n"
-                     "/*\n"
-                     "* comment\n"
-                     "*/\n";
-  std::string Expected = "// comment\n"
-                         "// comment \\\n"
-                         " comment continued\n"
-                         "/*\n"
-                         "* comment\n"
-                         "*/\n"
-                         "#include <vector>\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, MultipleBlockCommentsInOneLine) {
-  std::string Code = "/*\n"
-                     "* comment\n"
-                     "*/ /* comment\n"
-                     "*/\n"
-                     "\n\n"
-                     "/* c1 */ /*c2 */\n";
-  std::string Expected = "/*\n"
-                         "* comment\n"
-                         "*/ /* comment\n"
-                         "*/\n"
-                         "\n\n"
-                         "/* c1 */ /*c2 */\n"
-                         "#include <vector>\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, CodeAfterComments) {
-  std::string Code = "/*\n"
-                     "* comment\n"
-                     "*/ /* comment\n"
-                     "*/\n"
-                     "\n\n"
-                     "/* c1 */ /*c2 */\n"
-                     "\n"
-                     "int x;\n";
-  std::string Expected = "/*\n"
-                         "* comment\n"
-                         "*/ /* comment\n"
-                         "*/\n"
-                         "\n\n"
-                         "/* c1 */ /*c2 */\n"
-                         "\n"
-                         "#include <vector>\n"
-                         "int x;\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, FakeHeaderGuardIfDef) {
-  std::string Code = "// comment \n"
-                     "#ifdef X\n"
-                     "#define X\n";
-  std::string Expected = "// comment \n"
-                         "#include <vector>\n"
-                         "#ifdef X\n"
-                         "#define X\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, RealHeaderGuardAfterComments) {
-  std::string Code = "// comment \n"
-                     "#ifndef X\n"
-                     "#define X\n"
-                     "int x;\n"
-                     "#define Y 1\n";
-  std::string Expected = "// comment \n"
-                         "#ifndef X\n"
-                         "#define X\n"
-                         "#include <vector>\n"
-                         "int x;\n"
-                         "#define Y 1\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, IfNDefWithNoDefine) {
-  std::string Code = "// comment \n"
-                     "#ifndef X\n"
-                     "int x;\n"
-                     "#define Y 1\n";
-  std::string Expected = "// comment \n"
-                         "#include <vector>\n"
-                         "#ifndef X\n"
-                         "int x;\n"
-                         "#define Y 1\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, FakeHeaderGuard) {
-  std::string Code = "// comment \n"
-                     "#ifndef X\n"
-                     "#define 1\n";
-  std::string Expected = "// comment \n"
-                         "#include <vector>\n"
-                         "#ifndef X\n"
-                         "#define 1\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, HeaderGuardWithComment) {
-  std::string Code = "// comment \n"
-                     "#ifndef X // comment\n"
-                     "// comment\n"
-                     "/* comment\n"
-                     "*/\n"
-                     "/* comment */ #define X\n"
-                     "int x;\n"
-                     "#define Y 1\n";
-  std::string Expected = "// comment \n"
-                         "#ifndef X // comment\n"
-                         "// comment\n"
-                         "/* comment\n"
-                         "*/\n"
-                         "/* comment */ #define X\n"
-                         "#include <vector>\n"
-                         "int x;\n"
-                         "#define Y 1\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, EmptyCode) {
-  std::string Code = "";
-  std::string Expected = "#include <vector>\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, NoNewLineAtTheEndOfCode) {
-  std::string Code = "#include <map>";
-  std::string Expected = "#include <map>\n#include <vector>\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, NoNewLineAtTheEndOfCodeMultipleInsertions) {
-  std::string Code = "#include <map>";
-  std::string Expected =
-      "#include <map>\n#include <string>\n#include <vector>\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <string>"),
-                      createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, SkipExistingHeaders) {
-  std::string Code = "#include \"a.h\"\n"
-                     "#include <vector>\n";
-  std::string Expected = "#include \"a.h\"\n"
-                         "#include <vector>\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>"),
-                      createInsertion("#include \"a.h\"")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, AddIncludesWithDifferentForms) {
-  std::string Code = "#include \"a.h\"\n"
-                     "#include <vector>\n";
-  // FIXME: this might not be the best behavior.
-  std::string Expected = "#include \"a.h\"\n"
-                         "#include \"vector\"\n"
-                         "#include <vector>\n"
-                         "#include <a.h>\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include \"vector\""),
-                      createInsertion("#include <a.h>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
 TEST_F(CleanUpReplacementsTest, SimpleDeleteIncludes) {
   std::string Code = "#include \"abc.h\"\n"
                      "#include \"xyz.h\" // comment\n"
                      "#include \"xyz\"\n"
                      "int x;\n";
   std::string Expected = "#include \"xyz\"\n"
                          "int x;\n";
   tooling::Replacements Replaces =
       toReplacements({createDeletion("abc.h"), createDeletion("xyz.h")});
   EXPECT_EQ(Expected, apply(Code, Replaces));
 }
 
-TEST_F(CleanUpReplacementsTest, DeleteAllCode) {
-  std::string Code = "#include \"xyz.h\"\n"
-                     "#include <xyz.h>";
-  std::string Expected = "";
-  tooling::Replacements Replaces = toReplacements({createDeletion("xyz.h")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, DeleteAllIncludesWithSameNameIfNoType) {
-  std::string Code = "#include \"xyz.h\"\n"
-                     "#include \"xyz\"\n"
-                     "#include <xyz.h>\n";
-  std::string Expected = "#include \"xyz\"\n";
-  tooling::Replacements Replaces = toReplacements({createDeletion("xyz.h")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, OnlyDeleteHeaderWithType) {
-  std::string Code = "#include \"xyz.h\"\n"
-                     "#include \"xyz\"\n"
-                     "#include <xyz.h>";
-  std::string Expected = "#include \"xyz.h\"\n"
-                         "#include \"xyz\"\n";
-  tooling::Replacements Replaces = toReplacements({createDeletion("<xyz.h>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
 TEST_F(CleanUpReplacementsTest, InsertionAndDeleteHeader) {
   std::string Code = "#include \"a.h\"\n"
                      "\n"
                      "#include <vector>\n";
   std::string Expected = "#include \"a.h\"\n"
                          "\n"
                          "#include <map>\n";
   tooling::Replacements Replaces = toReplacements(
       {createDeletion("<vector>"), createInsertion("#include <map>")});
   EXPECT_EQ(Expected, apply(Code, Replaces));
 }
 
-TEST_F(CleanUpReplacementsTest, NoInsertionAfterCode) {
-  std::string Code = "#include \"a.h\"\n"
-                     "void f() {}\n"
-                     "#include \"b.h\"\n";
-  std::string Expected = "#include \"a.h\"\n"
-                         "#include \"c.h\"\n"
-                         "void f() {}\n"
-                         "#include \"b.h\"\n";
-  tooling::Replacements Replaces = toReplacements(
-      {createInsertion("#include \"c.h\"")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, NoInsertionInStringLiteral) {
-  std::string Code = "#include \"a.h\"\n"
-                     "const char[] = R\"(\n"
-                     "#include \"b.h\"\n"
-                     ")\";\n";
-  std::string Expected = "#include \"a.h\"\n"
-                         "#include \"c.h\"\n"
-                         "const char[] = R\"(\n"
-                         "#include \"b.h\"\n"
-                         ")\";\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include \"c.h\"")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, NoInsertionAfterOtherDirective) {
-  std::string Code = "#include \"a.h\"\n"
-                     "#ifdef X\n"
-                     "#include \"b.h\"\n"
-                     "#endif\n";
-  std::string Expected = "#include \"a.h\"\n"
-                         "#include \"c.h\"\n"
-                         "#ifdef X\n"
-                         "#include \"b.h\"\n"
-                         "#endif\n";
-  tooling::Replacements Replaces = toReplacements(
-      {createInsertion("#include \"c.h\"")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, CanInsertAfterLongSystemInclude) {
-  std::string Code = "#include \"a.h\"\n"
-                     "// comment\n\n"
-                     "#include <a/b/c/d/e.h>\n";
-  std::string Expected = "#include \"a.h\"\n"
-                         "// comment\n\n"
-                         "#include <a/b/c/d/e.h>\n"
-                         "#include <x.h>\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <x.h>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, CanInsertAfterComment) {
-  std::string Code = "#include \"a.h\"\n"
-                     "// Comment\n"
-                     "\n"
-                     "/* Comment */\n"
-                     "// Comment\n"
-                     "\n"
-                     "#include \"b.h\"\n";
-  std::string Expected = "#include \"a.h\"\n"
-                         "// Comment\n"
-                         "\n"
-                         "/* Comment */\n"
-                         "// Comment\n"
-                         "\n"
-                         "#include \"b.h\"\n"
-                         "#include \"c.h\"\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include \"c.h\"")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, LongCommentsInTheBeginningOfFile) {
-  std::string Code = "// Loooooooooooooooooooooooooong comment\n"
-                     "// Loooooooooooooooooooooooooong comment\n"
-                     "// Loooooooooooooooooooooooooong comment\n"
-                     "#include <string>\n"
-                     "#include <vector>\n"
-                     "\n"
-                     "#include \"a.h\"\n"
-                     "#include \"b.h\"\n";
-  std::string Expected = "// Loooooooooooooooooooooooooong comment\n"
-                         "// Loooooooooooooooooooooooooong comment\n"
-                         "// Loooooooooooooooooooooooooong comment\n"
-                         "#include <string>\n"
-                         "#include <vector>\n"
-                         "\n"
-                         "#include \"a.h\"\n"
-                         "#include \"b.h\"\n"
-                         "#include \"third.h\"\n";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include \"third.h\"")});
-  Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, CanDeleteAfterCode) {
-  std::string Code = "#include \"a.h\"\n"
-                     "void f() {}\n"
-                     "#include \"b.h\"\n";
-  std::string Expected = "#include \"a.h\"\n"
-                         "void f() {}\n";
-  tooling::Replacements Replaces = toReplacements({createDeletion("\"b.h\"")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
 } // end namespace
 } // end namespace format
 } // end namespace clang

unittests/Format/FormatTest.cpp

TEST_F(FormatTest, ParsesConfiguration) {
   BoostForeach.push_back("BOOST_FOREACH");
   CHECK_PARSE("ForEachMacros: [BOOST_FOREACH]", ForEachMacros, BoostForeach);
   std::vector<std::string> BoostAndQForeach;
   BoostAndQForeach.push_back("BOOST_FOREACH");
   BoostAndQForeach.push_back("Q_FOREACH");
   CHECK_PARSE("ForEachMacros: [BOOST_FOREACH, Q_FOREACH]", ForEachMacros,
               BoostAndQForeach);
 
-  Style.IncludeCategories.clear();
-  std::vector<FormatStyle::IncludeCategory> ExpectedCategories = {{"abc/.*", 2},
-                                                                  {".*", 1}};
+  Style.IncludeStyle.IncludeCategories.clear();
+  std::vector<tooling::IncludeStyle::IncludeCategory> ExpectedCategories = {
+      {"abc/.*", 2}, {".*", 1}};
   CHECK_PARSE("IncludeCategories:\n"
               "  - Regex: abc/.*\n"
               "    Priority: 2\n"
               "  - Regex: .*\n"
               "    Priority: 1",
-              IncludeCategories, ExpectedCategories);
-  CHECK_PARSE("IncludeIsMainRegex: 'abc$'", IncludeIsMainRegex, "abc$");
+              IncludeStyle.IncludeCategories, ExpectedCategories);
+  CHECK_PARSE("IncludeIsMainRegex: 'abc$'", IncludeStyle.IncludeIsMainRegex,
+              "abc$");
 
   Style.RawStringFormats.clear();
   std::vector<FormatStyle::RawStringFormat> ExpectedRawStringFormats = {
       {
           FormatStyle::LK_TextProto,
           {"pb", "proto"},
           {"PARSE_TEXT_PROTO"},
           /*CanonicalDelimiter=*/"",

unittests/Format/SortIncludesTest.cpp

 class SortIncludesTest : public ::testing::Test {
 protected:
   std::vector<tooling::Range> GetCodeRange(StringRef Code) {
     return std::vector<tooling::Range>(1, tooling::Range(0, Code.size()));
   }
 
   std::string sort(StringRef Code, std::vector<tooling::Range> Ranges,
                    StringRef FileName = "input.cc") {
-    auto Replaces = sortIncludes(Style, Code, Ranges, FileName);
+    auto Replaces = sortIncludes(FmtStyle, Code, Ranges, FileName);
     Ranges = tooling::calculateRangesAfterReplacements(Replaces, Ranges);
     auto Sorted = applyAllReplacements(Code, Replaces);
     EXPECT_TRUE(static_cast<bool>(Sorted));
     auto Result = applyAllReplacements(
-        *Sorted, reformat(Style, *Sorted, Ranges, FileName));
+        *Sorted, reformat(FmtStyle, *Sorted, Ranges, FileName));
     EXPECT_TRUE(static_cast<bool>(Result));
     return *Result;
   }
 
   std::string sort(StringRef Code, StringRef FileName = "input.cpp") {
     return sort(Code, GetCodeRange(Code), FileName);
   }
 
   unsigned newCursor(llvm::StringRef Code, unsigned Cursor) {
-    sortIncludes(Style, Code, GetCodeRange(Code), "input.cpp", &Cursor);
+    sortIncludes(FmtStyle, Code, GetCodeRange(Code), "input.cpp", &Cursor);
     return Cursor;
   }
 
-  FormatStyle Style = getLLVMStyle();
-
+  FormatStyle FmtStyle = getLLVMStyle();
+  tooling::IncludeStyle &Style = FmtStyle.IncludeStyle;
 };
 
 TEST_F(SortIncludesTest, BasicSorting) {
   EXPECT_EQ("#include \"a.h\"\n"
             "#include \"b.h\"\n"
             "#include \"c.h\"\n",
             sort("#include \"a.h\"\n"
                  "#include \"c.h\"\n"

 TEST_F(SortIncludesTest, NoReplacementsForValidIncludes) {
   // Identical #includes have led to a failure with an unstable sort.
   std::string Code = "#include <a>\n"
                      "#include <b>\n"
                      "#include <c>\n"
                      "#include <d>\n"
                      "#include <e>\n"
                      "#include <f>\n";
-  EXPECT_TRUE(sortIncludes(Style, Code, GetCodeRange(Code), "a.cc").empty());
+  EXPECT_TRUE(sortIncludes(FmtStyle, Code, GetCodeRange(Code), "a.cc").empty());
 }
 
 TEST_F(SortIncludesTest, SortedIncludesInMultipleBlocksAreMerged) {
-  Style.IncludeBlocks = FormatStyle::IBS_Merge;
+  Style.IncludeBlocks = tooling::IncludeStyle::IBS_Merge;
   EXPECT_EQ("#include \"a.h\"\n"
             "#include \"b.h\"\n"
             "#include \"c.h\"\n",
             sort("#include \"a.h\"\n"
                  "#include \"c.h\"\n"
                  "\n"
                  "\n"
                  "#include \"b.h\"\n"));
 
-  Style.IncludeBlocks = FormatStyle::IBS_Regroup;
+  Style.IncludeBlocks = tooling::IncludeStyle::IBS_Regroup;
   EXPECT_EQ("#include \"a.h\"\n"
             "#include \"b.h\"\n"
             "#include \"c.h\"\n",
             sort("#include \"a.h\"\n"
                  "#include \"c.h\"\n"
                  "\n"
                  "\n"
                  "#include \"b.h\"\n"));
EXPECT_EQ("#include <a>\n"
                  "// clang-format off\n"
                  "#include <b>\n"
                  "#include <a>\n"
                  "#include <c>\n"
                  "// clang-format on\n"));
 }
 
 TEST_F(SortIncludesTest, IncludeSortingCanBeDisabled) {
-  Style.SortIncludes = false;
+  FmtStyle.SortIncludes = false;
   EXPECT_EQ("#include \"a.h\"\n"
             "#include \"c.h\"\n"
             "#include \"b.h\"\n",
             sort("#include \"a.h\"\n"
                  "#include \"c.h\"\n"
                  "#include \"b.h\"\n"));
 }
 
EXPECT_EQ("#include \"a.h\"\n"
             "#include \"b.h\"\n",
             sort("#include \"a.h\"\n"
                  "#include \"c.h\"\n"
                  "\n"
                  "#include \"b.h\"\n"));
 }
 
 TEST_F(SortIncludesTest, SortsAllBlocksWhenMerging) {
-  Style.IncludeBlocks = FormatStyle::IBS_Merge;
+  Style.IncludeBlocks = tooling::IncludeStyle::IBS_Merge;
   EXPECT_EQ("#include \"a.h\"\n"
             "#include \"b.h\"\n"
             "#include \"c.h\"\n",
             sort("#include \"a.h\"\n"
                  "#include \"c.h\"\n"
                  "\n"
                  "#include \"b.h\"\n"));
 }
 
 TEST_F(SortIncludesTest, CommentsAlwaysSeparateGroups) {
   EXPECT_EQ("#include \"a.h\"\n"
             "#include \"c.h\"\n"
             "// comment\n"
             "#include \"b.h\"\n",
             sort("#include \"c.h\"\n"
                  "#include \"a.h\"\n"
                  "// comment\n"
                  "#include \"b.h\"\n"));
 
-  Style.IncludeBlocks = FormatStyle::IBS_Merge;
+  Style.IncludeBlocks = tooling::IncludeStyle::IBS_Merge;
   EXPECT_EQ("#include \"a.h\"\n"
             "#include \"c.h\"\n"
             "// comment\n"
             "#include \"b.h\"\n",
             sort("#include \"c.h\"\n"
                  "#include \"a.h\"\n"
                  "// comment\n"
                  "#include \"b.h\"\n"));
 
-  Style.IncludeBlocks = FormatStyle::IBS_Regroup;
+  Style.IncludeBlocks = tooling::IncludeStyle::IBS_Regroup;
   EXPECT_EQ("#include \"a.h\"\n"
             "#include \"c.h\"\n"
             "// comment\n"
             "#include \"b.h\"\n",
             sort("#include \"c.h\"\n"
                  "#include \"a.h\"\n"
                  "// comment\n"
                  "#include \"b.h\"\n"));
 }
 
 TEST_F(SortIncludesTest, HandlesAngledIncludesAsSeparateBlocks) {
   EXPECT_EQ("#include \"a.h\"\n"
             "#include \"c.h\"\n"
             "#include <b.h>\n"
             "#include <d.h>\n",
             sort("#include <d.h>\n"
                  "#include <b.h>\n"
                  "#include \"c.h\"\n"
                  "#include \"a.h\"\n"));
 
-  Style = getGoogleStyle(FormatStyle::LK_Cpp);
+  FmtStyle = getGoogleStyle(FormatStyle::LK_Cpp);
   EXPECT_EQ("#include <b.h>\n"
             "#include <d.h>\n"
             "#include \"a.h\"\n"
             "#include \"c.h\"\n",
             sort("#include <d.h>\n"
                  "#include <b.h>\n"
                  "#include \"c.h\"\n"
                  "#include \"a.h\"\n"));
 }
 
 TEST_F(SortIncludesTest, RegroupsAngledIncludesInSeparateBlocks) {
-  Style.IncludeBlocks = FormatStyle::IBS_Regroup;
+  Style.IncludeBlocks = tooling::IncludeStyle::IBS_Regroup;
   EXPECT_EQ("#include \"a.h\"\n"
             "#include \"c.h\"\n"
             "\n"
             "#include <b.h>\n"
             "#include <d.h>\n",
             sort("#include <d.h>\n"
                  "#include <b.h>\n"
                  "#include \"c.h\"\n"
EXPECT_EQ("#include \"a.h\"\n"
                  "#include \"a.h\"\n"
                  "#include \"c.h\"\n"
                  "#include \"llvm/a.h\"\n",
                  "a.cc"));
 }
 
 TEST_F(SortIncludesTest, RecognizeMainHeaderInAllGroups) {
   Style.IncludeIsMainRegex = "([-_](test|unittest))?$";
-  Style.IncludeBlocks = FormatStyle::IBS_Merge;
+  Style.IncludeBlocks = tooling::IncludeStyle::IBS_Merge;
 
   EXPECT_EQ("#include \"c.h\"\n"
             "#include \"a.h\"\n"
             "#include \"b.h\"\n",
             sort("#include \"b.h\"\n"
                  "\n"
                  "#include \"a.h\"\n"
                  "#include \"c.h\"\n",
                  "c.cc"));
 }
 
 TEST_F(SortIncludesTest, MainHeaderIsSeparatedWhenRegroupping) {
   Style.IncludeIsMainRegex = "([-_](test|unittest))?$";
-  Style.IncludeBlocks = FormatStyle::IBS_Regroup;
+  Style.IncludeBlocks = tooling::IncludeStyle::IBS_Regroup;
 
   EXPECT_EQ("#include \"a.h\"\n"
             "\n"
             "#include \"b.h\"\n"
             "#include \"c.h\"\n",
             sort("#include \"b.h\"\n"
                  "\n"
                  "#include \"a.h\"\n"
EXPECT_EQ("#include \"important_os_header.h\"\n"
             sort("#include \"important_os_header.h\"\n"
                  "#include \"c_main.h\"\n"
                  "#include \"a_other.h\"\n",
                  "c_main.cc"));
 }
 
 TEST_F(SortIncludesTest, PriorityGroupsAreSeparatedWhenRegroupping) {
   Style.IncludeCategories = {{".*important_os_header.*", -1}, {".*", 1}};
-  Style.IncludeBlocks = FormatStyle::IBS_Regroup;
+  Style.IncludeBlocks = tooling::IncludeStyle::IBS_Regroup;
 
   EXPECT_EQ("#include \"important_os_header.h\"\n"
             "\n"
             "#include \"c_main.h\"\n"
             "\n"
             "#include \"a_other.h\"\n",
             sort("#include \"c_main.h\"\n"
                  "#include \"a_other.h\"\n"
EXPECT_EQ("#include <a>\n"
             "#include <c>\n",
             sort("#include <a>\n"
                  "#include <b>\n"
                  "#include <b>\n"
                  "#include <b>\n"
                  "#include <b>\n"
                  "#include <c>\n"));
 
-  Style.IncludeBlocks = FormatStyle::IBS_Merge;
+  Style.IncludeBlocks = tooling::IncludeStyle::IBS_Merge;
   EXPECT_EQ("#include <a>\n"
             "#include <b>\n"
             "#include <c>\n",
             sort("#include <a>\n"
                  "#include <b>\n"
                  "\n"
                  "#include <b>\n"
                  "\n"
                  "#include <b>\n"
                  "#include <c>\n"));
 
-  Style.IncludeBlocks = FormatStyle::IBS_Regroup;
+  Style.IncludeBlocks = tooling::IncludeStyle::IBS_Regroup;
   EXPECT_EQ("#include <a>\n"
             "#include <b>\n"
             "#include <c>\n",
             sort("#include <a>\n"
                  "#include <b>\n"
                  "\n"
                  "#include <b>\n"
                  "\n"
                  "#include <b>\n"
                  "#include <c>\n"));
 }
 
 TEST_F(SortIncludesTest, SortAndDeduplicateIncludes) {
   EXPECT_EQ("#include <a>\n"
             "#include <b>\n"
             "#include <c>\n",
             sort("#include <b>\n"
                  "#include <a>\n"
                  "#include <b>\n"
                  "#include <b>\n"
                  "#include <c>\n"
                  "#include <b>\n"));
 
-  Style.IncludeBlocks = FormatStyle::IBS_Merge;
+  Style.IncludeBlocks = tooling::IncludeStyle::IBS_Merge;
   EXPECT_EQ("#include <a>\n"
             "#include <b>\n"
             "#include <c>\n",
             sort("#include <b>\n"
                  "#include <a>\n"
                  "\n"
                  "#include <b>\n"
                  "\n"
                  "#include <c>\n"
                  "#include <b>\n"));
 
-  Style.IncludeBlocks = FormatStyle::IBS_Regroup;
+  Style.IncludeBlocks = tooling::IncludeStyle::IBS_Regroup;
   EXPECT_EQ("#include <a>\n"
             "#include <b>\n"
             "#include <c>\n",
             sort("#include <b>\n"
                  "#include <a>\n"
                  "\n"
                  "#include <b>\n"
                  "\n"

 TEST_F(SortIncludesTest, ValidAffactedRangesAfterDeduplicatingIncludes) {
   std::string Code = "#include <a>\n"
                      "#include <b>\n"
                      "#include <a>\n"
                      "#include <a>\n"
                      "\n"
                      "   int     x ;";
   std::vector<tooling::Range> Ranges = {tooling::Range(0, 52)};
-  auto Replaces = sortIncludes(Style, Code, Ranges, "input.cpp");
+  auto Replaces = sortIncludes(FmtStyle, Code, Ranges, "input.cpp");
   Ranges = tooling::calculateRangesAfterReplacements(Replaces, Ranges);
   EXPECT_EQ(1u, Ranges.size());
   EXPECT_EQ(0u, Ranges[0].getOffset());
   EXPECT_EQ(26u, Ranges[0].getLength());
 }
 
 TEST_F(SortIncludesTest, DoNotSortLikelyXml) {
   EXPECT_EQ("<!--;\n"

unittests/Tooling/CMakeLists.txt

 add_clang_unittest(ToolingTests
   ASTSelectionTest.cpp
   CastExprTest.cpp
   CommentHandlerTest.cpp
   CompilationDatabaseTest.cpp
   DiagnosticsYamlTest.cpp
   ExecutionTest.cpp
   FixItTest.cpp
+  HeaderIncludesTest.cpp
   LexicallyOrderedRecursiveASTVisitorTest.cpp
   LookupTest.cpp
   QualTypeNamesTest.cpp
   RecursiveASTVisitorTests/Attr.cpp
   RecursiveASTVisitorTests/Class.cpp
   RecursiveASTVisitorTests/ConstructExpr.cpp
   RecursiveASTVisitorTests/CXXBoolLiteralExpr.cpp
   RecursiveASTVisitorTests/CXXMemberCall.cpp

unittests/Tooling/HeaderIncludesTest.cpp

+//===- unittest/Tooling/CleanupTest.cpp - Include insertion/deletion tests ===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/Tooling/Core/HeaderIncludes.h"
+#include "../Tooling/ReplacementTest.h"
+#include "../Tooling/RewriterTestContext.h"
+#include "clang/Format/Format.h"
+#include "clang/Tooling/Core/Replacement.h"
+
+#include "gtest/gtest.h"
+
+using clang::tooling::ReplacementTest;
+using clang::tooling::toReplacements;
+
+namespace clang {
+namespace tooling {
+namespace {
+
+class HeaderIncludesTest : public ::testing::Test {
+protected:
+  std::string insert(llvm::StringRef Code, llvm::StringRef Header) {
+    HeaderIncludes Includes(FileName, Code, Style);
+    auto R = Includes.insert(Header);
+    if (!R)
+      return Code;
+    auto Result = applyAllReplacements(Code, Replacements(*R));
+    EXPECT_TRUE(static_cast<bool>(Result));
+    return *Result;
+  }
+
+  std::string remove(llvm::StringRef Code, llvm::StringRef Header) {
+    HeaderIncludes Includes(FileName, Code, Style);
+    auto Replaces = Includes.remove(Header);
+    auto Result = applyAllReplacements(Code, Replaces);
+    EXPECT_TRUE(static_cast<bool>(Result));
+    return *Result;
+  }
+
+  const std::string FileName = "fix.cpp";
+  IncludeStyle Style = format::getLLVMStyle().IncludeStyle;
+};
+
+TEST_F(HeaderIncludesTest, NoExistingIncludeWithoutDefine) {
+  std::string Code = "int main() {}";
+  std::string Expected = "#include \"a.h\"\n"
+                         "int main() {}";
+  EXPECT_EQ(Expected, insert(Code, "a.h"));
+}
+
+TEST_F(HeaderIncludesTest, NoExistingIncludeWithDefine) {
+  std::string Code = "#ifndef A_H\n"
+                     "#define A_H\n"
+                     "class A {};\n"
+                     "#define MMM 123\n"
+                     "#endif";
+  std::string Expected = "#ifndef A_H\n"
+                         "#define A_H\n"
+                         "#include \"b.h\"\n"
+                         "class A {};\n"
+                         "#define MMM 123\n"
+                         "#endif";
+
+  EXPECT_EQ(Expected, insert(Code, "\"b.h\""));
+}
+
+TEST_F(HeaderIncludesTest, InsertBeforeCategoryWithLowerPriority) {
+  std::string Code = "#ifndef A_H\n"
+                     "#define A_H\n"
+                     "\n"
+                     "\n"
+                     "\n"
+                     "#include <vector>\n"
+                     "class A {};\n"
+                     "#define MMM 123\n"
+                     "#endif";
+  std::string Expected = "#ifndef A_H\n"
+                         "#define A_H\n"
+                         "\n"
+                         "\n"
+                         "\n"
+                         "#include \"a.h\"\n"
+                         "#include <vector>\n"
+                         "class A {};\n"
+                         "#define MMM 123\n"
+                         "#endif";
+
+  EXPECT_EQ(Expected, insert(Code, "a.h"));
+}
+
+TEST_F(HeaderIncludesTest, InsertAfterMainHeader) {
+  std::string Code = "#include \"fix.h\"\n"
+                     "\n"
+                     "int main() {}";
+  std::string Expected = "#include \"fix.h\"\n"
+                         "#include <a>\n"
+                         "\n"
+                         "int main() {}";
+  Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp)
+              .IncludeStyle;
+  EXPECT_EQ(Expected, insert(Code, "<a>"));
+}
+
+TEST_F(HeaderIncludesTest, InsertBeforeSystemHeaderLLVM) {
+  std::string Code = "#include <memory>\n"
+                     "\n"
+                     "int main() {}";
+  std::string Expected = "#include \"z.h\"\n"
+                         "#include <memory>\n"
+                         "\n"
+                         "int main() {}";
+  EXPECT_EQ(Expected, insert(Code, "z.h"));
+}
+
+TEST_F(HeaderIncludesTest, InsertAfterSystemHeaderGoogle) {
+  std::string Code = "#include <memory>\n"
+                     "\n"
+                     "int main() {}";
+  std::string Expected = "#include <memory>\n"
+                         "#include \"z.h\"\n"
+                         "\n"
+                         "int main() {}";
+  Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp)
+              .IncludeStyle;
+  EXPECT_EQ(Expected, insert(Code, "z.h"));
+}
+
+TEST_F(HeaderIncludesTest, InsertOneIncludeLLVMStyle) {
+  std::string Code = "#include \"x/fix.h\"\n"
+                     "#include \"a.h\"\n"
+                     "#include \"b.h\"\n"
+                     "#include \"clang/Format/Format.h\"\n"
+                     "#include <memory>\n";
+  std::string Expected = "#include \"x/fix.h\"\n"
+                         "#include \"a.h\"\n"
+                         "#include \"b.h\"\n"
+                         "#include \"clang/Format/Format.h\"\n"
+                         "#include \"llvm/x/y.h\"\n"
+                         "#include <memory>\n";
+  EXPECT_EQ(Expected, insert(Code, "llvm/x/y.h"));
+}
+
+TEST_F(HeaderIncludesTest, InsertIntoBlockSorted) {
+  std::string Code = "#include \"x/fix.h\"\n"
+                     "#include \"a.h\"\n"
+                     "#include \"c.h\"\n"
+                     "#include <memory>\n";
+  std::string Expected = "#include \"x/fix.h\"\n"
+                         "#include \"a.h\"\n"
+                         "#include \"b.h\"\n"
+                         "#include \"c.h\"\n"
+                         "#include <memory>\n";
+  EXPECT_EQ(Expected, insert(Code, "b.h"));
+}
+
+TEST_F(HeaderIncludesTest, InsertIntoFirstBlockOfSameKind) {
+  std::string Code = "#include \"x/fix.h\"\n"
+                     "#include \"c.h\"\n"
+                     "#include \"e.h\"\n"
+                     "#include \"f.h\"\n"
+                     "#include <memory>\n"
+                     "#include <vector>\n"
+                     "#include \"m.h\"\n"
+                     "#include \"n.h\"\n";
+  std::string Expected = "#include \"x/fix.h\"\n"
+                         "#include \"c.h\"\n"
+                         "#include \"d.h\"\n"
+                         "#include \"e.h\"\n"
+                         "#include \"f.h\"\n"
+                         "#include <memory>\n"
+                         "#include <vector>\n"
+                         "#include \"m.h\"\n"
+                         "#include \"n.h\"\n";
+  EXPECT_EQ(Expected, insert(Code, "d.h"));
+}
+
+TEST_F(HeaderIncludesTest, InsertIntoSystemBlockSorted) {
+  std::string Code = "#include \"x/fix.h\"\n"
+                     "#include \"a.h\"\n"
+                     "#include \"c.h\"\n"
+                     "#include <a>\n"
+                     "#include <z>\n";
+  std::string Expected = "#include \"x/fix.h\"\n"
+                         "#include \"a.h\"\n"
+                         "#include \"c.h\"\n"
+                         "#include <a>\n"
+                         "#include <vector>\n"
+                         "#include <z>\n";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
+TEST_F(HeaderIncludesTest, InsertNewSystemIncludeGoogleStyle) {
+  std::string Code = "#include \"x/fix.h\"\n"
+                     "\n"
+                     "#include \"y/a.h\"\n"
+                     "#include \"z/b.h\"\n";
+  // FIXME: inserting after the empty line following the main header might be
+  // preferred.
+  std::string Expected = "#include \"x/fix.h\"\n"
+                         "#include <vector>\n"
+                         "\n"
+                         "#include \"y/a.h\"\n"
+                         "#include \"z/b.h\"\n";
+  Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp)
+              .IncludeStyle;
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
+TEST_F(HeaderIncludesTest, NotConfusedByDefine) {
+  std::string Code = "void f() {}\n"
+                     "#define A \\\n"
+                     "  int i;";
+  std::string Expected = "#include <vector>\n"
+                         "void f() {}\n"
+                         "#define A \\\n"
+                         "  int i;";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
+TEST_F(HeaderIncludesTest, SkippedTopComment) {
+  std::string Code = "// comment\n"
+                     "\n"
+                     "   // comment\n";
+  std::string Expected = "// comment\n"
+                         "\n"
+                         "   // comment\n"
+                         "#include <vector>\n";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
+TEST_F(HeaderIncludesTest, SkippedMixedComments) {
+  std::string Code = "// comment\n"
+                     "// comment \\\n"
+                     " comment continued\n"
+                     "/*\n"
+                     "* comment\n"
+                     "*/\n";
+  std::string Expected = "// comment\n"
+                         "// comment \\\n"
+                         " comment continued\n"
+                         "/*\n"
+                         "* comment\n"
+                         "*/\n"
+                         "#include <vector>\n";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
+TEST_F(HeaderIncludesTest, MultipleBlockCommentsInOneLine) {
+  std::string Code = "/*\n"
+                     "* comment\n"
+                     "*/ /* comment\n"
+                     "*/\n"
+                     "\n\n"
+                     "/* c1 */ /*c2 */\n";
+  std::string Expected = "/*\n"
+                         "* comment\n"
+                         "*/ /* comment\n"
+                         "*/\n"
+                         "\n\n"
+                         "/* c1 */ /*c2 */\n"
+                         "#include <vector>\n";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
+TEST_F(HeaderIncludesTest, CodeAfterComments) {
+  std::string Code = "/*\n"
+                     "* comment\n"
+                     "*/ /* comment\n"
+                     "*/\n"
+                     "\n\n"
+                     "/* c1 */ /*c2 */\n"
+                     "\n"
+                     "int x;\n";
+  std::string Expected = "/*\n"
+                         "* comment\n"
+                         "*/ /* comment\n"
+                         "*/\n"
+                         "\n\n"
+                         "/* c1 */ /*c2 */\n"
+                         "\n"
+                         "#include <vector>\n"
+                         "int x;\n";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
+TEST_F(HeaderIncludesTest, FakeHeaderGuardIfDef) {
+  std::string Code = "// comment \n"
+                     "#ifdef X\n"
+                     "#define X\n";
+  std::string Expected = "// comment \n"
+                         "#include <vector>\n"
+                         "#ifdef X\n"
+                         "#define X\n";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
+TEST_F(HeaderIncludesTest, RealHeaderGuardAfterComments) {
+  std::string Code = "// comment \n"
+                     "#ifndef X\n"
+                     "#define X\n"
+                     "int x;\n"
+                     "#define Y 1\n";
+  std::string Expected = "// comment \n"
+                         "#ifndef X\n"
+                         "#define X\n"
+                         "#include <vector>\n"
+                         "int x;\n"
+                         "#define Y 1\n";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
+TEST_F(HeaderIncludesTest, IfNDefWithNoDefine) {
+  std::string Code = "// comment \n"
+                     "#ifndef X\n"
+                     "int x;\n"
+                     "#define Y 1\n";
+  std::string Expected = "// comment \n"
+                         "#include <vector>\n"
+                         "#ifndef X\n"
+                         "int x;\n"
+                         "#define Y 1\n";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
+TEST_F(HeaderIncludesTest, FakeHeaderGuard) {
+  std::string Code = "// comment \n"
+                     "#ifndef X\n"
+                     "#define 1\n";
+  std::string Expected = "// comment \n"
+                         "#include <vector>\n"
+                         "#ifndef X\n"
+                         "#define 1\n";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
+TEST_F(HeaderIncludesTest, HeaderGuardWithComment) {
+  std::string Code = "// comment \n"
+                     "#ifndef X // comment\n"
+                     "// comment\n"
+                     "/* comment\n"
+                     "*/\n"
+                     "/* comment */ #define X\n"
+                     "int x;\n"
+                     "#define Y 1\n";
+  std::string Expected = "// comment \n"
+                         "#ifndef X // comment\n"
+                         "// comment\n"
+                         "/* comment\n"
+                         "*/\n"
+                         "/* comment */ #define X\n"
+                         "#include <vector>\n"
+                         "int x;\n"
+                         "#define Y 1\n";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
+TEST_F(HeaderIncludesTest, EmptyCode) {
+  std::string Code = "";
+  std::string Expected = "#include <vector>\n";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+}
+
+TEST_F(HeaderIncludesTest, SkipExistingHeaders) {
+  std::string Code = "#include \"a.h\"\n"
+                     "#include <vector>\n";
+  std::string Expected = "#include \"a.h\"\n"
+                         "#include <vector>\n";
+  EXPECT_EQ(Expected, insert(Code, "<vector>"));
+  EXPECT_EQ(Expected, insert(Code, "a.h"));
+  EXPECT_EQ(Expected, insert(Code, "\"a.h\""));
+}
+
+TEST_F(HeaderIncludesTest, NoInsertionAfterCode) {
+  std::string Code = "#include \"a.h\"\n"
+                     "void f() {}\n"
+                     "#include \"b.h\"\n";
+  std::string Expected = "#include \"a.h\"\n"
+                         "#include \"c.h\"\n"
+                         "void f() {}\n"
+                         "#include \"b.h\"\n";
+  EXPECT_EQ(Expected, insert(Code, "\"c.h\""));
+}
+
+TEST_F(HeaderIncludesTest, NoInsertionInStringLiteral) {
+  std::string Code = "#include \"a.h\"\n"
+                     "const char[] = R\"(\n"
+                     "#include \"b.h\"\n"
+                     ")\";\n";
+  std::string Expected = "#include \"a.h\"\n"
+                         "#include \"c.h\"\n"
+                         "const char[] = R\"(\n"
+                         "#include \"b.h\"\n"
+                         ")\";\n";
+  EXPECT_EQ(Expected, insert(Code, "\"c.h\""));
+}
+
+TEST_F(HeaderIncludesTest, NoInsertionAfterOtherDirective) {
+  std::string Code = "#include \"a.h\"\n"
+                     "#ifdef X\n"
+                     "#include \"b.h\"\n"
+                     "#endif\n";
+  std::string Expected = "#include \"a.h\"\n"
+                         "#include \"c.h\"\n"
+                         "#ifdef X\n"
+                         "#include \"b.h\"\n"
+                         "#endif\n";
+  EXPECT_EQ(Expected, insert(Code, "\"c.h\""));
+}
+
+TEST_F(HeaderIncludesTest, CanInsertAfterLongSystemInclude) {
+  std::string Code = "#include \"a.h\"\n"
+                     "// comment\n\n"
+                     "#include <a/b/c/d/e.h>\n";
+  std::string Expected = "#include \"a.h\"\n"
+                         "// comment\n\n"
+                         "#include <a/b/c/d/e.h>\n"
+                         "#include <x.h>\n";
+  EXPECT_EQ(Expected, insert(Code, "<x.h>"));
+}
+
+TEST_F(HeaderIncludesTest, CanInsertAfterComment) {
+  std::string Code = "#include \"a.h\"\n"
+                     "// Comment\n"
+                     "\n"
+                     "/* Comment */\n"
+                     "// Comment\n"
+                     "\n"
+                     "#include \"b.h\"\n";
+  std::string Expected = "#include \"a.h\"\n"
+                         "// Comment\n"
+                         "\n"
+                         "/* Comment */\n"
+                         "// Comment\n"
+                         "\n"
+                         "#include \"b.h\"\n"
+                         "#include \"c.h\"\n";
+  EXPECT_EQ(Expected, insert(Code, "c.h"));
+}
+
+TEST_F(HeaderIncludesTest, LongCommentsInTheBeginningOfFile) {
+  std::string Code = "// Loooooooooooooooooooooooooong comment\n"
+                     "// Loooooooooooooooooooooooooong comment\n"
+                     "// Loooooooooooooooooooooooooong comment\n"
+                     "#include <string>\n"
+                     "#include <vector>\n"
+                     "\n"
+                     "#include \"a.h\"\n"
+                     "#include \"b.h\"\n";
+  std::string Expected = "// Loooooooooooooooooooooooooong comment\n"
+                         "// Loooooooooooooooooooooooooong comment\n"
+                         "// Loooooooooooooooooooooooooong comment\n"
+                         "#include <string>\n"
+                         "#include <vector>\n"
+                         "\n"
+                         "#include \"a.h\"\n"
+                         "#include \"b.h\"\n"
+                         "#include \"third.h\"\n";
+  Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp)
+              .IncludeStyle;
+  EXPECT_EQ(Expected, insert(Code, "third.h"));
+}
+
+
+TEST_F(HeaderIncludesTest, SimpleDeleteIncludes) {
+  std::string Code = "#include \"abc.h\"\n"
+                     "#include \"xyz.h\" // comment\n"
+                     "int x;\n";
+  std::string Expected = "#include \"abc.h\"\n"
+                         "int x;\n";
+  EXPECT_EQ(Expected, remove(Code, "\"xyz.h\""));
+}
+
+TEST_F(HeaderIncludesTest, DeleteQuotedOnly) {
+  std::string Code = "#include \"abc.h\"\n"
+                     "#include <abc.h>\n"
+                     "int x;\n";
+  std::string Expected = "#include <abc.h>\n"
+                         "int x;\n";
+  EXPECT_EQ(Expected, remove(Code, "\"abc.h\""));
+}
+
+TEST_F(HeaderIncludesTest, DeleteAllIncludesWithSameNameIfNoType) {
+  std::string Code = "#include \"xyz.h\"\n"
+                     "#include \"xyz\"\n"
+                     "#include <xyz.h>\n";
+  std::string Expected = "#include \"xyz\"\n";
+  EXPECT_EQ(Expected, remove(Code, "xyz.h"));
+}
+
+TEST_F(HeaderIncludesTest, DeleteAllCode) {
+  std::string Code = "#include \"xyz.h\"\n"
+                     "#include <xyz.h>";
+  std::string Expected = "";
+  EXPECT_EQ(Expected, remove(Code, "xyz.h"));
+}
+
+TEST_F(HeaderIncludesTest, CanDeleteAfterCode) {
+  std::string Code = "#include \"a.h\"\n"
+                     "void f() {}\n"
+                     "#include \"b.h\"\n";
+  std::string Expected = "#include \"a.h\"\n"
+                         "void f() {}\n";
+  EXPECT_EQ(Expected, remove(Code, "b.h"));
+}
+
+} // namespace
+} // namespace tooling
+} // namespace clang