From 8fa0e7868763af0ae2d06247779addb9cc80707f Mon Sep 17 00:00:00 2001
From: Cal Stephens <cal.stephens@airbnb.com>
Date: Fri, 13 Nov 2020 09:59:22 -0800
Subject: [PATCH] Increase strictness of grouped extensions

---
 Sources/Rules.swift                 | 23 ++++++++++--
 Tests/RulesTests+Organization.swift | 54 +++++++++++++++++++++++++++++
 Tests/XCTestManifests.swift         |  3 ++
 3 files changed, 78 insertions(+), 2 deletions(-)

diff --git a/Sources/Rules.swift b/Sources/Rules.swift
index 7547d492e..d97f9b52a 100644
--- a/Sources/Rules.swift
+++ b/Sources/Rules.swift
@@ -5125,9 +5125,28 @@ public struct _FormatRules {
                 //  extension Foo { } // This extension is "grouped" with its extending type
                 //  extension String { } // This extension is standalone (not grouped with any type)
                 //
-                let isGroupedWithExtendingType = declarations[..<index].contains(where: {
+                let isGroupedWithExtendingType: Bool
+                if let indexOfExtendingType = declarations[..<index].lastIndex(where: {
                     $0.name == typeName && ["class", "enum", "protocol", "struct", "typealias"].contains($0.keyword)
-                })
+                }) {
+                    let declarationsBetweenTypeAndExtension = declarations[indexOfExtendingType + 1 ..< index]
+                    isGroupedWithExtendingType = declarationsBetweenTypeAndExtension.allSatisfy {
+                        // Only treat the type and its extension as grouped if there aren't any other
+                        // types or type-like declaraitons between them
+                        if ["class", "enum", "protocol", "struct", "typealias"].contains($0.keyword) {
+                            return false
+                        }
+
+                        // Extensions extending other types also break the grouping
+                        if $0.keyword == "extension", $0.name != declaration.name {
+                            return false
+                        }
+
+                        return true
+                    }
+                } else {
+                    isGroupedWithExtendingType = false
+                }
 
                 if isGroupedWithExtendingType {
                     commentTemplate = "// \(formatter.options.groupedExtensionMarkComment)"
diff --git a/Tests/RulesTests+Organization.swift b/Tests/RulesTests+Organization.swift
index badc6acd7..601925a3f 100644
--- a/Tests/RulesTests+Organization.swift
+++ b/Tests/RulesTests+Organization.swift
@@ -2191,4 +2191,58 @@ extension RulesTests {
 
         testFormatting(for: input, rule: FormatRules.markTypes)
     }
+
+    func testDoesntUseGroupedMarkTemplateWhenSeparatedByOtherType() {
+        let input = """
+        // MARK: - MyComponent
+
+        class MyComponent {}
+
+        // MARK: - MyComponentContent
+
+        struct MyComponentContent {}
+
+        // MARK: - MyComponent + ContentConfigurableView
+
+        extension MyComponent: ContentConfigurableView {}
+        """
+
+        testFormatting(for: input, rule: FormatRules.markTypes)
+    }
+
+    func testUsesGroupedMarkTemplateWhenSeparatedByExtensionOfSameType() {
+        let input = """
+        // MARK: - MyComponent
+
+        class MyComponent {}
+
+        // MARK: Equatable
+
+        extension MyComponent: Equatable {}
+
+        // MARK: ContentConfigurableView
+
+        extension MyComponent: ContentConfigurableView {}
+        """
+
+        testFormatting(for: input, rule: FormatRules.markTypes)
+    }
+
+    func testDoesntUseGroupedMarkTemplateWhenSeparatedByExtensionOfOtherType() {
+        let input = """
+        // MARK: - MyComponent
+
+        class MyComponent {}
+
+        // MARK: - OtherComponent + Equatable
+
+        extension OtherComponent: Equatable {}
+
+        // MARK: - MyComponent + ContentConfigurableView
+
+        extension MyComponent: ContentConfigurableView {}
+        """
+
+        testFormatting(for: input, rule: FormatRules.markTypes)
+    }
 }
diff --git a/Tests/XCTestManifests.swift b/Tests/XCTestManifests.swift
index 0aec85d22..bf6a9c67d 100644
--- a/Tests/XCTestManifests.swift
+++ b/Tests/XCTestManifests.swift
@@ -688,6 +688,8 @@ extension RulesTests {
         ("testDoesntUpdateExtensionVisibilityWithInternalDeclarations", testDoesntUpdateExtensionVisibilityWithInternalDeclarations),
         ("testDoesntUpdateExtensionVisibilityWithoutMajorityBodyVisibility", testDoesntUpdateExtensionVisibilityWithoutMajorityBodyVisibility),
         ("testDoesntUpdatesExtensionThatHasLowerACLThanBodyDeclarations", testDoesntUpdatesExtensionThatHasLowerACLThanBodyDeclarations),
+        ("testDoesntUseGroupedMarkTemplateWhenSeparatedByExtensionOfOtherType", testDoesntUseGroupedMarkTemplateWhenSeparatedByExtensionOfOtherType),
+        ("testDoesntUseGroupedMarkTemplateWhenSeparatedByOtherType", testDoesntUseGroupedMarkTemplateWhenSeparatedByOtherType),
         ("testDoesntWrapReturnOnSingleLineFunctionDeclaration", testDoesntWrapReturnOnSingleLineFunctionDeclaration),
         ("testDontChangePrivateExtensionToFileprivate", testDontChangePrivateExtensionToFileprivate),
         ("testDontCorruptPartialFragment", testDontCorruptPartialFragment),
@@ -2005,6 +2007,7 @@ extension RulesTests {
         ("testUppercaseGroupedHexExponent", testUppercaseGroupedHexExponent),
         ("testUppercaseHexExponent", testUppercaseHexExponent),
         ("testUppercaseLiteralConvertedToLower", testUppercaseLiteralConvertedToLower),
+        ("testUsesGroupedMarkTemplateWhenSeparatedByExtensionOfSameType", testUsesGroupedMarkTemplateWhenSeparatedByExtensionOfSameType),
         ("testUseVoidOptionFalse", testUseVoidOptionFalse),
         ("testVarAttributeIsNotWrapped", testVarAttributeIsNotWrapped),
         ("testVarModifiersCorrected", testVarModifiersCorrected),