Fix crash in organizeDeclarations rule due to invalidated declaration…

… references (#1971)

Sources/Declaration.swift

@@ -247,11 +247,23 @@ final class TypeDeclaration: Declaration {
     func updateBody(to newBody: [Declaration]) {
         assert(!body.isEmpty)
 
-        // Store the expceted tokens associated with each declaration.
+        // Store the expected tokens associated with each declaration.
         // This is necessary since the declarations' range values will temporarily be invalid.
         var declarationTokens: [AnyHashable: [Token]] = [:]
+        var childDeclarationsNeedingUpdate = [AnyHashable: (originalIndexInParent: Int, originalTokens: [Token])]()
+
         for declaration in newBody {
             declarationTokens[declaration.identity] = Array(declaration.tokens)
+
+            // The body of this declaration won't be modified, but since we're update its range
+            // we have to also update the range of any children. Record the relative index of each child declaration
+            // so we can restore it later.
+            declaration.body?.forEachRecursiveDeclaration { childDeclaration in
+                let parent = declaration
+                let indexInParent = childDeclaration.range.lowerBound - parent.range.lowerBound
+                childDeclarationsNeedingUpdate[childDeclaration.identity] = (originalIndexInParent: indexInParent, originalTokens: Array(childDeclaration.tokens))
+                formatter.unregisterDeclaration(childDeclaration)
+            }
         }
 
         // Unlink the declarations and the formatter while we reorder the tokens
@@ -273,6 +285,20 @@ final class TypeDeclaration: Declaration {
 
             formatter.registerDeclaration(declaration)
             assert(Array(declaration.tokens) == tokens)
+            assert(declaration.isValid)
+
+            // Re-register each of the child declarations of this declaration
+            declaration.body?.forEachRecursiveDeclaration { childDeclaration in
+                guard let (originalIndexInParent, originalTokens) = childDeclarationsNeedingUpdate[childDeclaration.identity] else { return }
+                let parent = declaration
+                let newIndexInFile = parent.range.lowerBound + originalIndexInParent
+                let newRangeInFile = ClosedRange(newIndexInFile ..< (newIndexInFile + originalTokens.count))
+                childDeclaration.range = newRangeInFile
+
+                formatter.registerDeclaration(childDeclaration)
+                assert(Array(childDeclaration.tokens) == originalTokens)
+                assert(childDeclaration.isValid)
+            }
         }
 
         body = newBody

Tests/Rules/OrganizeDeclarationsTests.swift

@@ -3606,4 +3606,183 @@ class OrganizeDeclarationsTests: XCTestCase {
 
         testFormatting(for: input, rule: .organizeDeclarations)
     }
+
+    func testNoCrashWhenSortingNestedTypeDeclarations1() {
+        let input = """
+        public struct MyType {
+            var foo: Foo {
+                .foo
+            }
+
+            public let a: A
+            public let b: B
+            public let c: C
+            public let d: D
+            public let e: E
+
+            public enum Foo {
+                case foo
+                case bar
+                case baaz
+            }
+        }
+        """
+
+        let output = """
+        public struct MyType {
+
+            // MARK: Public
+
+            public enum Foo {
+                case foo
+                case bar
+                case baaz
+            }
+
+            public let a: A
+            public let b: B
+            public let c: C
+            public let d: D
+            public let e: E
+
+            // MARK: Internal
+
+            var foo: Foo {
+                .foo
+            }
+
+        }
+        """
+
+        let options = FormatOptions(organizeStructThreshold: 0)
+        testFormatting(for: input, output, rule: .organizeDeclarations, options: options, exclude: [.blankLinesAtStartOfScope, .blankLinesAtEndOfScope])
+    }
+
+    func testNoCrashWhenSortingNestedTypeDeclarations2() {
+        let input = """
+        public struct MyType {
+            public let a: A
+            public let b: B
+            public let c: C
+            public let d: D
+            public let e: E
+
+            public enum Foo {
+                case foo
+                case bar
+                case baaz
+            }
+        }
+        """
+
+        let output = """
+        public struct MyType {
+            public enum Foo {
+                case foo
+                case bar
+                case baaz
+            }
+
+            public let a: A
+            public let b: B
+            public let c: C
+            public let d: D
+            public let e: E
+
+        }
+        """
+
+        let options = FormatOptions(organizeStructThreshold: 0)
+        testFormatting(for: input, output, rule: .organizeDeclarations, options: options, exclude: [.blankLinesAtStartOfScope, .blankLinesAtEndOfScope])
+    }
+
+    func testSortsMultipleLayersOfNestedTypes() {
+        let input = """
+        public struct MyType {
+            public let a: A
+            public let b: B
+            public let c: C
+            public let d: D
+            public let e: E
+
+            public class Foo {
+                class Baaz {
+                    let b: B
+                    public let a: A
+
+                    public class Quux {
+                        let b: B
+                        public let a: A
+                    }
+                }
+
+                let bar: Bar
+                let baaz: Baaz
+
+                public class Bar {
+                    let b: B
+                    public let a: A
+                }
+            }
+        }
+        """
+
+        let output = """
+        public struct MyType {
+            public class Foo {
+
+                // MARK: Public
+
+                public class Bar {
+
+                    // MARK: Public
+
+                    public let a: A
+
+                    // MARK: Internal
+
+                    let b: B
+                }
+
+                // MARK: Internal
+
+                class Baaz {
+
+                    // MARK: Public
+
+                    public class Quux {
+
+                        // MARK: Public
+
+                        public let a: A
+
+                        // MARK: Internal
+
+                        let b: B
+                    }
+
+                    public let a: A
+
+                    // MARK: Internal
+
+                    let b: B
+
+                }
+
+                let bar: Bar
+                let baaz: Baaz
+
+            }
+
+            public let a: A
+            public let b: B
+            public let c: C
+            public let d: D
+            public let e: E
+
+        }
+        """
+
+        testFormatting(for: input, output, rule: .organizeDeclarations, exclude: [.blankLinesAtStartOfScope, .blankLinesAtEndOfScope])
+    }
 }