[clang-format] Refactor #include insertion/deletion functionality into a class.

[ioeric]

Summary

The class will be moved into libToolingCore as followup.

The new behaviors in this patch:

• New #include is inserted in the right position in a #include block to

preserver sorted #includes. This is best effort - only works when the
block is already sorted.

• When inserting multiple #includes to the end of a file which doesn't

end with a "\n" character, a "\n" will be prepended to each #include.
This is a special and rare case that was previously handled. This is now
relaxed to avoid complexity as it's rare in practice.

[ioeric]

There isn't actually as much code changed as it appears to be, but phabricator doens't understand the diff very well... git diff might be an easier way to review the patch.

[ilya-biryukov]

This change LG as an extraction of the helper functionality to be reused in clang, clang-tidy, etc.
However, I feel there are potential improvements both to the underlying code and the new APIs that we could make.

I left some comments, trying to focus on interface of the class and keeping the changes that would be required to address them to be purely cosmetic. As chatted offline, any non-trivial changes to the underlying implementation are out of scope of this patch.

preserver sorted #includes.

A typo in the change description:
s/preserver/preserve

This is best effort - only works when the
block is already sorted.

Could we describe behavior for the unsorted case in a sensible way? E.g. is it added to the end of the include list, added after closely related headers, added to a totally random place (but deterministic) place, etc?
It seems like an important case, and I believe we shouldn't completely ignore it and describe what the "best effort" means in that case.

When inserting multiple #includes to the end of a file which doesn't
end with a "\n" character, a "\n" will be prepended to each #include.
This is a special and rare case that was previously handled. This is now
relaxed to avoid complexity as it's rare in practice.

I don't quite get it, could you please elaborate on what has changed here exactly? Maybe give an example?

lib/Format/Format.cpp

1692–1693

I would personally keep the function non-const and not use mutable fields here.
Even though it's logically const, I would strive towards keeping the things const only if there are actually immutable
One particular problem that could be avoided is accidentally calling the const methods concurrently on different threads.

But up to you if you actually want to make this change.

2208

Maybe s/#include headers/#include directives/?
This is how they called in terminology of preprocessing.

2215

Maybe make an API more explicit about the style of includes that should be used?
Make it literally impossible to misuse the API:

enum class IncludeStyle { Quoted, Angle };
/// \p IncludeName is the include path to be inserted, \p Style specifies whether the included path should use quotes or angle brackets.
Optional<Replacement> insert(StringRef IncludeName, IncludeStyle Style = IncludeStyle::Quoted) const;

2220

Maybe use StringRef here?

2233

Have you considered changing the API slightly to allow iterating over all includes in the file and APIs to remove includes pointed by the iterators?
It would give more flexibility, allowing the clients to implement useful things like:

• Remove all #includes that start with a prefix, e.g. clang/Basic/.....
• Find and remove all duplicate #include directives.
• etc, etc.

The current implementation seems tied to a specific use-case that we have in clang-format, i.e. "preprocess once, remove headers in batches".
It feels wasteful to pay overhead of StringMap<vector<Include>> sometimes, e.g. if the code wants to insert just a single header or do other things, like removing duplicate headers.

I.e. I propose to extract the code that parses the source code and produces a list of #include directives with their respective ranges and filenames into a separate function/class.
This would. arguably, improve readability of the include insertion code.

That might not be a trivial change, though, since it requires rewriting some of the code used here, while the current patch seems to focus on simply extracting useful functionality from clang-format for reuse in clangd, clang-tidy, etc.
Let me know what you think.

2235

We should mention that remove only deletes exact matches and won't try to resolve filenames, patch up path separators, etc.

2236

Maybe also add a param of type Optional<IncludeStyle> and

• only remove includes with the specified style, if specified the value is specified
• remove both styles the value is llvm::None.

2237–2238

What do we use the FileName for?
On a higher level, it seems we shouldn't need it to implement what HeaderIncludes does.

2246

Could you please add a comment explaining what are those priorities and what this map is used for?

2246

Is this field always equal to Code.drop_front(MinInsertOffset)?
If so, maybe compute it when needed and don't store explicitly?

2248

Could you add a comment explaining how min and max offsets are calculated? I assume minoffset goes after the file comment and maxoffset is where non-preprocessor code starts, but it would be nice to have a short comment to know for sure.

[ioeric]

Thanks for reviewing this!

In D46180#1080822, @ilya-biryukov wrote:

This change LG as an extraction of the helper functionality to be reused in clang, clang-tidy, etc.
However, I feel there are potential improvements both to the underlying code and the new APIs that we could make.

I left some comments, trying to focus on interface of the class and keeping the changes that would be required to address them to be purely cosmetic. As chatted offline, any non-trivial changes to the underlying implementation are out of scope of this patch.

preserver sorted #includes.

A typo in the change description:
s/preserver/preserve

This is best effort - only works when the
block is already sorted.

Could we describe behavior for the unsorted case in a sensible way? E.g. is it added to the end of the include list, added after closely related headers, added to a totally random place (but deterministic) place, etc?
It seems like an important case, and I believe we shouldn't completely ignore it and describe what the "best effort" means in that case.

Sure. I added more comments to the API.

When inserting multiple #includes to the end of a file which doesn't
end with a "\n" character, a "\n" will be prepended to each #include.
This is a special and rare case that was previously handled. This is now
relaxed to avoid complexity as it's rare in practice.

I don't quite get it, could you please elaborate on what has changed here exactly? Maybe give an example?

This is really a corner cases that users might not need to know about. But an example is:
Code: "#include <a>" (without \n at the end). After inserting <x>, #include <a>\n#include <x>\n (this is good). However, if you insert another <y>, the code would become "#include <a>\n#include <x>\n\n#include <y>\n" (note the double newline!).

lib/Format/Format.cpp

1692–1693

Making this non-const would prevent us from making HeaderIncludes::insert const, which I do want to keep. Added a mutex to prevent race-condition...

2215

I think the risk is pretty low, so I'm not sure if it'd be worth spending another variable. From experience, the current callers would likely need to either carry another variable for each include or use the following pattern (as the IncludeName is often already quoted):

... = insert(Include.trime("\"<>"), QuoteStyle);

2220

We need to worry about life-cycle of the underlying string if we use StringRef, and I think it's not trivial in this case, so I think std::string would be a better fit here.

2233

I haven't thought think about these use cases.

These use cases don't seem to be very common though. Generally, I try to make the APIs easy to use for the use cases we have instead designing the APIs for all potential uses. An over-generalized API might be awkward to use for major use cases, and designing for the future also has opportunity cost...

That said, I think when the use cases do come up, it wouldn't be hard to add APIs for that (and change implementation if necessary). But I also think the current APIs would still be useful for users who only care about #include insertion/deletion.

2236

Again, I'm not quite sure if it's worth another parameter for this.

2237–2238

FileName is needed to decide whether a header is a main header.

[ioeric]

Addressed review comments.

[ilya-biryukov]

This is really a corner cases that users might not need to know about. But an example is:
Code: "#include <a>" (without \n at the end). After inserting <x>, #include <a>\n#include <x>\n (this is good). However, if you insert another <y>, the code would become "#include <a>\n#include <x>\n\n#include <y>\n" (note the double newline!).

Even though that's rare in practice, it seems like an actual bug.
How hard would be it be to fix it?

[ilya-biryukov]

lib/Format/Format.cpp

1692–1693

Actually keeping things const and mentioning in the class comments that the class is not thread-safe because it uses regexes seems like a better choice.
Most of the clients shouldn't pay for the mutex operations for the sake of the rare multi-threaded clients.
WDYT?

2215

I still don't feel encoding things in the strings is the right approach here.
It raises some questions that we should handle/assert: what if users don't quote the strings at all? what is the leading and trailing quotes don't match? and so on and so forth.
Being explicit about such things seems like a better choice.

From experience, the current callers would likely need to either carry another variable

It's not surprising, since the class is literally extracted from the middle of the include insertion implementation in clang-format.

2220

We used StringReffor Code, Filename, etc before.
Are they different? Is the lifetime of Include more complicated? When would Include point outside Code?

2233

Ok, I'm happy with leaving the current API as is, even though I do feel some refactoring would help both the readability and allow us to reuse the code for more use-cases.
But does seem out of scope of this patch, so we could wait for an actual use-case to pop up.

2236

The parameter in this case seems to be even more useful, since we have two different behaviors for quoted vs unquoted strings.
So it's really easy to misuse the API and make a mistake, thinking it does the right thing for you.

2237–2238

Can we pass bool IsMainHeader instead? Can we not store it as a field? It seems to be only needed in the constructor, right?

[ioeric]

• addressed more review comments.

lib/Format/Format.cpp

1692–1693

Honestly, I'm a bit torn on this. Either approach is not perfect but either seems acceptable. Rolled back to const+mutable and added the comment.

2215

Ok, I am convinced ;)

2236

After offline discussion, I was also convinced here :)

2237–2238

This can be only calculated at insert time when IncludeName is provided - we can only tell if a new include is main header when both FileName and name of the new header are provided.

[ioeric]

Friendly ping.

[ilya-biryukov]

I somehow missed the review email, sorry for the delayed response.

Just one nit and one question from me on changed behavior in the tests (quoted vs angled #include).
Otherwise LG, just wanted to make sure the change in behavior is intentional.

lib/Format/Format.cpp

2131

NIT: Maybe we could add some asserts to this function that the passed include name is actually properly quoted.
E.g. starts with '<' or '"', ends with the corresponding char, etc.

So long as this is part of Format.cpp, it's not terribly important, since it's only called on paths that come from regex matches.
If we make it a public later, having asserts would be really useful.
That said, we might address it later when/if we actually make this function public.

unittests/Format/CleanupTest.cpp

862

Are we sure we want this behavior change?
The API seems to allow us keeping the old one.

I know there's a FIXME there, but I just wanted to understand better what's the right thing to do here and why.

[ioeric]

• Addressed review comment.

[ilya-biryukov]

LGTM

lib/Format/Format.cpp

 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/YAMLTraits.h"
 #include <algorithm>
 #include <memory>
+#include <mutex>
 #include <string>
+#include <unordered_map>
 
 #define DEBUG_TYPE "format-formatter"
 
 using clang::format::FormatStyle;
 
 LLVM_YAML_IS_SEQUENCE_VECTOR(clang::format::FormatStyle::IncludeCategory)
 LLVM_YAML_IS_SEQUENCE_VECTOR(clang::format::FormatStyle::RawStringFormat)
 
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) {
+  // 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.IncludeCategories[i].Priority;
         break;
       }
     if (CheckMainHeader && IsMainFile && Ret > 0 && isMainHeader(IncludeName))
       Ret = 0;
bool isMainHeader(StringRef IncludeName) const {
     }
     return false;
   }
 
   const FormatStyle &Style;
   bool IsMainFile;
   StringRef FileName;
   StringRef FileStem;
-  SmallVector<llvm::Regex, 4> CategoryRegexs;
+  // Regex is not thread-safe.
+  mutable SmallVector<llvm::Regex, 4> CategoryRegexs;
 };
 
 const char IncludeRegexPattern[] =
     R"(^[\t\ ]*#[\t\ ]*(import|include)[^"<]*(["<][^">]*[">]))";
 
 } // anonymous namespace
 
 tooling::Replacements sortCppIncludes(const FormatStyle &Style, StringRef Code,
return getOffsetAfterTokenSequence(
         skipComments(Lex, Tok);
         unsigned MaxOffset = SM.getFileOffset(Tok.getLocation());
         while (checkAndConsumeInclusiveDirective(Lex, Tok))
           MaxOffset = SM.getFileOffset(Tok.getLocation());
         return MaxOffset;
       });
 }
 
-bool isDeletedHeader(llvm::StringRef HeaderName,
-                     const std::set<llvm::StringRef> &HeadersToDelete) {
-  return HeadersToDelete.count(HeaderName) ||
-         HeadersToDelete.count(HeaderName.trim("\"<>"));
+/// 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.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 not be inserted.
+  if (ExistingIncludes.find(IncludeName) != ExistingIncludes.end())
+    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,
                         const FormatStyle &Style) {
   if (!Style.isCpp())
     return Replaces;
 
   tooling::Replacements HeaderInsertions;
   std::set<llvm::StringRef> HeadersToDelete;
   tooling::Replacements Result;
   for (const auto &R : Replaces) {
     if (isHeaderInsertion(R)) {
       // Replacements from \p Replaces must be conflict-free already, so we can
       // simply consume the error.
       llvm::consumeError(HeaderInsertions.add(R));
     } else if (isHeaderDeletion(R)) {
       HeadersToDelete.insert(R.getReplacementText());
     } else if (R.getOffset() == UINT_MAX) {
       llvm::errs() << "Insertions other than header #include insertion are "
                       "not supported! "
                    << R.getReplacementText() << "\n";
     } else {
       llvm::consumeError(Result.add(R));
     }
   }
   if (HeaderInsertions.empty() && HeadersToDelete.empty())
     return Replaces;
 
-  llvm::Regex IncludeRegex(IncludeRegexPattern);
-  llvm::Regex DefineRegex(R"(^[\t\ ]*#[\t\ ]*define[\t\ ]*[^\\]*$)");
-  SmallVector<StringRef, 4> Matches;
 
   StringRef FileName = Replaces.begin()->getFilePath();
-  IncludeCategoryManager Categories(Style, FileName);
+  HeaderIncludes Includes(FileName, Code, Style);
 
-  // Record the offset of the end of the last include in each category.
-  std::map<int, int> CategoryEndOffsets;
-  // All possible priorities.
-  // Add 0 for main header and INT_MAX for headers that are not in any category.
-  std::set<int> Priorities = {0, INT_MAX};
-  for (const auto &Category : Style.IncludeCategories)
-    Priorities.insert(Category.Priority);
-  int FirstIncludeOffset = -1;
-  // All new headers should be inserted after this offset.
-  unsigned MinInsertOffset =
-      getOffsetAfterHeaderGuardsAndComments(FileName, Code, Style);
-  StringRef TrimmedCode = Code.drop_front(MinInsertOffset);
-  // Max insertion offset in the original code.
-  unsigned MaxInsertOffset =
-      MinInsertOffset +
-      getMaxHeaderInsertionOffset(FileName, TrimmedCode, Style);
-  SmallVector<StringRef, 32> Lines;
-  TrimmedCode.split(Lines, '\n');
-  unsigned Offset = MinInsertOffset;
-  unsigned NextLineOffset;
-  std::set<StringRef> ExistingIncludes;
-  for (auto Line : Lines) {
-    NextLineOffset = std::min(Code.size(), Offset + Line.size() + 1);
-    if (IncludeRegex.match(Line, &Matches)) {
-      // The header name with quotes or angle brackets.
-      StringRef IncludeName = Matches[2];
-      ExistingIncludes.insert(IncludeName);
-      // Only record the offset of current #include if we can insert after it.
-      if (Offset <= MaxInsertOffset) {
-        int Category = Categories.getIncludePriority(
-            IncludeName, /*CheckMainHeader=*/FirstIncludeOffset < 0);
-        CategoryEndOffsets[Category] = NextLineOffset;
-        if (FirstIncludeOffset < 0)
-          FirstIncludeOffset = Offset;
-      }
-      if (isDeletedHeader(IncludeName, HeadersToDelete)) {
-        // If this is the last line without trailing newline, we need to make
-        // sure we don't delete across the file boundary.
-        unsigned Length = std::min(Line.size() + 1, Code.size() - Offset);
-        llvm::Error Err =
-            Result.add(tooling::Replacement(FileName, Offset, Length, ""));
-        if (Err) {
-          // Ignore the deletion on conflict.
-          llvm::errs() << "Failed to add header deletion replacement for "
-                       << IncludeName << ": " << llvm::toString(std::move(Err))
-                       << "\n";
-        }
-      }
-    }
-    Offset = NextLineOffset;
-  }
+  for (const auto &Header : HeadersToDelete) {
+    tooling::Replacements Replaces =
+        Includes.remove(trimInclude(Header), Header.startswith("<"));
+    for (const auto &R : Replaces) {
+      auto Err = Result.add(R);
+                     << Header << ": " << llvm::toString(std::move(Err))
 
-  // 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)];
-
-  bool NeedNewLineAtEnd = !Code.empty() && Code.back() != '\n';
+  llvm::Regex IncludeRegex = llvm::Regex(IncludeRegexPattern);
+  llvm::SmallVector<StringRef, 4> Matches;
   for (const auto &R : HeaderInsertions) {
     auto IncludeDirective = R.getReplacementText();
     bool Matched = IncludeRegex.match(IncludeDirective, &Matches);
     assert(Matched && "Header insertion replacement must have replacement text "
                       "'#include ...'");
     (void)Matched;
     auto IncludeName = Matches[2];
-    if (ExistingIncludes.find(IncludeName) != ExistingIncludes.end()) {
-      DEBUG(llvm::dbgs() << "Skip adding existing include : " << IncludeName
-                         << "\n");
-      continue;
-    }
-    int Category =
-        Categories.getIncludePriority(IncludeName, /*CheckMainHeader=*/true);
-    Offset = CategoryEndOffsets[Category];
-    std::string NewInclude = !IncludeDirective.endswith("\n")
-                                 ? (IncludeDirective + "\n").str()
-                                 : IncludeDirective.str();
-    // When inserting headers at end of the code, also append '\n' to the code
-    // if it does not end with '\n'.
-    if (NeedNewLineAtEnd && Offset == Code.size()) {
-      NewInclude = "\n" + NewInclude;
-      NeedNewLineAtEnd = false;
-    }
-    auto NewReplace = tooling::Replacement(FileName, Offset, 0, NewInclude);
-    auto Err = Result.add(NewReplace);
+    auto Replace =
+        Includes.insert(trimInclude(IncludeName), IncludeName.startswith("<"));
+    if (Replace) {
+      auto Err = Result.add(*Replace);
       if (Err) {
         llvm::consumeError(std::move(Err));
-      unsigned NewOffset = Result.getShiftedCodePosition(Offset);
-      NewReplace = tooling::Replacement(FileName, NewOffset, 0, NewInclude);
-      Result = Result.merge(tooling::Replacements(NewReplace));
+        unsigned NewOffset = Result.getShiftedCodePosition(Replace->getOffset());
+        auto Shifted = tooling::Replacement(FileName, NewOffset, 0,
+                                            Replace->getReplacementText());
+        Result = Result.merge(tooling::Replacements(Shifted));
+      }
     }
   }
   return Result;
 }
 
 } // anonymous namespace
 
 llvm::Expected<tooling::Replacements>

unittests/Format/CleanupTest.cpp

std::string Expected = "#include \"x/fix.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"
                      "#include \"b.h\"\n"
+                     "#include \"z.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 \"z.h\"\n"
                          "#include \"clang/Format/Format.h\"\n"
-                         "#include <memory>\n"
-                         "#include <list>\n";
+                         "#include <list>\n"
+                         "#include <memory>\n";
   tooling::Replacements Replaces =
       toReplacements({createInsertion("#include <list>"),
                       createInsertion("#include \"new/new.h\"")});
   EXPECT_EQ(Expected, apply(Code, Replaces));
 }
 
 TEST_F(CleanUpReplacementsTest, InsertNewSystemIncludeGoogleStyle) {
   std::string Code = "#include \"x/fix.h\"\n"
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"
+                         "#include <vector>\n"
                          "\n"
+                         "#include \"x/x.h\"\n"
                          "#include \"y/a.h\"\n"
-                         "#include \"z/b.h\"\n"
-                         "#include \"x/x.h\"\n";
+                         "#include \"z/b.h\"\n";
   tooling::Replacements Replaces =
       toReplacements({createInsertion("#include <list>"),
                       createInsertion("#include \"x/x.h\"")});
   Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp);
   EXPECT_EQ(Expected, apply(Code, Replaces));
 }
 
 TEST_F(CleanUpReplacementsTest, InsertMultipleNewHeadersAndSortLLVM) {
TEST_F(CleanUpReplacementsTest, NoNewLineAtTheEndOfCode) {
   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#include <vector>\n";
+      "#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) {
+TEST_F(CleanUpReplacementsTest, NoAddSameIncludesWithDifferentForms) {
   std::string Code = "#include \"a.h\"\n"
                      "#include <vector>\n";
-  // FIXME: this might not be the best behavior.
   std::string Expected = "#include \"a.h\"\n"
-                         "#include \"vector\"\n"
-                         "#include <vector>\n"
-                         "#include <a.h>\n";
+                         "#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"
                      "#include \"xyz\"\n"
                      "int x;\n";
   std::string Expected = "#include \"xyz\"\n"
                          "int x;\n";
   tooling::Replacements Replaces =
       toReplacements({createDeletion("abc.h"), createDeletion("xyz.h")});
   EXPECT_EQ(Expected, apply(Code, Replaces));
 }
 
 TEST_F(CleanUpReplacementsTest, DeleteAllCode) {
   std::string Code = "#include \"xyz.h\"\n"
                      "#include <xyz.h>";
   std::string Expected = "";
-  tooling::Replacements Replaces = toReplacements({createDeletion("xyz.h")});
+  tooling::Replacements Replaces =
+      toReplacements({createDeletion("\"xyz.h\""), createDeletion("<xyz.h>")});
   EXPECT_EQ(Expected, apply(Code, Replaces));
 }
 
-TEST_F(CleanUpReplacementsTest, DeleteAllIncludesWithSameNameIfNoType) {
+TEST_F(CleanUpReplacementsTest, DeleteOnlyIncludesWithSameQuote) {
   std::string Code = "#include \"xyz.h\"\n"
                      "#include \"xyz\"\n"
                      "#include <xyz.h>\n";
-  std::string Expected = "#include \"xyz\"\n";
-  tooling::Replacements Replaces = toReplacements({createDeletion("xyz.h")});
+  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"