clang-format: Extend detection of the "main" #include to use the filename

Before, the first (non-system) header in a file was considered to be
the main include. This is conservative as it makes clang-format change
the #include order less often. Instead implement some basic usage of
the filename itself. With this patch, clang-format considers every
header to be a main include if the header file's basename is a prefix
to the filename the #include is in.

llvm-svn: 256148

2 files changed, 41 insertions, 23 deletions

diff --git a/clang/lib/Format/Format.cpp b/clang/lib/Format/Format.cpp
index 97195aa9462..c9058a587e6 100644
--- a/clang/lib/Format/Format.cpp
+++ b/clang/lib/Format/Format.cpp
@@ -1809,13 +1809,11 @@ tooling::Replacements sortIncludes(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 LookForMainHeader = FileName.endswith(".c") ||
-                           FileName.endswith(".cc") ||
-                           FileName.endswith(".cpp")||
-                           FileName.endswith(".c++")||
-                           FileName.endswith(".cxx") ||
-                           FileName.endswith(".m")||
-                           FileName.endswith(".mm");
+  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);
 
   // Create pre-compiled regular expressions for the #include categories.
   SmallVector<llvm::Regex, 4> CategoryRegexs;
@@ -1838,19 +1836,19 @@ tooling::Replacements sortIncludes(const FormatStyle &Style, StringRef Code,
     if (!FormattingOff && !Line.endswith("\\")) {
       if (IncludeRegex.match(Line, &Matches)) {
         StringRef IncludeName = Matches[2];
-        int Category;
-        if (LookForMainHeader && !IncludeName.startswith("<")) {
-          Category = 0;
-        } else {
-          Category = INT_MAX;
-          for (unsigned i = 0, e = CategoryRegexs.size(); i != e; ++i) {
-            if (CategoryRegexs[i].match(IncludeName)) {
-              Category = Style.IncludeCategories[i].Priority;
-              break;
-            }
+        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;
           }
         }
-        LookForMainHeader = false;
+        if (IsSource && Category > 0 && IncludeName.startswith("\"")) {
+          StringRef HeaderStem =
+              llvm::sys::path::stem(IncludeName.drop_front(1).drop_back(1));
+          if (FileStem.startswith(HeaderStem))
+            Category = 0;
+        }
         IncludesInBlock.push_back({IncludeName, Line, Prev, Category});
       } else if (!IncludesInBlock.empty()) {
         sortIncludes(Style, IncludesInBlock, Ranges, FileName, Replaces,
diff --git a/clang/unittests/Format/SortIncludesTest.cpp b/clang/unittests/Format/SortIncludesTest.cpp
index 9b0db349b04..ce83091b715 100644
--- a/clang/unittests/Format/SortIncludesTest.cpp
+++ b/clang/unittests/Format/SortIncludesTest.cpp
@@ -166,13 +166,21 @@ TEST_F(SortIncludesTest, LeavesMainHeaderFirst) {
             "#include \"c.h\"\n",
             sort("#include \"llvm/a.h\"\n"
                  "#include \"c.h\"\n"
-                 "#include \"b.h\"\n"));
+                 "#include \"b.h\"\n",
+                 "a.cc"));
   EXPECT_EQ("#include \"llvm/a.h\"\n"
             "#include \"b.h\"\n"
             "#include \"c.h\"\n",
             sort("#include \"llvm/a.h\"\n"
                  "#include \"c.h\"\n"
                  "#include \"b.h\"\n",
+                 "a_main.cc"));
+  EXPECT_EQ("#include \"llvm/input.h\"\n"
+            "#include \"b.h\"\n"
+            "#include \"c.h\"\n",
+            sort("#include \"llvm/input.h\"\n"
+                 "#include \"c.h\"\n"
+                 "#include \"b.h\"\n",
                  "input.mm"));
 
   // Don't do this in headers.
@@ -182,15 +190,27 @@ TEST_F(SortIncludesTest, LeavesMainHeaderFirst) {
             sort("#include \"llvm/a.h\"\n"
                  "#include \"c.h\"\n"
                  "#include \"b.h\"\n",
-                 "some_header.h"));
+                 "a.h"));
 }
 
 TEST_F(SortIncludesTest, NegativePriorities) {
   Style.IncludeCategories = {{".*important_os_header.*", -1}, {".*", 1}};
   EXPECT_EQ("#include \"important_os_header.h\"\n"
-            "#include \"a.h\"\n",
-            sort("#include \"a.h\"\n"
-                 "#include \"important_os_header.h\"\n"));
+            "#include \"c_main.h\"\n"
+            "#include \"a_other.h\"\n",
+            sort("#include \"c_main.h\"\n"
+                 "#include \"a_other.h\"\n"
+                 "#include \"important_os_header.h\"\n",
+                 "c_main.cc"));
+
+  // check stable when re-run
+  EXPECT_EQ("#include \"important_os_header.h\"\n"
+            "#include \"c_main.h\"\n"
+            "#include \"a_other.h\"\n",
+            sort("#include \"important_os_header.h\"\n"
+                 "#include \"c_main.h\"\n"
+                 "#include \"a_other.h\"\n",
+                 "c_main.cc"));
 }
 
 TEST_F(SortIncludesTest, CalculatesCorrectCursorPosition) {