[Format][Tooling] Fix HeaderIncludes::insert not respect the main-file header.

[kadircet]

clang/lib/Tooling/Inclusions/HeaderIncludes.cpp

339–341

instead of doing this unconditionally, can we do this until we find first include that looks like the main include? similarly we should also only look for main includes during ::insert if we didn't encounter one already during initial scan.

this ensures we give users a way out (in certain cases) when there are multiple alternatives, by respecting the order they have. it's also aligned with the behaviour of sortCppIncludes, which is used when formatting files elsewhere.

clang/unittests/Tooling/HeaderIncludesTest.cpp

146

can you also have test cases with competing includes, eg:

foo.cc, with contents:

#include "zoo/foo.h"
// insert "xoo/foo.h"

-> xoo/foo.h doesn't become main include

[hokein]

address comments

[kadircet]

clang/lib/Tooling/Inclusions/HeaderIncludes.cpp

369–370

sorry if I was unclear in the previous comment. but i was trying to say we should only consider includes we've encountered during the initial scan.

getting different results depending on the order of insertions being performed is likely gonna cause confusion elsewhere (e.g. depending on the first diagnostic you've in clangd, you'll get different order of includes). hence, let's keep this function stateless (which also gets rid of the concerns around mutable)

[hokein]

Make the MainFileFound stateless.

[hokein]

clang/lib/Tooling/Inclusions/HeaderIncludes.cpp

369–370

sorry for misunderstanding your comment. Make it stateless looks like a good idea.

clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp

void foo() {
     $DEF[[DEF]](a);
 
     $BAR[[BAR]](b);
 
     BAZ($Foo[[Foo]]);
 })cpp");
 
   TestTU TU;
-  TU.Filename = "foo.cpp";
+  TU.Filename = "main.cpp";
   TU.AdditionalFiles["a.h"] = guard("#include \"b.h\"");
   TU.AdditionalFiles["b.h"] = guard("void b();");
 
   TU.AdditionalFiles["dir/c.h"] = guard("#include \"d.h\"");
   TU.AdditionalFiles["dir/d.h"] =
       guard("namespace ns { struct Bar { void d(); }; }");
 
   TU.AdditionalFiles["system/e.h"] = guard("#include <f.h>");

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

 #include "aa.h"
 #include "ab.h"
 #include <e.h>
 )cpp");
 
   Results = {};
   Results.Missing.push_back("\"d.h\"");
   Code = R"cpp(#include "a.h")cpp";
-  // FIXME: this isn't correct, the main-file header d.h should be added before
-  // a.h.
   EXPECT_EQ(fixIncludes(Results, "d.cc", Code, format::getLLVMStyle()),
-R"cpp(#include "a.h"
-#include "d.h"
-)cpp");
+R"cpp(#include "d.h"
+#include "a.h")cpp");
 }
 
 MATCHER_P3(expandedAt, FileID, Offset, SM, "") {
   auto [ExpanedFileID, ExpandedOffset] = SM->getDecomposedExpansionLoc(arg);
   return ExpanedFileID == FileID && ExpandedOffset == Offset;
 }
 MATCHER_P3(spelledAt, FileID, Offset, SM, "") {
   auto [SpelledFileID, SpelledOffset] = SM->getDecomposedSpellingLoc(arg);

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

private:
   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;
+  // True if we find the main-file header in the Code.
+  bool MainIncludeFound;
   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;
 };
 
 } // namespace tooling
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLING_INCLUSIONS_HEADERINCLUDES_H

clang/lib/Tooling/Inclusions/HeaderIncludes.cpp

 HeaderIncludes::HeaderIncludes(StringRef FileName, StringRef Code,
                                const IncludeStyle &Style)
     : FileName(FileName), Code(Code), FirstIncludeOffset(-1),
       MinInsertOffset(
           getOffsetAfterHeaderGuardsAndComments(FileName, Code, Style)),
       MaxInsertOffset(MinInsertOffset +
                       getMaxHeaderInsertionOffset(
                           FileName, Code.drop_front(MinInsertOffset), Style)),
+      MainIncludeFound(false),
       Categories(Style, FileName) {
   // 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");
void HeaderIncludes::addExistingInclude(Include IncludeToAdd,
   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);
+        CurInclude.Name, /*CheckMainHeader=*/!MainIncludeFound);
+    if (Priority == 0)
+      MainIncludeFound = true;
     CategoryEndOffsets[Priority] = NextLineOffset;
     IncludesByPriority[Priority].push_back(&CurInclude);
     if (FirstIncludeOffset < 0)
       FirstIncludeOffset = CurInclude.R.getOffset();
   }
 }
 
 std::optional<tooling::Replacement>
for (const auto &Inc : It->second)
           ((IsAngled && StringRef(Inc.Name).startswith("<")) ||
            (!IsAngled && StringRef(Inc.Name).startswith("\""))))
         return std::nullopt;
   }
   std::string Quoted =
       std::string(llvm::formatv(IsAngled ? "<{0}>" : "\"{0}\"", IncludeName));
   StringRef QuotedName = Quoted;
   int Priority = Categories.getIncludePriority(
-      QuotedName, /*CheckMainHeader=*/FirstIncludeOffset < 0);
+      QuotedName, /*CheckMainHeader=*/!MainIncludeFound);
   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;
       }

clang/unittests/Tooling/HeaderIncludesTest.cpp

TEST_F(HeaderIncludesTest, InsertAfterMainHeader) {
 
   FileName = "fix_test.cu.cpp";
   EXPECT_EQ(Expected, insert(Code, "<a>"));
 
   FileName = "bar.cpp";
   EXPECT_NE(Expected, insert(Code, "<a>")) << "Not main header";
 }
 
+TEST_F(HeaderIncludesTest, InsertMainHeader) {
+  Style = format::getGoogleStyle(format::FormatStyle::LanguageKind::LK_Cpp)
+              .IncludeStyle;
+  FileName = "fix.cpp";
+  EXPECT_EQ(R"cpp(#include "fix.h"
+#include "a.h")cpp", insert("#include \"a.h\"", "\"fix.h\""));
+
+  // Respect the original main-file header.
+  EXPECT_EQ(R"cpp(#include "z/fix.h"
+#include "a/fix.h"
+)cpp", insert("#include \"z/fix.h\"", "\"a/fix.h\""));
+}
+
 TEST_F(HeaderIncludesTest, InsertBeforeSystemHeaderLLVM) {
   std::string Code = "#include <memory>\n"
                      "\n"
                      "int main() {}";
   std::string Expected = "#include \"z.h\"\n"
                          "#include <memory>\n"
                          "\n"
                          "int main() {}";