From 5409ed25eab4bdc76a379f9684b228b20cf49312 Mon Sep 17 00:00:00 2001
From: Nick Lockwood <nick@charcoaldesign.co.uk>
Date: Tue, 4 Oct 2016 18:29:29 +0100
Subject: [PATCH] Fixed spacing between block capture list and arguments

---
 SwiftFormat/Formatter.swift       |  2 +-
 SwiftFormat/Rules.swift           | 33 +++++++++++++++++++++++++++++--
 SwiftFormatTests/RulesTests.swift | 14 +++++++++++++
 3 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/SwiftFormat/Formatter.swift b/SwiftFormat/Formatter.swift
index f116ab601..77d090484 100644
--- a/SwiftFormat/Formatter.swift
+++ b/SwiftFormat/Formatter.swift
@@ -189,7 +189,7 @@ public class Formatter {
                 if token.type == .linebreak {
                     i -= 1
                 }
-            } else if matching(token) {
+            } else if scopeStack.count == 0 && matching(token) {
                 return i
             } else if token.type == .startOfScope {
                 scopeStack.append(token)
diff --git a/SwiftFormat/Rules.swift b/SwiftFormat/Rules.swift
index fe5594bce..565dd2c75 100644
--- a/SwiftFormat/Rules.swift
+++ b/SwiftFormat/Rules.swift
@@ -69,17 +69,46 @@ public func spaceAroundParens(_ formatter: Formatter) {
         }
     }
 
+    func isCaptureList(atIndex i: Int) -> Bool {
+        assert(formatter.tokens[i].string == "]")
+        guard let previousToken = formatter.previousToken(fromIndex: i + 1, matching: {
+            return !$0.isWhitespaceOrCommentOrLinebreak && ($0.type != .endOfScope || $0.string != "]")
+        }), previousToken.type == .startOfScope, previousToken.string == "{" else {
+            return false
+        }
+        guard let nextToken = formatter.nextToken(fromIndex: i, matching: {
+            return !$0.isWhitespaceOrCommentOrLinebreak && ($0.type != .startOfScope || $0.string != "(")
+        }), nextToken.type == .identifier, nextToken.string == "in" else {
+            return false
+        }
+        return true
+    }
+    
     formatter.forEachToken("(") { i, token in
         guard let previousToken = formatter.tokenAtIndex(i - 1) else {
             return
         }
         if previousToken.type == .identifier && spaceAfter(previousToken.string, index: i - 1) {
             formatter.insertToken(Token(.whitespace, " "), atIndex: i)
+        } else if previousToken.type == .endOfScope && previousToken.string == "]" {
+            if isCaptureList(atIndex: i - 1) {
+                formatter.insertToken(Token(.whitespace, " "), atIndex: i)
+            }
         } else if previousToken.type == .whitespace {
             if let token = formatter.tokenAtIndex(i - 2) {
-                if (token.type == .endOfScope && ["]", "}", ")", ">"].contains(token.string)) ||
-                    (token.type == .identifier && !spaceAfter(token.string, index: i - 2)) {
+                if token.type == .identifier && !spaceAfter(token.string, index: i - 2) {
                     formatter.removeTokenAtIndex(i - 1)
+                } else if token.type == .endOfScope {
+                    switch token.string {
+                    case "}", ")", ">":
+                        formatter.removeTokenAtIndex(i - 1)
+                    case "]":
+                        if !isCaptureList(atIndex: i - 2) {
+                            formatter.removeTokenAtIndex(i - 1)
+                        }
+                    default:
+                        break
+                    }
                 }
             }
         }
diff --git a/SwiftFormatTests/RulesTests.swift b/SwiftFormatTests/RulesTests.swift
index f9cc46dea..e4804808b 100644
--- a/SwiftFormatTests/RulesTests.swift
+++ b/SwiftFormatTests/RulesTests.swift
@@ -112,6 +112,20 @@ class RulesTests: XCTestCase {
         XCTAssertEqual(try! format(input, rules: [spaceAroundParens]), output)
         XCTAssertEqual(try! format(input + "\n", rules: defaultRules), output + "\n")
     }
+    
+    func testNoRemoveSpaceBetweenCaptureListAndArguments() {
+        let input = "{ [weak self] (foo) in }"
+        let output = "{ [weak self] (foo) in }"
+        XCTAssertEqual(try! format(input, rules: [spaceAroundParens]), output)
+        XCTAssertEqual(try! format(input + "\n", rules: defaultRules), output + "\n")
+    }
+    
+    func testAddSpaceBetweenCaptureListAndArguments() {
+        let input = "{ [weak self](foo) in }"
+        let output = "{ [weak self] (foo) in }"
+        XCTAssertEqual(try! format(input, rules: [spaceAroundParens]), output)
+        XCTAssertEqual(try! format(input + "\n", rules: defaultRules), output + "\n")
+    }
 
     func testSpaceBetweenClosingParenAndOpenBrace() {
         let input = "func foo(){ foo }"