From cedf93a632511a985baf1e281cafd95fafa455ad Mon Sep 17 00:00:00 2001
From: Cal Stephens <cal@calstephens.tech>
Date: Wed, 21 Feb 2024 10:26:05 -0800
Subject: [PATCH] Fix issue where extension body wouldn't be sorted if
 organizeDeclarations was enabled but excluded declaration type (#1628)

---
 Sources/Rules.swift                 | 25 +++++++---
 Tests/RulesTests+Organization.swift | 75 ++++++++++++++++++++++++-----
 2 files changed, 81 insertions(+), 19 deletions(-)

diff --git a/Sources/Rules.swift b/Sources/Rules.swift
index ebc66dc4..fc12b340 100644
--- a/Sources/Rules.swift
+++ b/Sources/Rules.swift
@@ -5940,7 +5940,8 @@ public struct _FormatRules {
         Sorts the body of declarations with // swiftformat:sort
         and declarations between // swiftformat:sort:begin and
         // swiftformat:sort:end comments.
-        """
+        """,
+        sharedOptions: ["organizetypes"]
     ) { formatter in
         formatter.forEachToken(
             where: { $0.isComment && $0.string.contains("swiftformat:sort") }
@@ -5966,20 +5967,28 @@ public struct _FormatRules {
 
             // For `:sort` directives, we sort the declarations
             // between the open and close brace of the following type
-            else if !commentToken.string.contains(":sort:"),
-                    // This part of the rule conflicts with the organizeDeclarations rule.
-                    // Instead, that rule manually implements support for the :sort directive.
-                    !formatter.options.enabledRules.contains(FormatRules.organizeDeclarations.name)
-            {
+            else if !commentToken.string.contains(":sort:") {
                 guard let typeOpenBrace = formatter.index(of: .startOfScope("{"), after: commentIndex),
                       let typeCloseBrace = formatter.endOfScope(at: typeOpenBrace),
                       let firstTypeBodyToken = formatter.index(of: .nonLinebreak, after: typeOpenBrace),
                       let lastTypeBodyToken = formatter.index(of: .nonLinebreak, before: typeCloseBrace),
+                      let declarationKeyword = formatter.lastSignificantKeyword(at: typeOpenBrace),
                       lastTypeBodyToken > typeOpenBrace
                 else { return }
 
-                rangeToSort = typeOpenBrace + 1 ... lastTypeBodyToken
-                numberOfLeadingLinebreaks = firstTypeBodyToken - typeOpenBrace - 1
+                // Sorting the body of a type conflicts with the `organizeDeclaration`
+                // keyword if enabled for this type of declaration. In that case,
+                // defer to the sorting implementation in `organizeDeclarations`.
+                if formatter.options.enabledRules.contains(FormatRules.organizeDeclarations.name),
+                   formatter.options.organizeTypes.contains(declarationKeyword)
+                {
+                    return
+                }
+
+                rangeToSort = firstTypeBodyToken ... lastTypeBodyToken
+                // We don't include any leading linebreaks in the range to sort,
+                // since `firstTypeBodyToken` is the first `nonLinebreak` in the body
+                numberOfLeadingLinebreaks = 0
             } else {
                 return
             }
diff --git a/Tests/RulesTests+Organization.swift b/Tests/RulesTests+Organization.swift
index 3b754366..19f4a2b0 100644
--- a/Tests/RulesTests+Organization.swift
+++ b/Tests/RulesTests+Organization.swift
@@ -3286,34 +3286,24 @@ class OrganizationTests: RulesTests {
         testFormatting(for: input, rule: FormatRules.sortDeclarations)
     }
 
-    func testTypeBodyWithBlankLines() {
+    func testSortsTypeBody() {
         let input = """
         // swiftformat:sort
         enum FeatureFlags {
-
             case upsellB
-
             case fooFeature
-
             case barFeature
-
             case upsellA
-
         }
         """
 
         let output = """
         // swiftformat:sort
         enum FeatureFlags {
-
             case barFeature
-
             case fooFeature
-
             case upsellA
-
             case upsellB
-
         }
         """
 
@@ -3522,6 +3512,69 @@ class OrganizationTests: RulesTests {
         testFormatting(for: input, rule: FormatRules.organizeDeclarations)
     }
 
+    func testSortDeclarationsSortsExtensionBody() {
+        let input = """
+        enum Namespace {}
+
+        // swiftformat:sort
+        extension Namespace {
+            static let foo = "foo"
+            public static let bar = "bar"
+            static let baaz = "baaz"
+        }
+        """
+
+        let output = """
+        enum Namespace {}
+
+        // swiftformat:sort
+        extension Namespace {
+            static let baaz = "baaz"
+            public static let bar = "bar"
+            static let foo = "foo"
+        }
+        """
+
+        // organizeTypes doesn't include "extension". So even though the
+        // organizeDeclarations rule is enabled, the extension should be
+        // sorted by the sortDeclarations rule.
+        let options = FormatOptions(organizeTypes: ["class"])
+        testFormatting(for: input, [output], rules: [FormatRules.sortDeclarations, FormatRules.organizeDeclarations], options: options)
+    }
+
+    func testOrganizeDeclarationsSortsExtensionBody() {
+        let input = """
+        enum Namespace {}
+
+        // swiftformat:sort
+        extension Namespace {
+            static let foo = "foo"
+            public static let bar = "bar"
+            static let baaz = "baaz"
+        }
+        """
+
+        let output = """
+        enum Namespace {}
+
+        // swiftformat:sort
+        extension Namespace {
+
+            // MARK: Public
+
+            public static let bar = "bar"
+
+            // MARK: Internal
+
+            static let baaz = "baaz"
+            static let foo = "foo"
+        }
+        """
+
+        let options = FormatOptions(organizeTypes: ["extension"])
+        testFormatting(for: input, output, rule: FormatRules.organizeDeclarations, options: options, exclude: ["blankLinesAtStartOfScope", "blankLinesAtEndOfScope"])
+    }
+
     // MARK: - sortTypealiases
 
     func testSortSingleLineTypealias() {