Fix various edge cases for spaceAroundBrackets

Sources/FormattingHelpers.swift

@@ -36,6 +36,54 @@ extension Formatter {
         return false
     }
 
+    /// Should the specified token be followed by a space if next token is an opening paren, bracket, etc?
+    func shouldInsertSpaceAfterToken(at index: Int) -> Bool? {
+        switch token(at: index) {
+        case let .keyword(keywordOrAttribute):
+            switch keywordOrAttribute {
+            case "@autoclosure":
+                if options.swiftVersion < "3",
+                   let nextIndex = self.index(of: .nonSpaceOrLinebreak, after: index),
+                   next(.nonSpaceOrCommentOrLinebreak, after: nextIndex) == .identifier("escaping")
+                {
+                    assert(tokens[nextIndex] == .startOfScope("("))
+                    return false
+                }
+                return true
+            case "@escaping", "@noescape", "@Sendable", "@MainActor":
+                return true
+            case _ where keywordOrAttribute.isAttribute:
+                if let i = self.index(of: .startOfScope("("), after: index) {
+                    return isParameterList(at: i)
+                }
+                return false
+            case "private", "fileprivate", "internal", "init", "subscript", "throws":
+                return false
+            case "await":
+                return options.swiftVersion >= "5.5" || options.swiftVersion == .undefined
+            default:
+                return !keywordOrAttribute.isMacroOrAttribute
+            }
+        case let .identifier(name):
+            switch name {
+            case "as", "is", "try": // not treated as keywords inside macro
+                return token(at: index - 1)?.isOperator(".") != true
+            case "unsafe":
+                return options.swiftVersion >= "6.2" || options.swiftVersion == .undefined
+            default:
+                return name.isKeywordInTypeContext && isTypePosition(at: index)
+            }
+        case .endOfScope("]"):
+            return isInClosureArguments(at: index)
+        case .endOfScope(")"):
+            return isAttribute(at: index)
+        case .number, .endOfScope("}"), .endOfScope(">"):
+            return false
+        default:
+            return nil
+        }
+    }
+
     /// remove self if possible
     func removeSelf(at i: Int, exclude: Set<String>, include: Set<String>? = nil) -> Bool {
         guard case let .identifier(selfKeyword) = tokens[i], ["self", "Self"].contains(selfKeyword) else {

Sources/Rules/SpaceAroundBrackets.swift

@@ -20,43 +20,26 @@ public extension FormatRule {
         help: "Add or remove space around square brackets."
     ) { formatter in
         formatter.forEach(.startOfScope("[")) { i, _ in
-            let index = i - 1
-            guard let prevToken = formatter.token(at: index) else {
-                return
-            }
-            switch prevToken {
-            case .identifier where prevToken.isKeywordInTypeContext && formatter.isTypePosition(at: index), .keyword:
-                formatter.insert(.space(" "), at: i)
-            case .space:
-                let index = i - 2
-                if let token = formatter.token(at: index) {
-                    switch token {
-                    case .identifier("as"), .identifier("is"), // not treated as keywords inside macro
-                         .identifier where token.isKeywordInTypeContext && formatter.isTypePosition(at: index),
-                         .identifier("try"), .keyword("try"):
-                        break
-                    case .identifier, .number, .endOfScope("]"), .endOfScope("}"), .endOfScope(")"):
-                        formatter.removeToken(at: i - 1)
-                    default:
-                        break
-                    }
-                }
+            let i = i - 1
+            switch formatter.token(at: i) {
+            case _ where formatter.shouldInsertSpaceAfterToken(at: i) == true:
+                formatter.insert(.space(" "), at: i + 1)
+            case .space where formatter.shouldInsertSpaceAfterToken(at: i - 1) == false:
+                formatter.removeToken(at: i)
             default:
                 break
             }
         }
         formatter.forEach(.endOfScope("]")) { i, _ in
-            guard let nextToken = formatter.token(at: i + 1) else {
-                return
-            }
-            switch nextToken {
+            let i = i + 1
+            switch formatter.token(at: i) {
             case .identifier, .keyword, .startOfScope("{"),
-                 .startOfScope("(") where formatter.isInClosureArguments(at: i):
-                formatter.insert(.space(" "), at: i + 1)
+                 .startOfScope("(") where formatter.isInClosureArguments(at: i - 1):
+                formatter.insert(.space(" "), at: i)
             case .space:
-                switch formatter.token(at: i + 2) {
-                case .startOfScope("(")? where !formatter.isInClosureArguments(at: i + 2), .startOfScope("[")?:
-                    formatter.removeToken(at: i + 1)
+                switch formatter.token(at: i + 1) {
+                case .startOfScope("(")? where !formatter.isInClosureArguments(at: i + 1), .startOfScope("[")?:
+                    formatter.removeToken(at: i)
                 default:
                     break
                 }

Sources/Rules/SpaceAroundParens.swift

@@ -56,47 +56,3 @@ public extension FormatRule {
         """
     }
 }
-
-extension Formatter {
-    func shouldInsertSpaceAfterToken(at index: Int) -> Bool? {
-        switch token(at: index) {
-        case let .keyword(keywordOrAttribute):
-            switch keywordOrAttribute {
-            case "@autoclosure":
-                if options.swiftVersion < "3",
-                   let nextIndex = self.index(of: .nonSpaceOrLinebreak, after: index),
-                   next(.nonSpaceOrCommentOrLinebreak, after: nextIndex) == .identifier("escaping")
-                {
-                    assert(tokens[nextIndex] == .startOfScope("("))
-                    return false
-                }
-                return true
-            case "@escaping", "@noescape", "@Sendable":
-                return true
-            case _ where keywordOrAttribute.isAttribute:
-                if let i = self.index(of: .startOfScope("("), after: index) {
-                    return isParameterList(at: i)
-                }
-                return false
-            case "private", "fileprivate", "internal", "init", "subscript", "throws":
-                return false
-            case "await":
-                return options.swiftVersion >= "5.5" || options.swiftVersion == .undefined
-            default:
-                return !keywordOrAttribute.isMacroOrAttribute
-            }
-        case .identifier("unsafe"):
-            return options.swiftVersion >= "6.2" || options.swiftVersion == .undefined
-        case let .identifier(name):
-            return name.isKeywordInTypeContext && isTypePosition(at: index)
-        case .endOfScope("]"):
-            return isInClosureArguments(at: index)
-        case .endOfScope(")"):
-            return isAttribute(at: index)
-        case .number, .endOfScope("}"), .endOfScope(">"):
-            return false
-        default:
-            return nil
-        }
-    }
-}

Tests/Rules/SpaceAroundBracketsTests.swift

@@ -188,4 +188,51 @@ final class SpaceAroundBracketsTests: XCTestCase {
         """
         testFormatting(for: input, rule: .spaceAroundBrackets)
     }
+
+    func testAddSpaceBetweenBracketAndAwait() {
+        let input = """
+        let foo = await[bar: 5]
+        """
+        let output = """
+        let foo = await [bar: 5]
+        """
+        testFormatting(for: input, output, rule: .spaceAroundBrackets)
+    }
+
+    func testAddSpaceBetweenParenAndAwaitForSwift5_5() {
+        let input = """
+        let foo = await[bar: 5]
+        """
+        let output = """
+        let foo = await [bar: 5]
+        """
+        testFormatting(for: input, output, rule: .spaceAroundBrackets,
+                       options: FormatOptions(swiftVersion: "5.5"))
+    }
+
+    func testNoAddSpaceBetweenParenAndAwaitForSwiftLessThan5_5() {
+        let input = """
+        let foo = await[bar: 5]
+        """
+        testFormatting(for: input, rule: .spaceAroundBrackets,
+                       options: FormatOptions(swiftVersion: "5.4.9"))
+    }
+
+    func testAddSpaceBetweenParenAndUnsafe() {
+        let input = """
+        unsafe[kinfo_proc](repeating: kinfo_proc(), count: length / MemoryLayout<kinfo_proc>.stride)
+        """
+        let output = """
+        unsafe [kinfo_proc](repeating: kinfo_proc(), count: length / MemoryLayout<kinfo_proc>.stride)
+        """
+        testFormatting(for: input, output, rule: .spaceAroundBrackets)
+    }
+
+    func testNoAddSpaceBetweenParenAndAwaitForSwiftLessThan6_2() {
+        let input = """
+        unsafe[kinfo_proc](repeating: kinfo_proc(), count: length / MemoryLayout<kinfo_proc>.stride)
+        """
+        testFormatting(for: input, rule: .spaceAroundBrackets,
+                       options: FormatOptions(swiftVersion: "6.1"))
+    }
 }