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

Summary: Also pull #include related style out of FormatStyle as tooling::IncludeStyle.

Reviewers: ilya-biryukov

Reviewed By: ilya-biryukov

Subscribers: klimek, mgorny, cfe-commits, djasper

Differential Revision: https://reviews.llvm.org/D46496

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@332287 91177308-0d34-0410-b5e6-96231b3b80d8

include/clang/Format/Format.h

@@ -19,6 +19,7 @@
 #include "clang/Tooling/Core/IncludeStyle.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/Support/Regex.h"
 #include <system_error>
 
 namespace clang {
@@ -1810,17 +1811,13 @@ formatReplacements(StringRef Code, const tooling::Replacements &Replaces,
 /// 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.
+/// The include manipulation is done via `tooling::HeaderInclude`, see its
+/// documentation for more details on how include insertion points are found and
+/// what edits are produced.
 llvm::Expected<tooling::Replacements>
 cleanupAroundReplacements(StringRef Code, const tooling::Replacements &Replaces,
                           const FormatStyle &Style);

include/clang/Tooling/Core/HeaderIncludes.h

@@ -0,0 +1,137 @@
+//===--- 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/Replacement.h"
+#include "clang/Tooling/Core/IncludeStyle.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/Regex.h"
+#include <unordered_map>
+
+namespace clang {
+namespace tooling {
+
+/// 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.
+  /// NOTE: this API is not thread-safe!
+  int getIncludePriority(StringRef IncludeName, bool CheckMainHeader) const;
+
+private:
+  bool isMainHeader(StringRef IncludeName) const;
+
+  const IncludeStyle Style;
+  bool IsMainFile;
+  std::string FileName;
+  // This refers to a substring in FileName.
+  StringRef FileStem;
+  // Regex is not thread-safe.
+  mutable SmallVector<llvm::Regex, 4> CategoryRegexs;
+};
+
+/// Generates replacements for inserting or deleting #include directives in a
+/// file.
+class HeaderIncludes {
+public:
+  HeaderIncludes(llvm::StringRef FileName, llvm::StringRef Code,
+                 const IncludeStyle &Style);
+
+  /// Inserts an #include directive of \p Header into the code. If \p IsAngled
+  /// is true, \p Header will be quoted with <> in the directive; otherwise, it
+  /// will be 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 of the same category. This respects the order of the existing
+  /// #includes in the block; if the existing #includes are not already sorted,
+  /// this will simply insert the #include in front of the first #include of the
+  /// same category in the code that should be sorted after \p IncludeName. If
+  /// \p IncludeName already exists (with exactly the same spelling), this
+  /// returns None.
+  llvm::Optional<tooling::Replacement> insert(llvm::StringRef Header,
+                                              bool IsAngled) const;
+
+  /// Removes all existing #includes of \p Header quoted with <> if \p IsAngled
+  /// is true or "" if \p IsAngled is false.
+  /// This doesn't resolve the header file path; it only deletes #includes with
+  /// exactly the same spelling.
+  tooling::Replacements remove(llvm::StringRef Header, bool IsAngled) const;
+
+private:
+  struct Include {
+    Include(StringRef Name, tooling::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.
+    tooling::Range R;
+  };
+
+  void addExistingInclude(Include IncludeToAdd, unsigned NextLineOffset);
+
+  std::string FileName;
+  std::string 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;
+
+  /// Map from priorities of #include categories to all #includes in the same
+  /// category. This is used to find #includes of the same category when
+  /// inserting new #includes. #includes in the same categories are sorted in
+  /// in the order they appear in the source file.
+  /// See comment for "FormatStyle::IncludeCategories" for details about include
+  /// priorities.
+  std::unordered_map<int, llvm::SmallVector<const Include *, 8>>
+      IncludesByPriority;
+
+  int FirstIncludeOffset;
+  // All new headers should be inserted after this offset (e.g. after header
+  // guards, file comment).
+  unsigned MinInsertOffset;
+  // Max insertion offset in the original code. For example, we want to avoid
+  // inserting new #includes into the actual code section (e.g. after a
+  // declaration).
+  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

@@ -31,6 +31,7 @@
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/VirtualFileSystem.h"
 #include "clang/Lex/Lexer.h"
+#include "clang/Tooling/Core/HeaderIncludes.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Allocator.h"
@@ -1676,62 +1677,6 @@ static void sortCppIncludes(const FormatStyle &Style,
 
 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.IncludeStyle.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.
-  // NOTE: this API is not thread-safe!
-  int 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.IncludeStyle.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.IncludeStyle.IncludeIsMainRegex).str(),
-          llvm::Regex::IgnoreCase);
-      if (MainIncludeRegex.match(FileStem))
-        return true;
-    }
-    return false;
-  }
-
-  const FormatStyle &Style;
-  bool IsMainFile;
-  StringRef FileName;
-  StringRef FileStem;
-  // Regex is not thread-safe.
-  mutable SmallVector<llvm::Regex, 4> CategoryRegexs;
-};
-
 const char IncludeRegexPattern[] =
     R"(^[\t\ ]*#[\t\ ]*(import|include)[^"<]*(["<][^">]*[">]))";
 
@@ -1755,7 +1700,7 @@ tooling::Replacements sortCppIncludes(const FormatStyle &Style, StringRef Code,
   //
   // 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;
@@ -1885,339 +1830,6 @@ 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) {
-  Environment Env(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;
-      });
-}
-
-/// Generates replacements for inserting or deleting #include directives in a
-/// file.
-class HeaderIncludes {
-public:
-  HeaderIncludes(llvm::StringRef FileName, llvm::StringRef Code,
-                 const FormatStyle &Style);
-
-  /// Inserts an #include directive of \p Header into the code. If \p IsAngled
-  /// is true, \p Header will be quoted with <> in the directive; otherwise, it
-  /// will be 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 of the same category. This respects the order of the existing
-  /// #includes in the block; if the existing #includes are not already sorted,
-  /// this will simply insert the #include in front of the first #include of the
-  /// same category in the code that should be sorted after \p IncludeName. If
-  /// \p IncludeName already exists (with exactly the same spelling), this
-  /// returns None.
-  llvm::Optional<tooling::Replacement> insert(llvm::StringRef Header,
-                                              bool IsAngled) const;
-
-  /// Removes all existing #includes of \p Header quoted with <> if \p IsAngled
-  /// is true or "" if \p IsAngled is false.
-  /// This doesn't resolve the header file path; it only deletes #includes with
-  /// exactly the same spelling.
-  tooling::Replacements remove(llvm::StringRef Header, bool IsAngled) const;
-
-private:
-  struct Include {
-    Include(StringRef Name, tooling::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.
-    tooling::Range R;
-  };
-
-  void addExistingInclude(Include IncludeToAdd, unsigned NextLineOffset);
-
-  std::string FileName;
-  std::string 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;
-
-  /// Map from priorities of #include categories to all #includes in the same
-  /// category. This is used to find #includes of the same category when
-  /// inserting new #includes. #includes in the same categories are sorted in
-  /// in the order they appear in the source file.
-  /// See comment for "FormatStyle::IncludeCategories" for details about include
-  /// priorities.
-  std::unordered_map<int, llvm::SmallVector<const Include *, 8>>
-      IncludesByPriority;
-
-  int FirstIncludeOffset;
-  // All new headers should be inserted after this offset (e.g. after header
-  // guards, file comment).
-  unsigned MinInsertOffset;
-  // Max insertion offset in the original code. For example, we want to avoid
-  // inserting new #includes into the actual code section (e.g. after a
-  // declaration).
-  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;
-};
-
-HeaderIncludes::HeaderIncludes(StringRef FileName, StringRef Code,
-                               const FormatStyle &Style)
-    : FileName(FileName), Code(Code), FirstIncludeOffset(-1),
-      MinInsertOffset(
-          getOffsetAfterHeaderGuardsAndComments(FileName, Code, Style)),
-      MaxInsertOffset(MinInsertOffset +
-                      getMaxHeaderInsertionOffset(
-                          FileName, Code.drop_front(MinInsertOffset), 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.IncludeStyle.IncludeCategories)
-    Priorities.insert(Category.Priority);
-  SmallVector<StringRef, 32> Lines;
-  Code.drop_front(MinInsertOffset).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],
-                  tooling::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)];
-}
-
-inline StringRef trimInclude(StringRef IncludeName) {
-  return IncludeName.trim("\"<>");
-}
-
-// \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<tooling::Replacement>
-HeaderIncludes::insert(llvm::StringRef IncludeName, bool IsAngled) const {
-  assert(IncludeName == trimInclude(IncludeName));
-  // If a <header> ("header") already exists in code, "header" (<header>) with
-  // different quotation will still be inserted.
-  // FIXME: figure out if this is the best behavior.
-  auto It = ExistingIncludes.find(IncludeName);
-  if (It != ExistingIncludes.end())
-    for (const auto &Inc : It->second)
-      if ((IsAngled && StringRef(Inc.Name).startswith("<")) ||
-          (!IsAngled && StringRef(Inc.Name).startswith("\"")))
-        return llvm::None;
-  std::string Quoted = IsAngled ? ("<" + 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();
-  // When inserting headers at end of the code, also append '\n' to the code
-  // if it does not end with '\n'.
-  // FIXME: when inserting multiple #includes at the end of code, only one
-  // newline should be added.
-  if (InsertOffset == Code.size() && (!Code.empty() && Code.back() != '\n'))
-    NewInclude = "\n" + NewInclude;
-  return tooling::Replacement(FileName, InsertOffset, 0, NewInclude);
-}
-
-tooling::Replacements HeaderIncludes::remove(llvm::StringRef IncludeName,
-                                             bool IsAngled) const {
-  assert(IncludeName == trimInclude(IncludeName));
-  tooling::Replacements Result;
-  auto Iter = ExistingIncludes.find(IncludeName);
-  if (Iter == ExistingIncludes.end())
-    return Result;
-  for (const auto &Inc : Iter->second) {
-    if ((IsAngled && StringRef(Inc.Name).startswith("\"")) ||
-        (!IsAngled && StringRef(Inc.Name).startswith("<")))
-      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;
-}
-
 // FIXME: insert empty lines between newly created blocks.
 tooling::Replacements
 fixCppIncludeInsertions(StringRef Code, const tooling::Replacements &Replaces,
@@ -2248,11 +1860,11 @@ fixCppIncludeInsertions(StringRef Code, const tooling::Replacements &Replaces,
 
 
   StringRef FileName = Replaces.begin()->getFilePath();
-  HeaderIncludes Includes(FileName, Code, Style);
+  tooling::HeaderIncludes Includes(FileName, Code, Style.IncludeStyle);
 
   for (const auto &Header : HeadersToDelete) {
     tooling::Replacements Replaces =
-        Includes.remove(trimInclude(Header), Header.startswith("<"));
+        Includes.remove(Header.trim("\"<>"), Header.startswith("<"));
     for (const auto &R : Replaces) {
       auto Err = Result.add(R);
       if (Err) {
@@ -2274,7 +1886,7 @@ fixCppIncludeInsertions(StringRef Code, const tooling::Replacements &Replaces,
     (void)Matched;
     auto IncludeName = Matches[2];
     auto Replace =
-        Includes.insert(trimInclude(IncludeName), IncludeName.startswith("<"));
+        Includes.insert(IncludeName.trim("\"<>"), IncludeName.startswith("<"));
     if (Replace) {
       auto Err = Result.add(*Replace);
       if (Err) {

lib/Tooling/Core/CMakeLists.txt

@@ -1,10 +1,11 @@
 set(LLVM_LINK_COMPONENTS support)
 
 add_clang_library(clangToolingCore
+  Diagnostic.cpp
+  HeaderIncludes.cpp
   IncludeStyle.cpp
   Lookup.cpp
   Replacement.cpp
-  Diagnostic.cpp
 
   LINK_LIBS
   clangAST

lib/Tooling/Core/HeaderIncludes.cpp

@@ -0,0 +1,330 @@
+//===--- 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/Basic/SourceManager.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) {
+  SourceManagerForFile VirtualSM(FileName, Code);
+  SourceManager &SM = VirtualSM.get();
+  Lexer Lex(SM.getMainFileID(), SM.getBuffer(SM.getMainFileID()), SM,
+            createLangOpts());
+  Token Tok;
+  // Get the first token.
+  Lex.LexFromRawLexer(Tok);
+  return GetOffsetAfterSequence(SM, 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)),
+      MaxInsertOffset(MinInsertOffset +
+                      getMaxHeaderInsertionOffset(
+                          FileName, Code.drop_front(MinInsertOffset), 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;
+  Code.drop_front(MinInsertOffset).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],
+                  tooling::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<tooling::Replacement>
+HeaderIncludes::insert(llvm::StringRef IncludeName, bool IsAngled) const {
+  assert(IncludeName == trimInclude(IncludeName));
+  // If a <header> ("header") already exists in code, "header" (<header>) with
+  // different quotation will still be inserted.
+  // FIXME: figure out if this is the best behavior.
+  auto It = ExistingIncludes.find(IncludeName);
+  if (It != ExistingIncludes.end())
+    for (const auto &Inc : It->second)
+      if ((IsAngled && StringRef(Inc.Name).startswith("<")) ||
+          (!IsAngled && StringRef(Inc.Name).startswith("\"")))
+        return llvm::None;
+  std::string Quoted = IsAngled ? ("<" + 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();
+  // When inserting headers at end of the code, also append '\n' to the code
+  // if it does not end with '\n'.
+  // FIXME: when inserting multiple #includes at the end of code, only one
+  // newline should be added.
+  if (InsertOffset == Code.size() && (!Code.empty() && Code.back() != '\n'))
+    NewInclude = "\n" + NewInclude;
+  return tooling::Replacement(FileName, InsertOffset, 0, NewInclude);
+}
+
+tooling::Replacements HeaderIncludes::remove(llvm::StringRef IncludeName,
+                                             bool IsAngled) const {
+  assert(IncludeName == trimInclude(IncludeName));
+  tooling::Replacements Result;
+  auto Iter = ExistingIncludes.find(IncludeName);
+  if (Iter == ExistingIncludes.end())
+    return Result;
+  for (const auto &Inc : Iter->second) {
+    if ((IsAngled && StringRef(Inc.Name).startswith("\"")) ||
+        (!IsAngled && StringRef(Inc.Name).startswith("<")))
+      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

@@ -358,175 +358,6 @@ TEST_F(CleanUpReplacementsTest, FixOnlyAffectedCodeAfterReplacements) {
   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, 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";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include \"b.h\"")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, 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";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include \"d.h\"")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, 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";
-  tooling::Replacements Replaces =
-      toReplacements({createInsertion("#include <vector>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-
 TEST_F(CleanUpReplacementsTest, InsertMultipleIncludesLLVMStyle) {
   std::string Code = "#include \"x/fix.h\"\n"
                      "#include \"a.h\"\n"
@@ -548,24 +379,6 @@ TEST_F(CleanUpReplacementsTest, InsertMultipleIncludesLLVMStyle) {
   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"
@@ -623,6 +436,19 @@ TEST_F(CleanUpReplacementsTest, InsertMultipleNewHeadersAndSortGoogle) {
   EXPECT_EQ(Expected, formatAndApply(Code, Replaces));
 }
 
+TEST_F(CleanUpReplacementsTest, NoNewLineAtTheEndOfCodeMultipleInsertions) {
+  std::string Code = "#include <map>";
+  // FIXME: a better behavior is to only append on newline to Code, but this
+  // case should be rare in practice.
+  std::string Expected =
+      "#include <map>\n#include <string>\n\n#include <vector>\n";
+  tooling::Replacements Replaces =
+      toReplacements({createInsertion("#include <string>"),
+                      createInsertion("#include <vector>")});
+  EXPECT_EQ(Expected, apply(Code, Replaces));
+}
+
+
 TEST_F(CleanUpReplacementsTest, FormatCorrectLineWhenHeadersAreInserted) {
   std::string Code = "\n"
                      "int x;\n"
@@ -648,227 +474,6 @@ TEST_F(CleanUpReplacementsTest, FormatCorrectLineWhenHeadersAreInserted) {
   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>";
-  // FIXME: a better behavior is to only append on newline to Code, but this
-  // case should be rare in practice.
-  std::string Expected =
-      "#include <map>\n#include <string>\n\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 <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, SimpleDeleteIncludes) {
   std::string Code = "#include \"abc.h\"\n"
                      "#include \"xyz.h\" // comment\n"
@@ -881,35 +486,6 @@ TEST_F(CleanUpReplacementsTest, SimpleDeleteIncludes) {
   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\""), createDeletion("<xyz.h>")});
-  EXPECT_EQ(Expected, apply(Code, Replaces));
-}
-
-TEST_F(CleanUpReplacementsTest, DeleteOnlyIncludesWithSameQuote) {
-  std::string Code = "#include \"xyz.h\"\n"
-                     "#include \"xyz\"\n"
-                     "#include <xyz.h>\n";
-  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, 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"
@@ -922,117 +498,6 @@ TEST_F(CleanUpReplacementsTest, InsertionAndDeleteHeader) {
   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/Tooling/CMakeLists.txt

@@ -18,6 +18,7 @@ add_clang_unittest(ToolingTests
   DiagnosticsYamlTest.cpp
   ExecutionTest.cpp
   FixItTest.cpp
+  HeaderIncludesTest.cpp
   LexicallyOrderedRecursiveASTVisitorTest.cpp
   LookupTest.cpp
   QualTypeNamesTest.cpp

unittests/Tooling/HeaderIncludesTest.cpp

@@ -0,0 +1,527 @@
+//===- 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);
+    assert(Header.startswith("\"") || Header.startswith("<"));
+    auto R = Includes.insert(Header.trim("\"<>"), Header.startswith("<"));
+    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);
+    assert(Header.startswith("\"") || Header.startswith("<"));
+    auto Replaces = Includes.remove(Header.trim("\"<>"), Header.startswith("<"));
+    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, NoNewLineAtTheEndOfCode) {
+  std::string Code = "#include <map>";
+  std::string Expected = "#include <map>\n#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\""));
+}
+
+TEST_F(HeaderIncludesTest, AddIncludesWithDifferentForms) {
+  std::string Code = "#include <vector>\n";
+  // FIXME: this might not be the best behavior.
+  std::string Expected = "#include \"vector\"\n"
+                         "#include <vector>\n";
+  EXPECT_EQ(Expected, insert(Code, "\"vector\""));
+}
+
+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, SimpleDeleteInclude) {
+  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, DeleteAllCode) {
+  std::string Code = "#include \"xyz.h\"\n";
+  std::string Expected = "";
+  EXPECT_EQ(Expected, remove(Code, "\"xyz.h\""));
+}
+
+TEST_F(HeaderIncludesTest, DeleteOnlyIncludesWithSameQuote) {
+  std::string Code = "#include \"xyz.h\"\n"
+                     "#include \"xyz\"\n"
+                     "#include <xyz.h>\n";
+  std::string Expected = "#include \"xyz.h\"\n"
+                         "#include \"xyz\"\n";
+  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