Fix spurious return removal

Author: nicklockwood

Sources/ParsingHelpers.swift

@@ -557,8 +557,7 @@ extension Formatter {
                 default:
                     return true
                 }
-            case .operator("->", .infix), .keyword("init"),
-                 .keyword("subscript"):
+            case .operator("->", .infix), .keyword("init"), .keyword("subscript"), .keyword("throws"):
                 return false
             case .endOfScope(">"):
                 guard let startIndex = index(of: .startOfScope("<"), before: prev) else {
@@ -721,7 +720,8 @@ extension Formatter {
             guard let nextToken = next(.nonSpaceOrComment, after: braceIndex),
                   !nextToken.isOperator(ofType: .infix),
                   !nextToken.isOperator(ofType: .postfix),
-                  nextToken != .startOfScope("(")
+                  nextToken != .startOfScope("("),
+                  nextToken != .startOfScope("{")
             else {
                 return isAfterBrace(index, braceIndex)
             }

Sources/Rules.swift

@@ -3283,12 +3283,13 @@ public struct _FormatRules {
         formatter.forEach(.startOfScope("{")) { startOfScopeIndex, _ in
             // Closures always supported implicit returns, but other types of scopes
             // only support implicit return in Swift 5.1+ (SE-0255)
-            if formatter.options.swiftVersion < "5.1", !formatter.isStartOfClosure(at: startOfScopeIndex) {
+            let isClosure = formatter.isStartOfClosure(at: startOfScopeIndex)
+            if formatter.options.swiftVersion < "5.1", !isClosure {
                 return
             }
 
             // Make sure this is a type of scope that supports implicit returns
-            if formatter.isConditionalStatement(at: startOfScopeIndex) ||
+            if !isClosure, formatter.isConditionalStatement(at: startOfScopeIndex) ||
                 ["do", "else", "catch"].contains(formatter.lastSignificantKeyword(at: startOfScopeIndex, excluding: ["throws"]))
             {
                 return
@@ -3304,7 +3305,7 @@ public struct _FormatRules {
             }
 
             // Make sure we aren't in a failable `init?`, where explicit return is required
-            if let lastSignificantKeywordIndex = formatter.indexOfLastSignificantKeyword(at: startOfScopeIndex),
+            if !isClosure, let lastSignificantKeywordIndex = formatter.indexOfLastSignificantKeyword(at: startOfScopeIndex),
                formatter.tokens[lastSignificantKeywordIndex] == .keyword("init"),
                let nextToken = formatter.next(.nonSpaceOrCommentOrLinebreak, after: lastSignificantKeywordIndex),
                nextToken == .operator("?", .postfix)

Tests/ParsingHelpersTests.swift

@@ -126,6 +126,14 @@ class ParsingHelpersTests: XCTestCase {
         XCTAssertFalse(formatter.isStartOfClosure(at: 18))
     }
 
+    func testClosureInIfCondition() {
+        let formatter = Formatter(tokenize("""
+        if let btn = btns.first { !$0.isHidden } {}
+        """))
+        XCTAssertTrue(formatter.isStartOfClosure(at: 12))
+        XCTAssertFalse(formatter.isStartOfClosure(at: 21))
+    }
+
     // functions
 
     func testFunctionBracesNotTreatedAsClosure() {
@@ -574,6 +582,14 @@ class ParsingHelpersTests: XCTestCase {
         XCTAssert(formatter.isStartOfClosure(at: 5))
     }
 
+    func testBraceAfterTypedThrows() {
+        let formatter = Formatter(tokenize("""
+        do throws(Foo) {} catch {}
+        """))
+        XCTAssertFalse(formatter.isStartOfClosure(at: 7))
+        XCTAssertFalse(formatter.isStartOfClosure(at: 12))
+    }
+
     // MARK: isAccessorKeyword
 
     func testDidSet() {
@@ -1650,6 +1666,14 @@ class ParsingHelpersTests: XCTestCase {
         }
     }
 
+    func testStartOfConditionalStatementConditionContainingUnParenthesizedClosure() {
+        let formatter = Formatter(tokenize("""
+        if let btn = btns.first { !$0.isHidden } {}
+        """))
+        XCTAssertEqual(formatter.startOfConditionalStatement(at: 12), 0)
+        XCTAssertEqual(formatter.startOfConditionalStatement(at: 21), 0)
+    }
+
     // MARK: isStartOfStatement
 
     func testAsyncAfterFuncNotTreatedAsStartOfStatement() {

Tests/RulesTests+Redundancy.swift

@@ -3008,6 +3008,43 @@ class RedundancyTests: RulesTests {
                        options: FormatOptions(swiftVersion: "5.1"), exclude: ["redundantProperty"])
     }
 
+    func testNoRemoveRequiredReturnInIfClosure() {
+        let input = """
+        func findButton() -> Button? {
+            let btns = [top, content, bottom]
+            if let btn = btns.first { !$0.isHidden && $0.alpha > 0.01 } {
+                return btn
+            }
+            return btns.first
+        }
+        """
+        let options = FormatOptions(swiftVersion: "5.1")
+        testFormatting(for: input, rule: FormatRules.redundantReturn, options: options)
+    }
+
+    func testRemoveRedundantReturnInIfClosure() {
+        let input = """
+        func findButton() -> Button? {
+            let btns = [top, content, bottom]
+            if let btn = btns.first { return !$0.isHidden && $0.alpha > 0.01 } {
+                print("hello")
+            }
+            return btns.first
+        }
+        """
+        let output = """
+        func findButton() -> Button? {
+            let btns = [top, content, bottom]
+            if let btn = btns.first { !$0.isHidden && $0.alpha > 0.01 } {
+                print("hello")
+            }
+            return btns.first
+        }
+        """
+        let options = FormatOptions(swiftVersion: "5.1")
+        testFormatting(for: input, output, rule: FormatRules.redundantReturn, options: options)
+    }
+
     func testDisableNextRedundantReturn() {
         let input = """
         func foo() -> Foo {