Support ObjC in IncludeInserter

[compositeprimes]

Summary

Update IncludeSorter/IncludeInserter to support objective-c google style (part 1):

1. Correctly consider .mm/.m extensions
2. Correctly categorize category headers.
3. Add support for generated files to go in a separate section of imports

[compositeprimes]

clang-format is complaining about my ordering of headers in my test case in IncludeInserterTest (but that's correct objc formatting). Not sure if there's a way to disable that on those lines

[alexfh]

LG

clang-tools-extra/clang-tidy/utils/IncludeSorter.h

 
 /// Class used by ``IncludeInserterCallback`` to record the names of the
 /// inclusions in a given source file being processed and generate the necessary
 /// commands to sort the inclusions according to the precedence encoded in
 /// ``IncludeKinds``.
 class IncludeSorter {
 public:
   /// Supported include styles.
-  enum IncludeStyle { IS_LLVM = 0, IS_Google = 1 };
+  enum IncludeStyle { IS_LLVM = 0, IS_Google = 1, IS_Google_ObjC };
 
   /// The classifications of inclusions, in the order they should be sorted.
   enum IncludeKinds {
     IK_MainTUInclude = 0,    ///< e.g. ``#include "foo.h"`` when editing foo.cc
     IK_CSystemInclude = 1,   ///< e.g. ``#include <stdio.h>``
     IK_CXXSystemInclude = 2, ///< e.g. ``#include <vector>``
     IK_NonSystemInclude = 3, ///< e.g. ``#include "bar.h"``
     IK_InvalidInclude = 4    ///< total number of valid ``IncludeKind``s

clang-tools-extra/clang-tidy/utils/IncludeSorter.cpp

 //===---------- IncludeSorter.cpp - clang-tidy ----------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "IncludeSorter.h"
+#include <algorithm>
 #include "clang/Basic/SourceManager.h"
 #include "clang/Lex/Lexer.h"
 
 namespace clang {
 namespace tidy {
 namespace utils {
 
 namespace {
 
 StringRef RemoveFirstSuffix(StringRef Str, ArrayRef<const char *> Suffixes) {
StringRef MakeCanonicalName(StringRef Str, IncludeSorter::IncludeStyle Style) {
   // E.g. tools/sort_includes.cc and tools/sort_includes_test.cc
   // would both canonicalize to tools/sort_includes and tools/sort_includes.h
   // (once canonicalized) will match as being the main include file associated
   // with the source files.
   if (Style == IncludeSorter::IS_LLVM) {
     return RemoveFirstSuffix(
         RemoveFirstSuffix(Str, {".cc", ".cpp", ".c", ".h", ".hpp"}), {"Test"});
   }
+  if (Style == IncludeSorter::IS_Google_ObjC) {
+    StringRef canonical =
+        RemoveFirstSuffix(RemoveFirstSuffix(Str, {".cc", ".cpp", ".c", ".h",
+                                                  ".hpp", ".mm", ".m"}),
+                          {"_unittest", "_regtest", "_test", "Test"});
+
+    // Objective-C categories have a `+suffix` format, but should be grouped
+    // with the file they are a category of.
+    return canonical.substr(
+        0, canonical.find_first_of('+', canonical.find_last_of('/')));
+  }
   return RemoveFirstSuffix(
       RemoveFirstSuffix(Str, {".cc", ".cpp", ".c", ".h", ".hpp"}),
       {"_unittest", "_regtest", "_test"});
 }
 
 // Scan to the end of the line and return the offset of the next line.
 size_t FindNextLine(const char *Text) {
   size_t EOLIndex = std::strcspn(Text, "\n");
if (IsAngled) {
     return IncludeFile.endswith(".h") ? IncludeSorter::IK_CSystemInclude
                                       : IncludeSorter::IK_CXXSystemInclude;
   }
   StringRef CanonicalInclude = MakeCanonicalName(IncludeFile, Style);
   if (CanonicalFile.endswith(CanonicalInclude)
       || CanonicalInclude.endswith(CanonicalFile)) {
     return IncludeSorter::IK_MainTUInclude;
   }
-  if (Style == IncludeSorter::IS_Google) {
+  if ((Style == IncludeSorter::IS_Google)
+      || (Style == IncludeSorter::IS_Google_ObjC)) {
     std::pair<StringRef, StringRef> Parts = CanonicalInclude.split("/public/");
     std::string AltCanonicalInclude =
         Parts.first.str() + "/internal/" + Parts.second.str();
     std::string ProtoCanonicalInclude =
         Parts.first.str() + "/proto/" + Parts.second.str();
 
     // Determine the kind of this inclusion.
     if (CanonicalFile.equals(AltCanonicalInclude) ||
         CanonicalFile.equals(ProtoCanonicalInclude)) {
       return IncludeSorter::IK_MainTUInclude;
     }
   }
   return IncludeSorter::IK_NonSystemInclude;
 }
 
+
+int CompareHeaders(StringRef LHS, StringRef RHS,
+                   IncludeSorter::IncludeStyle Style) {
+  if (Style == IncludeSorter::IncludeStyle::IS_Google_ObjC) {
+    const std::pair<const char *, const char *> &mismatch =
+        std::mismatch(LHS.begin(), LHS.end(), RHS.begin());
+    if ((mismatch.first != LHS.end()) && (mismatch.second != RHS.end())) {
+      if ((*mismatch.first == '.') && (*mismatch.second == '+')) {
+        return -1;
+      }
+      if ((*mismatch.first == '+') && (*mismatch.second == '.')) {
+        return 1;
+      }
+    }
+  }
+  return LHS.compare(RHS);
+}
+
 } // namespace
 
 IncludeSorter::IncludeSorter(const SourceManager *SourceMgr,
                              const FileID FileID, StringRef FileName,
                              IncludeStyle Style)
     : SourceMgr(SourceMgr), Style(Style), CurrentFileID(FileID),
       CanonicalFile(MakeCanonicalName(FileName, Style)) {}
 
void IncludeSorter::AddInclude(StringRef FileName, bool IsAngled,
   IncludeKinds Kind =
       DetermineIncludeKind(CanonicalFile, FileName, IsAngled, Style);
   if (Kind != IK_InvalidInclude)
     IncludeBucket[Kind].push_back(FileName.str());
 }
 
 Optional<FixItHint> IncludeSorter::CreateIncludeInsertion(StringRef FileName,
                                                           bool IsAngled) {
-  std::string IncludeStmt =
-      IsAngled ? llvm::Twine("#include <" + FileName + ">\n").str()
+  std::string IncludeStmt;
+  if (Style == IncludeStyle::IS_Google_ObjC) {
+    IncludeStmt = IsAngled
+                      ? llvm::Twine("#import <" + FileName + ">\n").str()
+                      : llvm::Twine("#import \"" + FileName + "\"\n").str();
+  } else {
+    IncludeStmt = IsAngled
+                      ? llvm::Twine("#include <" + FileName + ">\n").str()
                       : llvm::Twine("#include \"" + FileName + "\"\n").str();
+  }
   if (SourceLocations.empty()) {
     // If there are no includes in this file, add it in the first line.
     // FIXME: insert after the file comment or the header guard, if present.
     IncludeStmt.append("\n");
     return FixItHint::CreateInsertion(
         SourceMgr->getLocForStartOfFile(CurrentFileID), IncludeStmt);
   }
 
   auto IncludeKind =
       DetermineIncludeKind(CanonicalFile, FileName, IsAngled, Style);
 
   if (!IncludeBucket[IncludeKind].empty()) {
     for (const std::string &IncludeEntry : IncludeBucket[IncludeKind]) {
-      if (FileName < IncludeEntry) {
+      if (CompareHeaders(FileName, IncludeEntry, Style) < 0) {
         const auto &Location = IncludeLocations[IncludeEntry][0];
         return FixItHint::CreateInsertion(Location.getBegin(), IncludeStmt);
       } else if (FileName == IncludeEntry) {
         return llvm::None;
       }
     }
     // FileName comes after all include entries in bucket, insert it after
     // last.

 }
 
 } // namespace utils
 
 llvm::ArrayRef<std::pair<utils::IncludeSorter::IncludeStyle, StringRef>>
 OptionEnumMapping<utils::IncludeSorter::IncludeStyle>::getEnumMapping() {
   static constexpr std::pair<utils::IncludeSorter::IncludeStyle, StringRef>
       Mapping[] = {{utils::IncludeSorter::IS_LLVM, "llvm"},
-                   {utils::IncludeSorter::IS_Google, "google"}};
+                   {utils::IncludeSorter::IS_Google, "google"},
+                   {utils::IncludeSorter::IS_Google_ObjC, "google-objc"}};
   return makeArrayRef(Mapping);
 }
 } // namespace tidy
 } // namespace clang

clang-tools-extra/unittests/clang-tidy/IncludeInserterTest.cpp

 #if !defined(_WIN32)
 
 namespace clang {
 namespace tidy {
 namespace {
 
 class IncludeInserterCheckBase : public ClangTidyCheck {
 public:
-  IncludeInserterCheckBase(StringRef CheckName, ClangTidyContext *Context)
-      : ClangTidyCheck(CheckName, Context) {}
+  IncludeInserterCheckBase(StringRef CheckName, ClangTidyContext *Context,
+                           utils::IncludeSorter::IncludeStyle Style =
+                               utils::IncludeSorter::IS_Google)
+      : ClangTidyCheck(CheckName, Context), Inserter(Style) {}
 
   void registerPPCallbacks(const SourceManager &SM, Preprocessor *PP,
                            Preprocessor *ModuleExpanderPP) override {
     Inserter.registerPreprocessor(PP);
   }
 
   void registerMatchers(ast_matchers::MatchFinder *Finder) override {
     Finder->addMatcher(ast_matchers::declStmt().bind("stmt"), this);
   }
 
   void check(const ast_matchers::MatchFinder::MatchResult &Result) override {
     auto Diag = diag(Result.Nodes.getNodeAs<DeclStmt>("stmt")->getBeginLoc(),
                      "foo, bar");
     for (StringRef Header : headersToInclude()) {
       Diag << Inserter.createMainFileIncludeInsertion(Header);
     }
   }
 
   virtual std::vector<StringRef> headersToInclude() const = 0;
 
-  utils::IncludeInserter Inserter{utils::IncludeSorter::IS_Google};
+  utils::IncludeInserter Inserter;
 };
 
 class NonSystemHeaderInserterCheck : public IncludeInserterCheckBase {
 public:
   NonSystemHeaderInserterCheck(StringRef CheckName, ClangTidyContext *Context)
       : IncludeInserterCheckBase(CheckName, Context) {}
 
   std::vector<StringRef> headersToInclude() const override {
public:
   InvalidIncludeInserterCheck(StringRef CheckName, ClangTidyContext *Context)
       : IncludeInserterCheckBase(CheckName, Context) {}
 
   std::vector<StringRef> headersToInclude() const override {
     return {"a.h", "<stdlib.h", "cstdlib>", "b.h", "<c.h>", "<d>"};
   }
 };
 
+class ObjCEarlyInAlphabetHeaderInserterCheck : public IncludeInserterCheckBase {
+public:
+  ObjCEarlyInAlphabetHeaderInserterCheck(StringRef CheckName,
+                                         ClangTidyContext *Context)
+      : IncludeInserterCheckBase(CheckName, Context,
+                                 utils::IncludeSorter::IS_Google_ObjC) {}
+
+  std::vector<StringRef> headersToInclude() const override {
+    return {"a/header.h"};
+  }
+};
+
+class ObjCCategoryHeaderInserterCheck : public IncludeInserterCheckBase {
+public:
+  ObjCCategoryHeaderInserterCheck(StringRef CheckName,
+                                         ClangTidyContext *Context)
+      : IncludeInserterCheckBase(CheckName, Context,
+                                 utils::IncludeSorter::IS_Google_ObjC) {}
+
+  std::vector<StringRef> headersToInclude() const override {
+    return {"clang_tidy/tests/insert_includes_test_header+foo.h"};
+  }
+};
+
 template <typename Check>
 std::string runCheckOnCode(StringRef Code, StringRef Filename) {
   std::vector<ClangTidyError> Errors;
   return test::runCheckOnCode<Check>(Code, &Errors, Filename, None,
                                      ClangTidyOptions(),
                                      {// Main file include
                                       {"clang_tidy/tests/"
                                        "insert_includes_test_header.h",
                                        "\n"},
+                                      // ObjC category.
+                                      {"clang_tidy/tests/"
+                                       "insert_includes_test_header+foo.h",
+                                       "\n"},
                                       // Non system headers
                                       {"a/header.h", "\n"},
                                       {"path/to/a/header.h", "\n"},
                                       {"path/to/z/header.h", "\n"},
                                       {"path/to/header.h", "\n"},
                                       {"path/to/header2.h", "\n"},
                                       // Fake system headers.
                                       {"stdlib.h", "\n"},
void foo() {
   int a = 0;
 })";
 
   EXPECT_EQ(PostCode, runCheckOnCode<InvalidIncludeInserterCheck>(
                           PreCode, "clang_tidy/tests/"
                                    "insert_includes_test_header.cc"));
 }
 
+TEST(IncludeInserterTest, InsertHeaderObjectiveC) {
+  const char *PreCode = R"(
+#import "clang_tidy/tests/insert_includes_test_header.h"
+
+void foo() {
+  int a = 0;
+})";
+  const char *PostCode = R"(
+#import "clang_tidy/tests/insert_includes_test_header.h"
+
+#import "a/header.h"
+
+void foo() {
+  int a = 0;
+})";
+
+  EXPECT_EQ(PostCode, runCheckOnCode<ObjCEarlyInAlphabetHeaderInserterCheck>(
+                          PreCode, "repo/clang_tidy/tests/"
+                                   "insert_includes_test_header.mm"));
+}
+
+TEST(IncludeInserterTest, InsertCategoryHeaderObjectiveC) {
+  const char *PreCode = R"(
+#import "clang_tidy/tests/insert_includes_test_header.h"
+
+void foo() {
+  int a = 0;
+})";
+  const char *PostCode = R"(
+#import "clang_tidy/tests/insert_includes_test_header.h"
+#import "clang_tidy/tests/insert_includes_test_header+foo.h"
+
+void foo() {
+  int a = 0;
+})";
+
+  EXPECT_EQ(PostCode, runCheckOnCode<ObjCCategoryHeaderInserterCheck>(
+                          PreCode, "devtools/cymbal/clang_tidy/tests/"
+                                   "insert_includes_test_header.mm"));
+}
+
 } // anonymous namespace
 } // namespace tidy
 } // namespace clang
 
 #endif