[clang][include cleaner] Fix MainHeader insertion issue

clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp

@@ -26,7 +26,6 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
-#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/Testing/Annotations/Annotations.h"
@@ -398,7 +397,7 @@ TEST_F(AnalyzeTest, SpellingIncludesWithSymlinks) {
 }
 
 // Make sure that the references to implicit operator new/delete are reported as
-// ambigious.
+// ambiguous.
 TEST_F(AnalyzeTest, ImplicitOperatorNewDeleteNotMissing) {
   ExtraFS = llvm::makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>();
   ExtraFS->addFile("header.h",
@@ -706,5 +705,83 @@ TEST_F(WalkUsedTest, MacroConcat) {
       AllOf(Contains(Pair(Code.point("bar"), UnorderedElementsAre(Header))),
             Contains(Pair(Code.point("xyz"), UnorderedElementsAre(Header)))));
 }
+
+TEST(FixIncludes, MainHeaderGrouping) {
+  AnalysisResults Results;
+  Results.Missing.push_back({"\"b.h\"", Header("\"b.h\"")});
+  Results.Missing.push_back({"\"foo.h\"", Header("\"foo.h\"")});
+  Results.Missing.push_back({"<vector>", Header("<vector>")});
+  Results.Missing.push_back({"\"a.h\"", Header("\"a.h\"")});
+
+  format::FormatStyle Style = format::getLLVMStyle();
+  Style.Language = format::FormatStyle::LK_Cpp;
+
+  std::string Code = R"cpp(
+void test();
+)cpp";
+
+  std::string Fixed = fixIncludes(Results, "foo.cc", Code, Style);
+  EXPECT_EQ(Fixed, "\n#include \"foo.h\"\n#include \"a.h\"\n#include "
+                   "\"b.h\"\n#include <vector>\nvoid test();\n");
+}
+
+TEST(FixIncludes, MultipleInsertionsAndDeletions) {
+  AnalysisResults Results;
+  Include UnusedInc;
+  UnusedInc.Spelled = "unused.h";
+  UnusedInc.Line = 1;
+  Results.Unused.push_back(&UnusedInc);
+
+  Results.Missing.push_back({"\"a.h\"", Header("\"a.h\"")});
+  Results.Missing.push_back({"<foo>", Header("<foo>")});
+
+  format::FormatStyle Style = format::getLLVMStyle();
+  Style.Language = format::FormatStyle::LK_Cpp;
+
+  std::string Code = R"cpp(#include "unused.h"
+
+void test();
+)cpp";
+
+  std::string Fixed = fixIncludes(Results, "test.cc", Code, Style);
+  EXPECT_EQ(Fixed, R"cpp(#include "a.h"
+#include <foo>
+
+void test();
+)cpp");
+}
+
+TEST(FixIncludes, MultipleInsertionsSameOffset) {
+  AnalysisResults Results;
+  Results.Missing.emplace_back("\"a.h\"", Header(""));
+  Results.Missing.emplace_back("\"b.h\"", Header(""));
+
+  // Empty code guarantees HeaderIncludes chooses offset 0 for both.
+  llvm::StringRef Code = "";
+
+  // Should concatenate them without conflict errors in Replacements::add
+  EXPECT_EQ(fixIncludes(Results, "d.cc", Code, format::getLLVMStyle()),
+            "#include \"a.h\"\n#include \"b.h\"\n");
+}
+
+TEST(FixIncludes, MissingIncludesSortingAndGrouping) {
+  AnalysisResults Results;
+  Results.Missing.push_back({"\"b.h\"", Header("\"b.h\"")});
+  Results.Missing.push_back({"\"a.h\"", Header("\"a.h\"")});
+  Results.Missing.push_back({"<foo>", Header("<foo>")});
+
+  format::FormatStyle Style = format::getLLVMStyle();
+  Style.Language = format::FormatStyle::LK_Cpp;
+
+  std::string Code = R"cpp(
+void bar();
+)cpp";
+
+  std::string Fixed = fixIncludes(Results, "test.cc", Code, Style);
+  EXPECT_EQ(
+      Fixed,
+      "\n#include \"a.h\"\n#include \"b.h\"\n#include <foo>\nvoid bar();\n");
+}
+
 } // namespace
 } // namespace clang::include_cleaner

clang/include/clang/Tooling/Inclusions/HeaderIncludes.h

@@ -9,13 +9,16 @@
 #ifndef LLVM_CLANG_TOOLING_INCLUSIONS_HEADERINCLUDES_H
 #define LLVM_CLANG_TOOLING_INCLUSIONS_HEADERINCLUDES_H
 
-#include "clang/Basic/SourceManager.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "clang/Tooling/Inclusions/IncludeStyle.h"
-#include "llvm/Support/Path.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/Support/Regex.h"
 #include <list>
 #include <optional>
+#include <set>
+#include <string>
 #include <unordered_map>
 
 namespace clang {
@@ -64,7 +67,7 @@ class HeaderIncludes {
   /// default. These code sections include:
   ///   - raw string literals (containing #include).
   ///   - #if blocks.
-  ///   - Special #include's among declarations (e.g. functions).
+  ///   - Special #includes 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
@@ -77,6 +80,38 @@ class HeaderIncludes {
                                              bool IsAngled,
                                              IncludeDirective Directive) const;
 
+  /// Represents a single header directive to be inserted in a batch operation.
+  ///
+  /// Usage:
+  ///   - HeaderToInsert("<vector>") -> inserts #include <vector>
+  ///     (auto-detects angled)
+  ///   - HeaderToInsert("\"foo.h\"") -> inserts #include "foo.h"
+  ///     (auto-detects quoted)
+  ///   - HeaderToInsert("<foo>", IncludeDirective::Import) -> inserts #import
+  ///     <foo>
+  ///   - HeaderToInsert("foo.h", IncludeDirective::Include, /*IsAngled=*/false)
+  ///     -> explicit IsAngled
+  struct HeaderToInsert {
+    enum class QuoteStyle { AUTO, ANGLED, QUOTED };
+
+    // The header name, with any surrounding quotes or brackets removed.
+    std::string Header;
+    // Whether to insert #include or #import.
+    IncludeDirective Directive;
+    // Whether to use <> or "" for the header. This can be set explicitly with
+    // QuoteStyle::ANGLED or QuoteStyle::QUOTED, or auto-detected based on
+    // `RawOrSpelledHeader` with QuoteStyle::AUTO.
+    bool IsAngled;
+
+    HeaderToInsert(llvm::StringRef RawOrSpelledHeader,
+                   IncludeDirective Directive = IncludeDirective::Include,
+                   QuoteStyle QuoteStyle = QuoteStyle::AUTO);
+  };
+
+  /// Inserts a batch of headers into the code, sorting and grouping them
+  /// according to IncludeStyle and returning the replacements.
+  tooling::Replacements insert(llvm::ArrayRef<HeaderToInsert> Headers) const;
+
   /// Removes all existing #includes and #imports 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 and
@@ -88,7 +123,7 @@ class HeaderIncludes {
 
 private:
   struct Include {
-    Include(StringRef Name, tooling::Range R, IncludeDirective D)
+    Include(llvm::StringRef Name, tooling::Range R, IncludeDirective D)
         : Name(Name), R(R), Directive(D) {}
 
     // An include header quoted with either <> or "".

clang/lib/Format/Format.cpp

@@ -4211,6 +4211,7 @@ fixCppIncludeInsertions(StringRef Code, const tooling::Replacements &Replaces,
   }
 
   SmallVector<StringRef, 4> Matches;
+  SmallVector<tooling::HeaderIncludes::HeaderToInsert, 4> HeadersToInsert;
   for (const auto &R : HeaderInsertions) {
     auto IncludeDirective = R.getReplacementText();
     bool Matched =
@@ -4219,19 +4220,17 @@ fixCppIncludeInsertions(StringRef Code, const tooling::Replacements &Replaces,
                       "'#include ...'");
     (void)Matched;
     auto IncludeName = Matches[2];
-    auto Replace =
-        Includes.insert(IncludeName.trim("\"<>"), IncludeName.starts_with("<"),
-                        tooling::IncludeDirective::Include);
-    if (Replace) {
-      auto Err = Result.add(*Replace);
-      if (Err) {
-        consumeError(std::move(Err));
-        unsigned NewOffset =
-            Result.getShiftedCodePosition(Replace->getOffset());
-        auto Shifted = tooling::Replacement(FileName, NewOffset, 0,
-                                            Replace->getReplacementText());
-        Result = Result.merge(tooling::Replacements(Shifted));
-      }
+    HeadersToInsert.emplace_back(IncludeName,
+                                 tooling::IncludeDirective::Include);
+  }
+  for (const auto &Replace : Includes.insert(HeadersToInsert)) {
+    auto Err = Result.add(Replace);
+    if (Err) {
+      consumeError(std::move(Err));
+      unsigned NewOffset = Result.getShiftedCodePosition(Replace.getOffset());
+      auto Shifted = tooling::Replacement(FileName, NewOffset, 0,
+                                          Replace.getReplacementText());
+      Result = Result.merge(tooling::Replacements(Shifted));
     }
   }
   return Result;

clang/lib/Tooling/Inclusions/HeaderIncludes.cpp

@@ -14,6 +14,8 @@
 #include "clang/Lex/Token.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "clang/Tooling/Inclusions/IncludeStyle.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
@@ -25,10 +27,12 @@
 #include <cassert>
 #include <climits>
 #include <functional>
+#include <iterator>
 #include <optional>
 #include <string>
 #include <type_traits>
 #include <utility>
+#include <vector>
 
 namespace clang {
 namespace tooling {
@@ -512,11 +516,83 @@ HeaderIncludes::insert(llvm::StringRef Header, bool IsAngled,
   return tooling::Replacement(FileName, InsertOffset, 0, NewInclude);
 }
 
-tooling::Replacements HeaderIncludes::remove(llvm::StringRef IncludeName,
+HeaderIncludes::HeaderToInsert::HeaderToInsert(
+    llvm::StringRef RawOrSpelledHeader, IncludeDirective Directive,
+    QuoteStyle QuoteStyle)
+    : Directive(Directive) {
+  if (RawOrSpelledHeader.starts_with("<")) {
+    Header = RawOrSpelledHeader.trim("<>").str();
+    this->IsAngled = QuoteStyle != QuoteStyle::QUOTED;
+  } else if (RawOrSpelledHeader.starts_with("\"")) {
+    Header = RawOrSpelledHeader.trim("\"").str();
+    this->IsAngled = QuoteStyle == QuoteStyle::ANGLED;
+  }
+}
+
+tooling::Replacements
+HeaderIncludes::insert(llvm::ArrayRef<HeaderToInsert> Headers) const {
+  tooling::Replacements Result;
+  if (Headers.empty())
+    return Result;
+
+  std::vector<HeaderToInsert> SortedHeaders = Headers.vec();
+  llvm::stable_sort(SortedHeaders, [&](const HeaderToInsert &L,
+                                       const HeaderToInsert &R) {
+    std::string QuotedL =
+        std::string(llvm::formatv(L.IsAngled ? "<{0}>" : "\"{0}\"", L.Header));
+    std::string QuotedR =
+        std::string(llvm::formatv(R.IsAngled ? "<{0}>" : "\"{0}\"", R.Header));
+    int PriorityL = Categories.getIncludePriority(
+        QuotedL, /*CheckMainHeader=*/!MainIncludeFound);
+    int PriorityR = Categories.getIncludePriority(
+        QuotedR, /*CheckMainHeader=*/!MainIncludeFound);
+    if (PriorityL != PriorityR)
+      return PriorityL < PriorityR;
+    if (L.Header != R.Header)
+      return L.Header < R.Header;
+    if (L.IsAngled != R.IsAngled)
+      return L.IsAngled < R.IsAngled;
+    return L.Directive > R.Directive;
+  });
+  SortedHeaders.erase(
+      std::unique(SortedHeaders.begin(), SortedHeaders.end(),
+                  [](const HeaderToInsert &L, const HeaderToInsert &R) {
+                    return L.Header == R.Header && L.IsAngled == R.IsAngled;
+                  }),
+      SortedHeaders.end());
+
+  struct InsertionInfo {
+    std::string Text;
+    unsigned Length = 0;
+  };
+  llvm::DenseMap<unsigned, InsertionInfo> InsertionsByOffset;
+
+  for (const auto &H : SortedHeaders) {
+    if (auto Insertion = insert(H.Header, H.IsAngled, H.Directive)) {
+      auto &Info = InsertionsByOffset[Insertion->getOffset()];
+      Info.Text += Insertion->getReplacementText();
+      if (Insertion->getLength() > 0) {
+        assert(Info.Length == 0 && "Multiple replacements at same offset?");
+        Info.Length = Insertion->getLength();
+      }
+    }
+  }
+
+  for (const auto &Entry : InsertionsByOffset) {
+    const auto &Info = Entry.second;
+    const unsigned Offset = Entry.first;
+    cantFail(Result.add(
+        tooling::Replacement(FileName, Offset, Info.Length, Info.Text)));
+  }
+
+  return Result;
+}
+
+tooling::Replacements HeaderIncludes::remove(llvm::StringRef Header,
                                              bool IsAngled) const {
-  assert(IncludeName == trimInclude(IncludeName));
+  assert(Header == trimInclude(Header));
   tooling::Replacements Result;
-  auto Iter = ExistingIncludes.find(IncludeName);
+  auto Iter = ExistingIncludes.find(Header);
   if (Iter == ExistingIncludes.end())
     return Result;
   for (const auto &Inc : Iter->second) {