[Lint/Format] Extend return omission rule to include closures

Author: xedin

Sources/SwiftFormat/Core/Pipelines+Generated.swift

@@ -60,6 +60,11 @@ class LintPipeline: SyntaxVisitor {
     return .visitChildren
   }
 
+  override func visit(_ node: ClosureExprSyntax) -> SyntaxVisitorContinueKind {
+    visitIfEnabled(OmitReturns.visit, for: node)
+    return .visitChildren
+  }
+
   override func visit(_ node: ClosureParameterSyntax) -> SyntaxVisitorContinueKind {
     visitIfEnabled(NoLeadingUnderscores.visit, for: node)
     return .visitChildren

Sources/SwiftFormat/Rules/OmitReturns.swift

@@ -37,6 +37,20 @@ public final class OmitReturns: SyntaxFormatRule {
     return decl
   }
 
+  public override func visit(_ node: ClosureExprSyntax) -> ExprSyntax {
+    let expr = super.visit(node)
+
+    // test { return ... }
+    if var closure = expr.as(ClosureExprSyntax.self),
+       let `return` = containsSingleReturn(closure.statements) {
+       closure.statements = unwrapReturnStmt(`return`)
+       diagnose(.omitReturnStatement, on: `return`, severity: .refactoring)
+       return ExprSyntax(closure)
+    }
+
+    return expr
+  }
+
   private func containsSingleReturn(_ body: CodeBlockItemListSyntax) -> ReturnStmtSyntax? {
     if let element = body.firstAndOnly?.as(CodeBlockItemSyntax.self),
        let ret = element.item.as(ReturnStmtSyntax.self),

Tests/SwiftFormatTests/Rules/OmitReturnsTests.swift

@@ -15,4 +15,19 @@ final class OmitReturnsTests: LintOrFormatRuleTestCase {
         }
         """)
   }
+
+  func testOmitReturnInClosure() {
+    XCTAssertFormatting(
+      OmitReturns.self,
+      input: """
+        vals.filter {
+          return $0.count == 1
+        }
+        """,
+      expected: """
+        vals.filter {
+          $0.count == 1
+        }
+        """)
+  }
 }