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

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.

Reviewers: klimek, djasper

Subscribers: klimek, cfe-commits, bkramer

Differential Revision: http://reviews.llvm.org/D20734

llvm-svn: 271276

clang/include/clang/Format/Format.h

@@ -773,6 +773,8 @@ tooling::Replacements formatReplacements(StringRef Code,
 
 /// \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);

clang/lib/Format/Format.cpp

@@ -1262,6 +1262,65 @@ static void sortCppIncludes(const FormatStyle &Style,
                                        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,
@@ -1269,8 +1328,7 @@ tooling::Replacements sortCppIncludes(const FormatStyle &Style, StringRef Code,
                                       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;
 
@@ -1281,19 +1339,9 @@ 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.
-  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 (;;) {
@@ -1310,26 +1358,11 @@ tooling::Replacements sortCppIncludes(const FormatStyle &Style, StringRef Code,
     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)
+          MainIncludeFound = true;
         IncludesInBlock.push_back({IncludeName, Line, Prev, Category});
       } else if (!IncludesInBlock.empty()) {
         sortCppIncludes(Style, IncludesInBlock, Ranges, FileName, Replaces,
@@ -1402,6 +1435,113 @@ tooling::Replacements formatReplacements(StringRef Code,
   return processReplacements(Reformat, Code, SortedReplaces, Style);
 }
 
+namespace {
+
+inline bool isHeaderInsertion(const tooling::Replacement &Replace) {
+  return Replace.getOffset() == UINT_MAX &&
+         llvm::Regex(IncludeRegexPattern).match(Replace.getReplacementText());
+}
+
+// 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) {
+  if (Style.Language != FormatStyle::LanguageKind::LK_Cpp)
+    return Replaces;
+
+  tooling::Replacements HeaderInsertions;
+  for (const auto &R : Replaces) {
+    if (isHeaderInsertion(R))
+      HeaderInsertions.insert(R);
+    else if (R.getOffset() == UINT_MAX)
+      llvm::errs() << "Insertions other than header #include insertion are "
+                      "not supported! "
+                   << R.getReplacementText() << "\n";
+  }
+  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) {
@@ -1412,7 +1552,10 @@ cleanupAroundReplacements(StringRef Code, const tooling::Replacements &Replaces,
                     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 =
+      fixCppIncludeInsertions(Code, Replaces, Style);
+  return processReplacements(Cleanup, Code, NewReplaces, Style);
 }
 
 tooling::Replacements reformat(const FormatStyle &Style, SourceManager &SM,

clang/unittests/Format/CleanupTest.cpp

@@ -243,13 +243,39 @@ TEST_F(CleanupTest, CtorInitializerInNamespace) {
 
 class CleanUpReplacementsTest : public ::testing::Test {
 protected:
-  tooling::Replacement createReplacement(SourceLocation Start, unsigned Length,
-                                         llvm::StringRef ReplacementText) {
-    return tooling::Replacement(Context.Sources, Start, Length,
-                                ReplacementText);
+  tooling::Replacement createReplacement(unsigned Offset, unsigned Length,
+                                         StringRef Text) {
+    return tooling::Replacement(FileName, Offset, Length, Text);
   }
 
-  RewriterTestContext Context;
+  tooling::Replacement createInsertion(StringRef HeaderName) {
+    return createReplacement(UINT_MAX, 0, HeaderName);
+  }
+
+  inline std::string apply(StringRef Code,
+                           const tooling::Replacements Replaces) {
+    return applyAllReplacements(
+        Code, cleanupAroundReplacements(Code, Replaces, Style));
+  }
+
+  inline std::string formatAndApply(StringRef Code,
+                                    const tooling::Replacements Replaces) {
+    return applyAllReplacements(
+        Code,
+        formatReplacements(
+            Code, cleanupAroundReplacements(Code, Replaces, Style), Style));
+  }
+
+  int getOffset(StringRef Code, int Line, int Column) {
+    RewriterTestContext Context;
+    FileID ID = Context.createInMemoryFile(FileName, Code);
+    auto DecomposedLocation =
+        Context.Sources.getDecomposedLoc(Context.getLocation(ID, Line, Column));
+    return DecomposedLocation.second;
+  }
+
+  const std::string FileName = "fix.cpp";
+  FormatStyle Style = getLLVMStyle();
 };
 
 TEST_F(CleanUpReplacementsTest, FixOnlyAffectedCodeAfterReplacements) {
@@ -268,17 +294,247 @@ TEST_F(CleanUpReplacementsTest, FixOnlyAffectedCodeAfterReplacements) {
                          "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, ""));
+  tooling::Replacements Replaces = {
+      createReplacement(getOffset(Code, 3, 3), 6, ""),
+      createReplacement(getOffset(Code, 9, 34), 6, "")};
 
-  format::FormatStyle Style = format::getLLVMStyle();
-  auto FinalReplaces = formatReplacements(
-      Code, cleanupAroundReplacements(Code, Replaces, Style), Style);
-  EXPECT_EQ(Expected, applyAllReplacements(Code, FinalReplaces));
+  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 = {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 = {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 = {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 = {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 = {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 = {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 = {createInsertion("#include \"d.h\""),
+                                    createInsertion("#include \"llvm/x/y.h\"")};
+  EXPECT_EQ(Expected, apply(Code, Replaces));
+}
+
+TEST_F(CleanUpReplacementsTest, InsertMultipleIncludesLLVMStyle) {
+  std::string Code = "#include \"x/fix.h\"\n"
+                     "#include \"a.h\"\n"
+                     "#include \"b.h\"\n"
+                     "#include \"clang/Format/Format.h\"\n"
+                     "#include <memory>\n";
+  std::string Expected = "#include \"x/fix.h\"\n"
+                         "#include \"a.h\"\n"
+                         "#include \"b.h\"\n"
+                         "#include \"new/new.h\"\n"
+                         "#include \"clang/Format/Format.h\"\n"
+                         "#include <memory>\n"
+                         "#include <list>\n";
+  tooling::Replacements Replaces = {createInsertion("#include <list>"),
+                                    createInsertion("#include \"new/new.h\"")};
+  EXPECT_EQ(Expected, apply(Code, Replaces));
+}
+
+TEST_F(CleanUpReplacementsTest, InsertNewSystemIncludeGoogleStyle) {
+  std::string Code = "#include \"x/fix.h\"\n"
+                     "\n"
+                     "#include \"y/a.h\"\n"
+                     "#include \"z/b.h\"\n";
+  // FIXME: inserting after the empty line following the main header might be
+  // prefered.
+  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 = {createInsertion("#include <vector>")};
+  Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
+  EXPECT_EQ(Expected, apply(Code, Replaces));
+}
+
+TEST_F(CleanUpReplacementsTest, InsertMultipleIncludesGoogleStyle) {
+  std::string Code = "#include \"x/fix.h\"\n"
+                     "\n"
+                     "#include <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";
+  tooling::Replacements Replaces = {createInsertion("#include <list>"),
+                                    createInsertion("#include \"x/x.h\"")};
+  Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
+  EXPECT_EQ(Expected, apply(Code, Replaces));
+}
+
+TEST_F(CleanUpReplacementsTest, InsertMultipleNewHeadersAndSortLLVM) {
+  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;";
+  tooling::Replacements Replaces = {createInsertion("#include \"a.h\""),
+                                    createInsertion("#include \"c.h\""),
+                                    createInsertion("#include \"b.h\""),
+                                    createInsertion("#include <vector>"),
+                                    createInsertion("#include <list>"),
+                                    createInsertion("#include \"fix.h\"")};
+  EXPECT_EQ(Expected, formatAndApply(Code, Replaces));
+}
+
+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;";
+  tooling::Replacements Replaces = {createInsertion("#include \"a.h\""),
+                                    createInsertion("#include \"c.h\""),
+                                    createInsertion("#include \"b.h\""),
+                                    createInsertion("#include <vector>"),
+                                    createInsertion("#include <list>"),
+                                    createInsertion("#include \"fix.h\"")};
+  Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
+  EXPECT_EQ(Expected, formatAndApply(Code, Replaces));
+}
+
+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;";
+  tooling::Replacements Replaces = {
+      createReplacement(getOffset(Code, 3, 8), 1, "b"),
+      createInsertion("#include <vector>"),
+      createInsertion("#include <list>"),
+      createInsertion("#include \"clang/x/x.h\""),
+      createInsertion("#include \"y.h\""),
+      createInsertion("#include \"x.h\"")};
+  EXPECT_EQ(Expected, formatAndApply(Code, Replaces));
+}
+
+TEST_F(CleanUpReplacementsTest, NotConfusedByDefine) {
+  std::string Code = "void f() {}\n"
+                     "#define A \\\n"
+                     "  int i;";
+  std::string Expected = "#include <vector>\n"
+                         "void f() {}\n"
+                         "#define A \\\n"
+                         "  int i;";
+  tooling::Replacements Replaces = {createInsertion("#include <vector>")};
+  EXPECT_EQ(Expected, formatAndApply(Code, Replaces));
 }
 
 } // end namespace