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

[Haojian Wu via cfe-commits]

Author: Haojian Wu
Date: 2023-07-13T10:41:46+02:00
New Revision: 85c6d57eb304482264964778d89fd0554df07592

URL: https://github.com/llvm/llvm-project/commit/85c6d57eb304482264964778d89fd0554df07592
DIFF: https://github.com/llvm/llvm-project/commit/85c6d57eb304482264964778d89fd0554df07592.diff

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

Differential Revision: https://reviews.llvm.org/D154963

Added: 
    

Modified: 
    clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp
    clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp
    clang/include/clang/Tooling/Inclusions/HeaderIncludes.h
    clang/lib/Tooling/Inclusions/HeaderIncludes.cpp
    clang/unittests/Tooling/HeaderIncludesTest.cpp

Removed: 
    


################################################################################
diff  --git a/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp b/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp
index 6397f32c525f1d..aefdba05dace7a 100644
--- a/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp
+++ b/clang-tools-extra/clangd/unittests/IncludeCleanerTests.cpp
@@ -214,7 +214,7 @@ TEST(IncludeCleaner, GenerateMissingHeaderDiags) {
 })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();");
 

diff  --git a/clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp b/clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp
index 64dd6bc1e84cae..3ad632a18c563a 100644
--- a/clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp
+++ b/clang-tools-extra/include-cleaner/unittests/AnalysisTest.cpp
@@ -311,12 +311,9 @@ R"cpp(#include "d.h"
   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, "") {

diff  --git a/clang/include/clang/Tooling/Inclusions/HeaderIncludes.h b/clang/include/clang/Tooling/Inclusions/HeaderIncludes.h
index 9e08046d2e33a8..d5439dd2c84eb7 100644
--- a/clang/include/clang/Tooling/Inclusions/HeaderIncludes.h
+++ b/clang/include/clang/Tooling/Inclusions/HeaderIncludes.h
@@ -128,6 +128,8 @@ class HeaderIncludes {
   // 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;

diff  --git a/clang/lib/Tooling/Inclusions/HeaderIncludes.cpp b/clang/lib/Tooling/Inclusions/HeaderIncludes.cpp
index b37c32a9ef62e6..15a2024c478887 100644
--- a/clang/lib/Tooling/Inclusions/HeaderIncludes.cpp
+++ b/clang/lib/Tooling/Inclusions/HeaderIncludes.cpp
@@ -276,6 +276,7 @@ HeaderIncludes::HeaderIncludes(StringRef FileName, StringRef Code,
       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.
@@ -335,7 +336,9 @@ void HeaderIncludes::addExistingInclude(Include IncludeToAdd,
   // 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)
@@ -362,7 +365,7 @@ HeaderIncludes::insert(llvm::StringRef IncludeName, bool IsAngled,
       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

diff  --git a/clang/unittests/Tooling/HeaderIncludesTest.cpp b/clang/unittests/Tooling/HeaderIncludesTest.cpp
index cd48d38cdb0a8c..256aa825554c50 100644
--- a/clang/unittests/Tooling/HeaderIncludesTest.cpp
+++ b/clang/unittests/Tooling/HeaderIncludesTest.cpp
@@ -143,6 +143,19 @@ TEST_F(HeaderIncludesTest, InsertAfterMainHeader) {
   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"