Fix another case of spurious return removal

nicklockwood · Jul 20, 2024 · 310bb19 · 310bb19
1 parent e4965d9
commit 310bb19
Show file tree

Hide file tree

Showing 3 changed files with 77 additions and 39 deletions.
diff --git a/Sources/ParsingHelpers.swift b/Sources/ParsingHelpers.swift
@@ -707,27 +707,6 @@ extension Formatter {
             return nil
         }
 
-        func isAfterBrace(_ index: Int, _ i: Int) -> Bool {
-            if let scopeStart = lastIndex(of: .startOfScope, in: index ..< i) {
-                return isAfterBrace(index, scopeStart)
-            }
-            guard let braceIndex = lastIndex(
-                of: .endOfScope("}"),
-                in: index ..< i
-            ) else {
-                return false
-            }
-            guard let nextToken = next(.nonSpaceOrComment, after: braceIndex),
-                  !nextToken.isOperator(ofType: .infix),
-                  !nextToken.isOperator(ofType: .postfix),
-                  nextToken != .startOfScope("("),
-                  nextToken != .startOfScope("{")
-            else {
-                return isAfterBrace(index, braceIndex)
-            }
-            return true
-        }
-
         if tokens[index] == .keyword("case"), let i = self.index(
             of: .nonSpaceOrCommentOrLinebreak,
             before: index,
@@ -746,16 +725,13 @@ extension Formatter {
             switch tokens[prevIndex] {
             case let .keyword(name) where
                 ["if", "guard", "while", "for", "case", "catch"].contains(name):
-                fallthrough
+                return prevIndex
             case .delimiter(","):
-                return isAfterBrace(prevIndex, i) ? nil : prevIndex
+                return startOfConditionalStatement(at: prevIndex)
             default:
                 return nil
             }
         case "if", "guard", "while", "for", "case", "where", "switch":
-            if isAfterBrace(index, i) {
-                return nil
-            }
             return index
         default:
             return nil
@@ -778,25 +754,39 @@ extension Formatter {
         else {
             return nil
         }
+
+        func isAfterBrace(_ index: Int, _ i: Int) -> Bool {
+            if let scopeStart = lastIndex(of: .startOfScope, in: index ..< i) {
+                return isAfterBrace(index, scopeStart)
+            }
+            guard let braceIndex = lastIndex(
+                of: .endOfScope("}"),
+                in: index ..< i
+            ) else {
+                return false
+            }
+            guard let nextToken = next(.nonSpaceOrComment, after: braceIndex),
+                  !nextToken.isOperator(ofType: .infix),
+                  !nextToken.isOperator(ofType: .postfix),
+                  nextToken != .startOfScope("("),
+                  nextToken != .startOfScope("{")
+            else {
+                return isAfterBrace(index, braceIndex)
+            }
+            return true
+        }
+
+        if isAfterBrace(index, i) {
+            return nil
+        }
+
         switch keyword {
         case let name where name.hasPrefix("#") || excluding.contains(name):
             fallthrough
         case "in", "is", "as", "try", "await":
             return indexOfLastSignificantKeyword(at: index - 1, excluding: excluding)
         default:
-            guard let braceIndex = self.index(of: .startOfScope("{"), in: index ..< i),
-                  let endIndex = endOfScope(at: braceIndex),
-                  next(.nonSpaceOrComment, after: endIndex) != .startOfScope("(")
-            else {
-                return index
-            }
-            if keyword == "if" || ["var", "let"].contains(keyword) &&
-                last(.nonSpaceOrCommentOrLinebreak, before: index) == .keyword("if"),
-                self.index(of: .startOfScope("{"), in: endIndex ..< i) == nil
-            {
-                return index
-            }
-            return nil
+            return index
         }
     }
 

diff --git a/Tests/ParsingHelpersTests.swift b/Tests/ParsingHelpersTests.swift
@@ -134,6 +134,14 @@ class ParsingHelpersTests: XCTestCase {
         XCTAssertFalse(formatter.isStartOfClosure(at: 21))
     }
 
+    func testClosureInIfCondition2() {
+        let formatter = Formatter(tokenize("""
+        if let foo, let btn = btns.first { !$0.isHidden } {}
+        """))
+        XCTAssertTrue(formatter.isStartOfClosure(at: 17))
+        XCTAssertFalse(formatter.isStartOfClosure(at: 26))
+    }
+
     // functions
 
     func testFunctionBracesNotTreatedAsClosure() {
@@ -473,6 +481,14 @@ class ParsingHelpersTests: XCTestCase {
         XCTAssertFalse(formatter.isStartOfClosure(at: 22))
     }
 
+    func testClosureInsideIfCondition3() {
+        let formatter = Formatter(tokenize("""
+        if baz, let foo = bar(), { x == y }() {}
+        """))
+        XCTAssert(formatter.isStartOfClosure(at: 16))
+        XCTAssertFalse(formatter.isStartOfClosure(at: 28))
+    }
+
     func testClosureAfterGenericType() {
         let formatter = Formatter(tokenize("let foo = Foo<String> {}"))
         XCTAssert(formatter.isStartOfClosure(at: 11))
@@ -590,6 +606,24 @@ class ParsingHelpersTests: XCTestCase {
         XCTAssertFalse(formatter.isStartOfClosure(at: 12))
     }
 
+    // MARK: isConditionalStatement
+
+    func testIfConditionContainingClosure() {
+        let formatter = Formatter(tokenize("""
+        if let btn = btns.first { !$0.isHidden } {}
+        """))
+        XCTAssertTrue(formatter.isConditionalStatement(at: 12))
+        XCTAssertTrue(formatter.isConditionalStatement(at: 21))
+    }
+
+    func testIfConditionContainingClosure2() {
+        let formatter = Formatter(tokenize("""
+        if let foo, let btn = btns.first { !$0.isHidden } {}
+        """))
+        XCTAssertTrue(formatter.isConditionalStatement(at: 17))
+        XCTAssertTrue(formatter.isConditionalStatement(at: 26))
+    }
+
     // MARK: isAccessorKeyword
 
     func testDidSet() {

diff --git a/Tests/RulesTests+Redundancy.swift b/Tests/RulesTests+Redundancy.swift
@@ -3022,6 +3022,20 @@ class RedundancyTests: RulesTests {
         testFormatting(for: input, rule: FormatRules.redundantReturn, options: options)
     }
 
+    func testNoRemoveRequiredReturnInIfClosure2() {
+        let input = """
+        func findButton() -> Button? {
+            let btns = [top, content, bottom]
+            if let foo, 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? {