[Lint/Format] Add a rule to omit return from functions, closures, subscripts, and variables

Sources/SwiftFormat/Core/Pipelines+Generated.swift

@@ -60,6 +60,11 @@ class LintPipeline: SyntaxVisitor {
     return .visitChildren
   }
 
+  override func visit(_ node: ClosureExprSyntax) -> SyntaxVisitorContinueKind {
+    visitIfEnabled(OmitExplicitReturns.visit, for: node)
+    return .visitChildren
+  }
+
   override func visit(_ node: ClosureParameterSyntax) -> SyntaxVisitorContinueKind {
     visitIfEnabled(NoLeadingUnderscores.visit, for: node)
     return .visitChildren
@@ -146,6 +151,7 @@ class LintPipeline: SyntaxVisitor {
     visitIfEnabled(AlwaysUseLowerCamelCase.visit, for: node)
     visitIfEnabled(BeginDocumentationCommentWithOneLineSummary.visit, for: node)
     visitIfEnabled(NoLeadingUnderscores.visit, for: node)
+    visitIfEnabled(OmitExplicitReturns.visit, for: node)
     visitIfEnabled(UseTripleSlashForDocumentationComments.visit, for: node)
     visitIfEnabled(ValidateDocumentationComments.visit, for: node)
     return .visitChildren
@@ -226,6 +232,7 @@ class LintPipeline: SyntaxVisitor {
   }
 
   override func visit(_ node: PatternBindingSyntax) -> SyntaxVisitorContinueKind {
+    visitIfEnabled(OmitExplicitReturns.visit, for: node)
     visitIfEnabled(UseSingleLinePropertyGetter.visit, for: node)
     return .visitChildren
   }
@@ -275,6 +282,7 @@ class LintPipeline: SyntaxVisitor {
   override func visit(_ node: SubscriptDeclSyntax) -> SyntaxVisitorContinueKind {
     visitIfEnabled(AllPublicDeclarationsHaveDocumentation.visit, for: node)
     visitIfEnabled(BeginDocumentationCommentWithOneLineSummary.visit, for: node)
+    visitIfEnabled(OmitExplicitReturns.visit, for: node)
     visitIfEnabled(UseTripleSlashForDocumentationComments.visit, for: node)
     return .visitChildren
   }
@@ -343,6 +351,7 @@ extension FormatPipeline {
     node = NoLabelsInCasePatterns(context: context).rewrite(node)
     node = NoParensAroundConditions(context: context).rewrite(node)
     node = NoVoidReturnOnFunctionSignature(context: context).rewrite(node)
+    node = OmitExplicitReturns(context: context).rewrite(node)
     node = OneCasePerLine(context: context).rewrite(node)
     node = OneVariableDeclarationPerLine(context: context).rewrite(node)
     node = OrderedImports(context: context).rewrite(node)

Sources/SwiftFormat/Core/RuleNameCache+Generated.swift

@@ -37,6 +37,7 @@ public let ruleNameCache: [ObjectIdentifier: String] = [
   ObjectIdentifier(NoLeadingUnderscores.self): "NoLeadingUnderscores",
   ObjectIdentifier(NoParensAroundConditions.self): "NoParensAroundConditions",
   ObjectIdentifier(NoVoidReturnOnFunctionSignature.self): "NoVoidReturnOnFunctionSignature",
+  ObjectIdentifier(OmitExplicitReturns.self): "OmitExplicitReturns",
   ObjectIdentifier(OneCasePerLine.self): "OneCasePerLine",
   ObjectIdentifier(OneVariableDeclarationPerLine.self): "OneVariableDeclarationPerLine",
   ObjectIdentifier(OnlyOneTrailingClosureArgument.self): "OnlyOneTrailingClosureArgument",

Sources/SwiftFormat/Rules/OmitExplicitReturns.swift

@@ -0,0 +1,148 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftSyntax
+
+/// Single-expression functions, closures, subscripts can omit `return` statement.
+///
+/// Lint: `func <name>() { return ... }` and similar single expression constructs will yield a lint error.
+///
+/// Format: `func <name>() { return ... }` constructs will be replaced with
+///         equivalent `func <name>() { ... }` constructs.
+@_spi(Rules)
+public final class OmitExplicitReturns: SyntaxFormatRule {
+  public override class var isOptIn: Bool { return true }
+
+  public override func visit(_ node: FunctionDeclSyntax) -> DeclSyntax {
+    let decl = super.visit(node)
+
+    // func <name>() -> <Type> { return ... }
+    guard var funcDecl = decl.as(FunctionDeclSyntax.self),
+          let body = funcDecl.body,
+          let returnStmt = containsSingleReturn(body.statements) else {
+      return decl
+    }
+
+    funcDecl.body?.statements = rewrapReturnedExpression(returnStmt)
+    diagnose(.omitReturnStatement, on: returnStmt, severity: .refactoring)
+    return DeclSyntax(funcDecl)
+  }
+
+  public override func visit(_ node: SubscriptDeclSyntax) -> DeclSyntax {
+    let decl = super.visit(node)
+
+    guard var subscriptDecl = decl.as(SubscriptDeclSyntax.self),
+          let accessorBlock = subscriptDecl.accessorBlock,
+          // We are assuming valid Swift code here where only
+          // one `get { ... }` is allowed.
+          let transformed = transformAccessorBlock(accessorBlock) else {
+      return decl
+    }
+
+    subscriptDecl.accessorBlock = transformed
+    return DeclSyntax(subscriptDecl)
+  }
+
+  public override func visit(_ node: PatternBindingSyntax) -> PatternBindingSyntax {
+    var binding = node
+
+    guard let accessorBlock = binding.accessorBlock,
+          let transformed = transformAccessorBlock(accessorBlock) else {
+      return node
+    }
+
+    binding.accessorBlock = transformed
+    return binding
+  }
+
+  public override func visit(_ node: ClosureExprSyntax) -> ExprSyntax {
+    let expr = super.visit(node)
+
+    // test { return ... }
+    guard var closureExpr = expr.as(ClosureExprSyntax.self),
+          let returnStmt = containsSingleReturn(closureExpr.statements) else {
+       return expr
+    }
+
+    closureExpr.statements = rewrapReturnedExpression(returnStmt)
+    diagnose(.omitReturnStatement, on: returnStmt, severity: .refactoring)
+    return ExprSyntax(closureExpr)
+  }
+
+  private func transformAccessorBlock(_ accessorBlock: AccessorBlockSyntax) -> AccessorBlockSyntax? {
+    // We are assuming valid Swift code here where only
+    // one `get { ... }` is allowed.
+    switch accessorBlock.accessors {
+    case .accessors(let accessors):
+      guard var getter = accessors.filter({
+        $0.accessorSpecifier.tokenKind == .keyword(.get)
+      }).first else {
+        return nil
+      }
+
+      guard let body = getter.body,
+            let returnStmt = containsSingleReturn(body.statements) else {
+        return nil
+      }
+
+      guard let getterAt = accessors.firstIndex(where: {
+        $0.accessorSpecifier.tokenKind == .keyword(.get)
+      }) else {
+        return nil
+      }
+
+      getter.body?.statements = rewrapReturnedExpression(returnStmt)
+
+      diagnose(.omitReturnStatement, on: returnStmt, severity: .refactoring)
+
+      var newBlock = accessorBlock
+      newBlock.accessors = .accessors(accessors.with(\.[getterAt], getter))
+      return newBlock
+
+    case .getter(let getter):
+      guard let returnStmt = containsSingleReturn(getter) else {
+        return nil
+      }
+
+      diagnose(.omitReturnStatement, on: returnStmt, severity: .refactoring)
+
+      var newBlock = accessorBlock
+      newBlock.accessors = .getter(rewrapReturnedExpression(returnStmt))
+      return newBlock
+    }
+  }
+
+  private func containsSingleReturn(_ body: CodeBlockItemListSyntax) -> ReturnStmtSyntax? {
+    guard let element = body.firstAndOnly?.as(CodeBlockItemSyntax.self),
+       let returnStmt = element.item.as(ReturnStmtSyntax.self) else
+        {
+      return nil
+    }
+
+    return !returnStmt.children(viewMode: .all).isEmpty && returnStmt.expression != nil ? returnStmt : nil
+  }
+
+  private func rewrapReturnedExpression(_ returnStmt: ReturnStmtSyntax) -> CodeBlockItemListSyntax {
+    CodeBlockItemListSyntax([
+      CodeBlockItemSyntax(
+        leadingTrivia: returnStmt.leadingTrivia,
+        item: .expr(returnStmt.expression!),
+        semicolon: nil,
+        trailingTrivia: returnStmt.trailingTrivia)
+    ])
+  }
+}
+
+extension Finding.Message {
+  public static let omitReturnStatement: Finding.Message =
+    "'return' can be omitted because body consists of a single expression"
+}

Sources/SwiftFormatConfiguration/RuleRegistry+Generated.swift

@@ -36,6 +36,7 @@ enum RuleRegistry {
     "NoLeadingUnderscores": false,
     "NoParensAroundConditions": true,
     "NoVoidReturnOnFunctionSignature": true,
+    "OmitExplicitReturns": false,
     "OneCasePerLine": true,
     "OneVariableDeclarationPerLine": true,
     "OnlyOneTrailingClosureArgument": true,

Tests/SwiftFormatTests/Rules/OmitReturnsTests.swift

@@ -0,0 +1,119 @@
+import _SwiftFormatTestSupport
+
+@_spi(Rules) import SwiftFormat
+
+final class OmitReturnsTests: LintOrFormatRuleTestCase {
+  func testOmitReturnInFunction() {
+    assertFormatting(
+      OmitExplicitReturns.self,
+      input: """
+        func test() -> Bool {
+          1️⃣return false
+        }
+      """,
+      expected: """
+        func test() -> Bool {
+          false
+        }
+      """,
+      findings: [
+        FindingSpec("1️⃣", message: "'return' can be omitted because body consists of a single expression")
+      ])
+  }
+
+  func testOmitReturnInClosure() {
+    assertFormatting(
+      OmitExplicitReturns.self,
+      input: """
+        vals.filter {
+          1️⃣return $0.count == 1
+        }
+      """,
+      expected: """
+        vals.filter {
+          $0.count == 1
+        }
+      """,
+      findings: [
+        FindingSpec("1️⃣", message: "'return' can be omitted because body consists of a single expression")
+      ])
+  }
+
+  func testOmitReturnInSubscript() {
+    assertFormatting(
+      OmitExplicitReturns.self,
+      input: """
+        struct Test {
+          subscript(x: Int) -> Bool {
+            1️⃣return false
+          }
+        }
+
+        struct Test {
+          subscript(x: Int) -> Bool {
+            get {
+              2️⃣return false
+            }
+            set { }
+          }
+        }
+      """,
+      expected: """
+        struct Test {
+          subscript(x: Int) -> Bool {
+            false
+          }
+        }
+
+        struct Test {
+          subscript(x: Int) -> Bool {
+            get {
+              false
+            }
+            set { }
+          }
+        }
+      """,
+      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 testOmitReturnInComputedVars() {
+    assertFormatting(
+      OmitExplicitReturns.self,
+      input: """
+        var x: Int {
+          1️⃣return 42
+        }
+
+        struct Test {
+          var x: Int {
+            get {
+              2️⃣return 42
+            }
+            set { }
+          }
+        }
+      """,
+      expected: """
+        var x: Int {
+          42
+        }
+
+        struct Test {
+          var x: Int {
+            get {
+              42
+            }
+            set { }
+          }
+        }
+      """,
+      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")
+      ])
+  }
+}