From 6f0b35a002fa3057dd76e845d7af5901cf7e9c7f Mon Sep 17 00:00:00 2001
From: TTOzzi <devttozzi@gmail.com>
Date: Sun, 16 Feb 2025 01:25:22 +0900
Subject: [PATCH] Improve the formatting logic of OmitExplicitReturns rule

---
 .../Rules/OmitExplicitReturns.swift           |  4 +-
 .../Rules/OmitReturnsTests.swift              | 42 +++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/Sources/SwiftFormat/Rules/OmitExplicitReturns.swift b/Sources/SwiftFormat/Rules/OmitExplicitReturns.swift
index 63c65fc2a..d9808732c 100644
--- a/Sources/SwiftFormat/Rules/OmitExplicitReturns.swift
+++ b/Sources/SwiftFormat/Rules/OmitExplicitReturns.swift
@@ -60,7 +60,7 @@ public final class OmitExplicitReturns: SyntaxFormatRule {
     guard let accessorBlock = binding.accessorBlock,
       let transformed = transformAccessorBlock(accessorBlock)
     else {
-      return node
+      return super.visit(node)
     }
 
     binding.accessorBlock = transformed
@@ -145,7 +145,7 @@ public final class OmitExplicitReturns: SyntaxFormatRule {
     CodeBlockItemListSyntax([
       CodeBlockItemSyntax(
         leadingTrivia: returnStmt.leadingTrivia,
-        item: .expr(returnStmt.expression!),
+        item: .expr(returnStmt.expression!.detached.with(\.trailingTrivia, [])),
         semicolon: nil,
         trailingTrivia: returnStmt.trailingTrivia
       )
diff --git a/Tests/SwiftFormatTests/Rules/OmitReturnsTests.swift b/Tests/SwiftFormatTests/Rules/OmitReturnsTests.swift
index ae0d84188..74e41dbf9 100644
--- a/Tests/SwiftFormatTests/Rules/OmitReturnsTests.swift
+++ b/Tests/SwiftFormatTests/Rules/OmitReturnsTests.swift
@@ -119,4 +119,46 @@ final class OmitReturnsTests: LintOrFormatRuleTestCase {
       ]
     )
   }
+
+  func testInVariableBindings() {
+    assertFormatting(
+      OmitExplicitReturns.self,
+      input: """
+          var f = l.filter { 1️⃣return $0.a != o }
+          var bar = l.filter { 
+            2️⃣return $0.a != o
+          }
+        """,
+      expected: """
+          var f = l.filter { $0.a != o }
+          var bar = l.filter { 
+            $0.a != o
+          }
+        """,
+      findings: [
+        FindingSpec("1️⃣", message: "'return' can be omitted because body consists of a single expression"),
+        FindingSpec("2️⃣", message: "'return' can be omitted because body consists of a single expression"),
+      ]
+    )
+  }
+
+  func testInVariableBindingWithTrailingTrivia() {
+    assertFormatting(
+      OmitExplicitReturns.self,
+      input: """
+          var f = l.filter { 
+            1️⃣return $0.a != o // comment
+          }
+        """,
+      expected: """
+          var f = l.filter { 
+            $0.a != o // comment
+          }
+        """,
+      findings: [
+        FindingSpec("1️⃣", message: "'return' can be omitted because body consists of a single expression")
+      ]
+    )
+  }
+
 }