Improve redundantReturn logic

nicklockwood · Jul 10, 2024 · dd689bd · dd689bd
1 parent 250a876
commit dd689bd
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 10 deletions.
diff --git a/Sources/ParsingHelpers.swift b/Sources/ParsingHelpers.swift
@@ -696,14 +696,14 @@ extension Formatter {
     }
 
     /// Returns true if the token at the specified index is part of a conditional statement
-    func isConditionalStatement(at i: Int) -> Bool {
-        startOfConditionalStatement(at: i) != nil
+    func isConditionalStatement(at i: Int, excluding: Set<String> = []) -> Bool {
+        startOfConditionalStatement(at: i, excluding: excluding) != nil
     }
 
     /// If the token at the specified index is part of a conditional statement, returns the index of the first
     /// token in the statement (e.g. `if`, `guard`, `while`, etc.), otherwise returns nil
-    func startOfConditionalStatement(at i: Int) -> Int? {
-        guard var index = indexOfLastSignificantKeyword(at: i, excluding: ["else"]) else {
+    func startOfConditionalStatement(at i: Int, excluding: Set<String> = []) -> Int? {
+        guard var index = indexOfLastSignificantKeyword(at: i, excluding: excluding.union(["else"])) else {
             return nil
         }
 
@@ -762,7 +762,7 @@ extension Formatter {
         }
     }
 
-    func lastSignificantKeyword(at i: Int, excluding: [String] = []) -> String? {
+    func lastSignificantKeyword(at i: Int, excluding: Set<String> = []) -> String? {
         guard let index = indexOfLastSignificantKeyword(at: i, excluding: excluding),
               case let .keyword(keyword) = tokens[index]
         else {
@@ -771,7 +771,7 @@ extension Formatter {
         return keyword
     }
 
-    func indexOfLastSignificantKeyword(at i: Int, excluding: [String] = []) -> Int? {
+    func indexOfLastSignificantKeyword(at i: Int, excluding: Set<String> = []) -> Int? {
         guard let token = token(at: i),
               let index = token.isKeyword ? i : index(of: .keyword, before: i),
               case let .keyword(keyword) = tokens[index]

diff --git a/Sources/Rules.swift b/Sources/Rules.swift
@@ -3258,8 +3258,10 @@ public struct _FormatRules {
                     return
                 }
             }
-            let endIndex = formatter.endOfScope(at: i)
-            if let endIndex = endIndex, formatter.tokens[i + 1 ..< endIndex].contains(.keyword("return")) {
+            // Don't remove return if it's followed by more code
+            guard let endIndex = formatter.endOfScope(at: i),
+                  formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: i) == endIndex
+            else {
                 return
             }
             if formatter.index(of: .nonSpaceOrLinebreak, after: i) == endIndex,
@@ -3289,7 +3291,7 @@ public struct _FormatRules {
             }
 
             // Make sure this is a type of scope that supports implicit returns
-            if !isClosure, formatter.isConditionalStatement(at: startOfScopeIndex) ||
+            if !isClosure, formatter.isConditionalStatement(at: startOfScopeIndex, excluding: ["where"]) ||
                 ["do", "else", "catch"].contains(formatter.lastSignificantKeyword(at: startOfScopeIndex, excluding: ["throws"]))
             {
                 return
@@ -3304,8 +3306,9 @@ public struct _FormatRules {
                 return
             }
 
-            // Make sure we aren't in a failable `init?`, where explicit return is required
+            // Make sure we aren't in a failable `init?`, where explicit return is required unless it's the only statement
             if !isClosure, let lastSignificantKeywordIndex = formatter.indexOfLastSignificantKeyword(at: startOfScopeIndex),
+               formatter.next(.nonSpaceOrCommentOrLinebreak, after: startOfScopeIndex) != .keyword("return"),
                formatter.tokens[lastSignificantKeywordIndex] == .keyword("init"),
                let nextToken = formatter.next(.nonSpaceOrCommentOrLinebreak, after: lastSignificantKeywordIndex),
                nextToken == .operator("?", .postfix)