[clang-format] insert new #includes into correct blocks when cleaning up Replacement with cleanupAroundReplacements().

[ioeric]

Summary

When a replacement's offset is set to UINT_MAX or -1U, it is treated as
a header insertion replacement by cleanupAroundReplacements(). The new #include
directive is then inserted into the correct block.

[djasper]

lib/Format/Format.cpp

1444

But you should pass it as an ArrayRef or a SmallVectorImpl.

1446

nit: no braces

1451

From the docs in Format.h: "If none of the regular expressions match, INT_MAX is assigned". If that is true for an existing #include as well as the inserted one, I don't see a reason, why those should belong to the same block. So, I think you should return INT_MAX here and remove all the special-casing below.

1464

We already have an include regex in this file, I don't think we should have two. Can you pull out the string into a constant?

1468

There is a lot of duplicated code between this and sortCppIncludes. Pull common code out into functions.

1487

There is no guarantee that people assign categories sequentially, I can easily see people assigning category "1000" or "10000" to be able to add more categories later. I think this should be a map and I think you then also don't really need the vector. Also note that categories can be negative, which I think would break this.

1491

I am guessing that AfterFirstDefine is used so that we insert #includes after header guards. I think we should name it AfterHeaderGuard than and either make the behavior more strict or at least add a FIXME for it (only in header files, only if the #define is after an #ifndef, ...)

1515

I think if you assign the EndOffset of all lower categories that aren't set yet, that will simplify the logic below.

1542

What happens if there is no category with higher priority, but there is a category with lower priority. Shouldn't you then insert the #include right before that?

1581

I think you can remove this FIXME.. Seems obvious.

1584

Why do you split out all the header insertion replacements here instead of just ignoring the ones that aren't header insertions in fixCppIncludeInsertions?

1587

nit: No braces.

[ioeric]

• Addressed reviewer's comments.

[ioeric]

lib/Format/Format.cpp

1584

I think this enables us to exit early in fixCppIncludeInsertions if there is no header insertion replacement at all. Also, it is easier to split out header insertion replacements and then merge fixed replacements with other replacements instead of doing this in place, which IMO might not be more efficient since there are still set deletion and insertion for each header insertion replacement.

[bkramer]

lib/Format/Format.cpp

1283

Typo: IncldueName

1502

Typo: intialized

[djasper]

include/clang/Format/Format.h

778

Why don't we just stick with UINT_MAX everywhere?

lib/Format/Format.cpp

1456

Unused

1465

Why not just a set?

1487

nit: s/more strict/stricter/

1493

Are you sure that this is what's implemented? The way I see it is:

• Ensure that EndOffset[Highest] is always populated.
• If EndOffset[Priority] isn't set, use the next higher value that is set, up to EndOffset[Highest].

1512

I'd leave all of this logic out. This only affects headers with no current #includes, which is rare to begin with and on top of that, this is more to do with whether or not we want to automatically add new lines between the different blocks. As it currently is, this would do both:

#define HEADER_GUARD

#include <a> // inserted

and

#define HEADER_GUARD
#include <a> // inserted
#include <b>

Which might be the right behavior, but I think this is the wrong place to worry about it. In the long run, we should probably make clang-format inserted empty lines around the #include blocks if desired.

1524

Add a comment:

// 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.

1530

This should be done together with the check for the offset being -1u. Eventually we are going to have something else that we are going to insert (e.g. using declarations) and we don't want to confuse those. And even more so, if all replacements at offset -1u are not #includes, we don't need to execute most of this function.

1537

Why is CheckMainHeader false?

1540

No parentheses.

1549

Hm.. Not happy with the naming. There is no real difference between fixHeaderInsertions and fixCppIncludeInsertions. You could argue that the "Cpp" part is the difference, but the actual difference is that one iterates over an entire set of replacements whereas the other only modifies header insertions.

I think I'd just merge the two functions, but that probably also ties back to my wish not to separate out the two replacement sets ;-).

[ioeric]

• Addressed reviewers' comments.

lib/Format/Format.cpp

1549

How about fixInsertionReplacements?

[djasper]

lib/Format/Format.cpp

1284

Nit: Otherwise, returns the priority of the matching category or INT_MAX.

1502

nit: I think you should be using "I" and "E".

1507

Can you just use "std::prev(I)" instead of Prev?

1549

So, not knowing anything of it, what's the difference between fixCppIncludeInsertions and fixInsertionReplacements and when do you need to call what. To me the fact that it's that difficult to find a name for it is telling..

Also, I still don't understand why you think it is better to separate the two functions. In that way you are currently implementing it, it is certainly less efficient. You are always copying all the replacements, which (in the way you are doing it) is O(N lg N), even if there isn't a single header insertion. Maybe at least use set_difference (similar to what I am showing below).

So, how I would structure it:

• Pull out a function isHeaderInsertion(const Replacement &). Easy naming, good to pull this functionality out.
• At the start of fixCppIncludeInsertions, write:

tooling::Replacements HeaderInsertions;
for (const auto &R : Replaces)
  if (isHeaderInsertion(R))
    HeaderInsertions.insert(R);
if (HeaderInsertions.empty())
  return Replaces;

tooling::Replacements Result;
std::set_difference(Replaces.begin(), Replaces.end(),
                    HeaderInsertions.begin(), HeaderInsertions.end(),
                    Result.begin());

• Do what the function currently does with HeaderInsertions and add the results back to Result.

[ioeric]

• Addressed reviewer's comments.

lib/Format/Format.cpp

1549

I really appreciate the way you structure it, which does make the code clearer.

Just one coveat: std::set_difference() is not applicable on Replacements since it calls std::copy() on Replacement, which Replacement does not support.

I think another way to approach it would be to delete header insertion replacement and insert new replacements back if we can assume that there are only few header insertion replacements.

[ioeric]

• Removed inline from isHeaderInsertion.

[djasper]

lib/Format/Format.cpp

1548

So, add a copy constructor to Replacement?

[ioeric]

• Use std::set_difference in fixCppIncludeInsertions()

[djasper]

lib/Format/Format.cpp

1444

This error output doesn't belong here. I think we can just remove it. If you'd prefer to keep it, add it to the loop in fixCppIncludeInsertions():

for (const auto &R : Replaces) {
  if (isHeaderInsertion(R))
    HeaderInsertions.insert(R);
  else if (R.getOffset() == UINT_MAX)
    log::errs() << ...
}

1556

I'd move this check into fixCppIncludeInsertions.

unittests/Format/CleanupTest.cpp

310

I'd pull out a lot of these environment setup things into abstractions in the test fixture. Maybe all you need is two functions like:

insert(Code, Replaces);

and

insertAndFormat(Code, Replaces);

?

311

Can you just make this:

tooling::Replacements Replaces = {
    tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"b.h\"")};

?

510

I'd create a function:

tooling::Replacement createInsertion(StringRef HeaderName) {
  return tooling::Replacement("fix.cpp", UINT_MAX, 0, HeaderName);
}

[ioeric]

• Addressed commments.

unittests/Format/CleanupTest.cpp

310

I had tried to pull out the environment setup...but it didn't seem to save much code without sacrificing flexibility.

In some test cases, we also need to add replacements other than header insertions, and those replacements construction uses FileID, so createInMemoryFile needs to stay in test functions.

FormatStyle differs among test cases, so it has to stay.

All we save might just be the two lines which could be either cleanup or cleanup+format.\

Any better ideas?

[djasper]

unittests/Format/CleanupTest.cpp

310

Well, the only reason you seem to be using the FileID or "Context" for that matter is to translate between line/col and offset. I'd just pull that functionality out into a separate function (which does it based on the "Code" input) or not at all (hard-coding the offset doesn't seem *that* bad).

510

Have you seen this comment?

[ioeric]

• Addressed reviewer's comments.

unittests/Format/CleanupTest.cpp

310

Ohh, I see! Thanks!

510

Sorry that I missed this one...

[djasper]

A few nitpicks, but otherwise looks good.

lib/Format/Format.cpp

1287

I'd consistently try to use i and e for integers vs. I and E for iterators.

unittests/Format/CleanupTest.cpp

257

I'd probably make "Style" a class member, defaulting to LLVM Style that you can then overwrite for other styles.

271

I'd consider making "Code" a member variable. But it has pros and cons.

297

Why are you starting and ending the header guards with a double underscore? I think this is actually reserved for built-in identifiers. Maybe just stick with the LLVM naming scheme?

[ioeric]

• Nits fixed.

include/clang/Format/Format.h

 /// \brief Returns the replacements corresponding to applying and formatting
 /// \p Replaces.
 tooling::Replacements 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.
+/// This also inserts a C++ #include directive into the correct block if the
+/// replacement corresponding to the header insertion has offset UINT_MAX.
 tooling::Replacements
 cleanupAroundReplacements(StringRef Code, const tooling::Replacements &Replaces,
                           const FormatStyle &Style);
 
 /// \brief Reformats the given \p Ranges in the file \p ID.
 ///
 /// Each range is extended on either end to its next bigger logic unit, i.e.
 /// everything that might influence its formatting or might be influenced by its
 /// formatting.

lib/Format/Format.cpp

static void sortCppIncludes(const FormatStyle &Style,
   assert(result.size() ==
          Includes.back().Offset + Includes.back().Text.size() -
              Includes.front().Offset);
 
   Replaces.insert(tooling::Replacement(FileName, Includes.front().Offset,
                                        result.size(), result));
 }
 
+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);
+    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)) {
+      llvm::Regex MainIncludeRegex(
+          (HeaderStem + Style.IncludeIsMainRegex).str());
+      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(
-      R"(^[\t\ ]*#[\t\ ]*(import|include)[^"<]*(["<][^">]*[">]))");
+  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.
-  bool IsSource = FileName.endswith(".c") || FileName.endswith(".cc") ||
-                  FileName.endswith(".cpp") || FileName.endswith(".c++") ||
-                  FileName.endswith(".cxx") || FileName.endswith(".m") ||
-                  FileName.endswith(".mm");
-  StringRef FileStem = llvm::sys::path::stem(FileName);
+  IncludeCategoryManager Categories(Style, FileName);
   bool FirstIncludeBlock = true;
   bool MainIncludeFound = false;
-
-  // Create pre-compiled regular expressions for the #include categories.
-  SmallVector<llvm::Regex, 4> CategoryRegexs;
-  for (const auto &Category : Style.IncludeCategories)
-    CategoryRegexs.emplace_back(Category.Regex);
-
   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;
 
     if (!FormattingOff && !Line.endswith("\\")) {
       if (IncludeRegex.match(Line, &Matches)) {
         StringRef IncludeName = Matches[2];
-        int Category = INT_MAX;
-        for (unsigned i = 0, e = CategoryRegexs.size(); i != e; ++i) {
-          if (CategoryRegexs[i].match(IncludeName)) {
-            Category = Style.IncludeCategories[i].Priority;
-            break;
-          }
-        }
-        if (IsSource && !MainIncludeFound && Category > 0 &&
-            FirstIncludeBlock && IncludeName.startswith("\"")) {
-          StringRef HeaderStem =
-              llvm::sys::path::stem(IncludeName.drop_front(1).drop_back(1));
-          if (FileStem.startswith(HeaderStem)) {
-            llvm::Regex MainIncludeRegex(
-                (HeaderStem + Style.IncludeIsMainRegex).str());
-            if (MainIncludeRegex.match(FileStem)) {
-              Category = 0;
-              MainIncludeFound = true;
-            }
-          }
-        }
+        int Category = Categories.getIncludePriority(
+            IncludeName,
+            /*CheckMainHeader=*/!MainIncludeFound && FirstIncludeBlock);
+        if (Category == 0)
         IncludesInBlock.push_back({IncludeName, Line, Prev, Category});
       } else if (!IncludesInBlock.empty()) {
         sortCppIncludes(Style, IncludesInBlock, Ranges, FileName, Replaces,
                         Cursor);
         IncludesInBlock.clear();
         FirstIncludeBlock = false;
       }
       Prev = Pos + 1;
tooling::Replacements formatReplacements(StringRef Code,
   auto Reformat = [](const FormatStyle &Style, StringRef Code,
                      std::vector<tooling::Range> Ranges,
                      StringRef FileName) -> tooling::Replacements {
     return reformat(Style, Code, Ranges, FileName);
   };
   return processReplacements(Reformat, Code, SortedReplaces, Style);
 }
 
+namespace {
+
+bool isHeaderInsertion(const tooling::Replacement &Replace) {
+  if (Replace.getOffset() != UINT_MAX)
+    return false;
+  if (!llvm::Regex(IncludeRegexPattern).match(Replace.getReplacementText())) {
+    llvm::errs() << "Insertions other than header #include insertion are "
+                    "not supported! "
+                 << Replace.getReplacementText() << "\n";
+    return false;
+  }
+  return true;
+}
+
+// FIXME: do not insert headers into conditional #include blocks, e.g. #includes
+// surrounded by compile condition "#if...".
+// FIXME: do not insert existing headers.
+// FIXME: insert empty lines between newly created blocks.
+tooling::Replacements
+fixCppIncludeInsertions(StringRef Code, const tooling::Replacements &Replaces,
+                        const FormatStyle &Style) {
+  tooling::Replacements HeaderInsertions;
+  for (const auto &R : Replaces)
+    if (isHeaderInsertion(R))
+      HeaderInsertions.insert(R);
+  if (HeaderInsertions.empty())
+    return Replaces;
+  tooling::Replacements Result;
+  std::set_difference(Replaces.begin(), Replaces.end(),
+                      HeaderInsertions.begin(), HeaderInsertions.end(),
+                      std::inserter(Result, Result.begin()));
+
+  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);
+
+  // 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;
+  int Offset = 0;
+  int AfterHeaderGuard = 0;
+  SmallVector<StringRef, 32> Lines;
+  Code.split(Lines, '\n');
+  for (auto Line : Lines) {
+    if (IncludeRegex.match(Line, &Matches)) {
+      StringRef IncludeName = Matches[2];
+      int Category = Categories.getIncludePriority(
+          IncludeName, /*CheckMainHeader=*/FirstIncludeOffset < 0);
+      CategoryEndOffsets[Category] = Offset + Line.size() + 1;
+      if (FirstIncludeOffset < 0)
+        FirstIncludeOffset = Offset;
+    }
+    Offset += Line.size() + 1;
+    // FIXME: make header guard matching stricter, e.g. consider #ifndef.
+    if (AfterHeaderGuard == 0 && DefineRegex.match(Line))
+      AfterHeaderGuard = Offset;
+  }
+
+  // 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].
+  // FIXME: skip comment section in the beginning of the code if there is no
+  // existing #include and #define.
+  auto Highest = Priorities.begin();
+  if (CategoryEndOffsets.find(*Highest) == CategoryEndOffsets.end()) {
+    if (FirstIncludeOffset >= 0)
+      CategoryEndOffsets[*Highest] = FirstIncludeOffset;
+    else
+      CategoryEndOffsets[*Highest] = AfterHeaderGuard;
+  }
+  // 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 &R : HeaderInsertions) {
+    auto IncludeDirective = R.getReplacementText();
+    bool Matched = IncludeRegex.match(IncludeDirective, &Matches);
+    assert(Matched && "Header insertion replacement must have replacement text "
+                      "'#include ...'");
+    auto IncludeName = Matches[2];
+    int Category =
+        Categories.getIncludePriority(IncludeName, /*CheckMainHeader=*/true);
+    Offset = CategoryEndOffsets[Category];
+    std::string NewInclude = !IncludeDirective.endswith("\n")
+                                 ? (IncludeDirective + "\n").str()
+                                 : IncludeDirective.str();
+    Result.insert(tooling::Replacement(FileName, Offset, 0, NewInclude));
+  }
+  return Result;
+}
+
+} // anonymous namespace
+
 tooling::Replacements
 cleanupAroundReplacements(StringRef Code, const tooling::Replacements &Replaces,
                           const FormatStyle &Style) {
   // We need to use lambda function here since there are two versions of
   // `cleanup`.
   auto Cleanup = [](const FormatStyle &Style, StringRef Code,
                     std::vector<tooling::Range> Ranges,
                     StringRef FileName) -> tooling::Replacements {
     return cleanup(Style, Code, Ranges, FileName);
   };
-  return processReplacements(Cleanup, Code, Replaces, Style);
+  // Make header insertion replacements insert new headers into correct blocks.
+  tooling::Replacements NewReplaces =
+      (Style.Language != FormatStyle::LanguageKind::LK_Cpp)
+          ? Replaces
+          : fixCppIncludeInsertions(Code, Replaces, Style);
+  return processReplacements(Cleanup, Code, NewReplaces, Style);
 }
 
 tooling::Replacements reformat(const FormatStyle &Style, SourceManager &SM,
                                FileID ID, ArrayRef<CharSourceRange> Ranges,
                                bool *IncompleteFormat) {
   FormatStyle Expanded = expandPresets(Style);
   if (Expanded.DisableFormat)
     return tooling::Replacements();
 
   Environment Env(SM, ID, Ranges);
   Formatter Format(Env, Expanded, IncompleteFormat);
   return Format.process();
 }
 
 tooling::Replacements reformat(const FormatStyle &Style, StringRef Code,
                                ArrayRef<tooling::Range> Ranges,
                                StringRef FileName, bool *IncompleteFormat) {
   FormatStyle Expanded = expandPresets(Style);
   if (Expanded.DisableFormat)
     return tooling::Replacements();
 
   std::unique_ptr<Environment> Env =
       Environment::CreateVirtualEnvironment(Code, FileName, Ranges);
   Formatter Format(*Env, Expanded, IncompleteFormat);
   return Format.process();
 }
 
 tooling::Replacements cleanup(const FormatStyle &Style, SourceManager &SM,
                               FileID ID, ArrayRef<CharSourceRange> Ranges) {
   Environment Env(SM, ID, Ranges);
   Cleaner Clean(Env, Style);
   return Clean.process();
 }
 
 tooling::Replacements cleanup(const FormatStyle &Style, StringRef Code,
                               ArrayRef<tooling::Range> Ranges,

unittests/Format/CleanupTest.cpp

return tooling::Replacement(Context.Sources, Start, Length,
                                 ReplacementText);
   }
 
   RewriterTestContext Context;
 };
 
 TEST_F(CleanUpReplacementsTest, FixOnlyAffectedCodeAfterReplacements) {
   std::string Code = "namespace A {\n"
                      "namespace B {\n"
                      "  int x;\n"
                      "} // namespace B\n"
                      "} // namespace A\n"
                      "\n"
                      "namespace C {\n"
                      "namespace D { int i; }\n"
                      "inline namespace E { namespace { int y; } }\n"
                      "int x=     0;"
                      "}";
   std::string Expected = "\n\nnamespace C {\n"
                          "namespace D { int i; }\n\n"
                          "int x=     0;"
                          "}";
   FileID ID = Context.createInMemoryFile("fix.cpp", Code);
   tooling::Replacements Replaces;
   Replaces.insert(tooling::Replacement(Context.Sources,
                                        Context.getLocation(ID, 3, 3), 6, ""));
   Replaces.insert(tooling::Replacement(Context.Sources,
                                        Context.getLocation(ID, 9, 34), 6, ""));
 
   format::FormatStyle Style = format::getLLVMStyle();
   auto FinalReplaces = formatReplacements(
       Code, cleanupAroundReplacements(Code, Replaces, Style), Style);
   EXPECT_EQ(Expected, applyAllReplacements(Code, FinalReplaces));
 }
 
+TEST_F(CleanUpReplacementsTest, NoExistingIncludeWithoutDefine) {
+  std::string Code = "int main() {}";
+  std::string Expected = "#include \"a.h\"\n"
+                         "int main() {}";
+  Context.createInMemoryFile("fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"a.h\""));
+  format::FormatStyle Style = format::getLLVMStyle();
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(Code, NewReplaces));
+}
+
+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";
+
+  Context.createInMemoryFile("fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"b.h\""));
+  format::FormatStyle Style = format::getLLVMStyle();
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(Code, NewReplaces));
+}
+
+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";
+
+  Context.createInMemoryFile("fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"a.h\""));
+  format::FormatStyle Style = format::getLLVMStyle();
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(Code, NewReplaces));
+}
+
+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() {}";
+  Context.createInMemoryFile("fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include <a>"));
+  format::FormatStyle Style =
+      format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(Code, NewReplaces));
+}
+
+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() {}";
+  Context.createInMemoryFile("fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"z.h\""));
+  format::FormatStyle Style = format::getLLVMStyle();
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(Code, NewReplaces));
+}
+
+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() {}";
+  Context.createInMemoryFile("fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"z.h\""));
+  format::FormatStyle Style =
+      format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(Code, NewReplaces));
+}
+
+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";
+  Context.createInMemoryFile("fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"d.h\""));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"llvm/x/y.h\""));
+  format::FormatStyle Style = format::getLLVMStyle();
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(Code, NewReplaces));
+}
+
+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";
+  Context.createInMemoryFile("fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include <list>"));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"new/new.h\""));
+  format::FormatStyle Style = format::getLLVMStyle();
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(Code, NewReplaces));
+}
+
+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
+  // prefered.
+  std::string Expected = "#include \"x/fix.h\"\n"
+                         "#include <vector>\n"
+                         "\n"
+                         "#include \"y/a.h\"\n"
+                         "#include \"z/b.h\"\n";
+  Context.createInMemoryFile("x/fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include <vector>"));
+  format::FormatStyle Style =
+      format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(Code, NewReplaces));
+}
+
+TEST_F(CleanUpReplacementsTest, InsertMultipleIncludesGoogleStyle) {
+  std::string Code = "#include \"x/fix.h\"\n"
+                     "\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 <list>\n"
+                         "\n"
+                         "#include \"y/a.h\"\n"
+                         "#include \"z/b.h\"\n"
+                         "#include \"x/x.h\"\n";
+  Context.createInMemoryFile("x/fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include <list>"));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"x/x.h\""));
+  format::FormatStyle Style =
+      format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(Code, NewReplaces));
+}
+
+TEST_F(CleanUpReplacementsTest, InsertMultipleNewHeadersAndSortLLVM) {
+  std::string Code = "\nint x;";
+  std::string Expected = "#include \"fix.h\"\n"
+                         "#include \"a.h\"\n"
+                         "#include \"b.h\"\n"
+                         "#include \"c.h\"\n"
+                         "#include <list>\n"
+                         "#include <vector>\n"
+                         "\nint x;";
+  Context.createInMemoryFile("fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"a.h\""));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"c.h\""));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"b.h\""));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include <vector>"));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include <list>"));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"fix.h\""));
+  format::FormatStyle Style = format::getLLVMStyle();
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(
+                          Code, formatReplacements(Code, NewReplaces, Style)));
+}
+
+TEST_F(CleanUpReplacementsTest, InsertMultipleNewHeadersAndSortGoogle) {
+  std::string Code = "\nint x;";
+  std::string Expected = "#include \"fix.h\"\n"
+                         "#include <list>\n"
+                         "#include <vector>\n"
+                         "#include \"a.h\"\n"
+                         "#include \"b.h\"\n"
+                         "#include \"c.h\"\n"
+                         "\nint x;";
+  Context.createInMemoryFile("fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"a.h\""));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"c.h\""));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"b.h\""));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include <vector>"));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include <list>"));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"fix.h\""));
+  format::FormatStyle Style =
+      format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(
+                          Code, formatReplacements(Code, NewReplaces, Style)));
+}
+
+TEST_F(CleanUpReplacementsTest, FormatCorrectLineWhenHeadersAreInserted) {
+  std::string Code = "\n"
+                     "int    a;\n"
+                     "int    a;\n"
+                     "int    a;";
+
+  std::string Expected = "#include \"x.h\"\n"
+                         "#include \"y.h\"\n"
+                         "#include \"clang/x/x.h\"\n"
+                         "#include <list>\n"
+                         "#include <vector>\n"
+                         "\n"
+                         "int    a;\n"
+                         "int b;\n"
+                         "int    a;";
+  FileID ID = Context.createInMemoryFile("fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(tooling::Replacement(Context.Sources,
+                                       Context.getLocation(ID, 3, 8), 1, "b"));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include <vector>"));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include <list>"));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"clang/x/x.h\""));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"y.h\""));
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include \"x.h\""));
+  format::FormatStyle Style = format::getLLVMStyle();
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(
+                          Code, formatReplacements(Code, NewReplaces, Style)));
+}
+
+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;";
+  Context.createInMemoryFile("fix.cpp", Code);
+  tooling::Replacements Replaces;
+  Replaces.insert(
+      tooling::Replacement("fix.cpp", UINT_MAX, 0, "#include <vector>"));
+  format::FormatStyle Style = format::getLLVMStyle();
+  auto NewReplaces = cleanupAroundReplacements(Code, Replaces, Style);
+  EXPECT_EQ(Expected, applyAllReplacements(Code, NewReplaces));
+}
+
 } // end namespace
 } // end namespace format
 } // end namespace clang