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

Author: xedin

Sources/SwiftFormat/Core/Pipelines+Generated.swift

@@ -281,6 +281,7 @@ class LintPipeline: SyntaxVisitor {
   override func visit(_ node: SubscriptDeclSyntax) -> SyntaxVisitorContinueKind {
     visitIfEnabled(AllPublicDeclarationsHaveDocumentation.visit, for: node)
     visitIfEnabled(BeginDocumentationCommentWithOneLineSummary.visit, for: node)
+    visitIfEnabled(OmitReturns.visit, for: node)
     visitIfEnabled(UseTripleSlashForDocumentationComments.visit, for: node)
     return .visitChildren
   }

Sources/SwiftFormat/Rules/OmitReturns.swift

@@ -37,6 +37,69 @@ public final class OmitReturns: SyntaxFormatRule {
     return decl
   }
 
+  public override func visit(_ node: SubscriptDeclSyntax) -> DeclSyntax {
+    let decl = super.visit(node)
+
+    guard var `subscript` = decl.as(SubscriptDeclSyntax.self) else {
+      return decl
+    }
+
+    if let accessorBlock = `subscript`.accessorBlock {
+      // 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 decl
+        }
+
+        guard let body = getter.body,
+              let `return` = containsSingleReturn(body.statements) else {
+          return decl
+        }
+
+        guard let getterAt = accessors.firstIndex(where: {
+          $0.accessorSpecifier.tokenKind == .keyword(.get)
+        }) else {
+          return decl
+        }
+
+        getter.body?.statements = unwrapReturnStmt(`return`)
+
+        `subscript`.accessorBlock = .init(
+            leadingTrivia: accessorBlock.leadingTrivia,
+            leftBrace: accessorBlock.leftBrace,
+            accessors: .accessors(accessors.with(\.[getterAt], getter)),
+            rightBrace: accessorBlock.rightBrace,
+            trailingTrivia: accessorBlock.trailingTrivia)
+
+        diagnose(.omitReturnStatement, on: `return`, severity: .refactoring)
+
+        return DeclSyntax(`subscript`)
+
+      case .getter(let getter):
+        guard let `return` = containsSingleReturn(getter) else {
+          return decl
+        }
+        
+        `subscript`.accessorBlock = .init(
+            leadingTrivia: accessorBlock.leadingTrivia,
+            leftBrace: accessorBlock.leftBrace,
+            accessors: .getter(unwrapReturnStmt(`return`)),
+            rightBrace: accessorBlock.rightBrace,
+            trailingTrivia: accessorBlock.trailingTrivia)
+
+        diagnose(.omitReturnStatement, on: `return`, severity: .refactoring)
+
+        return DeclSyntax(`subscript`)
+      }
+    }
+
+    return decl
+  }
+
   public override func visit(_ node: ClosureExprSyntax) -> ExprSyntax {
     let expr = super.visit(node)
 

Tests/SwiftFormatTests/Rules/OmitReturnsTests.swift

@@ -30,4 +30,46 @@ final class OmitReturnsTests: LintOrFormatRuleTestCase {
         }
         """)
   }
+
+  func testOmitReturnInSubscript() {
+      XCTAssertFormatting(
+        OmitReturns.self,
+        input: """
+          struct Test {
+            subscript(x: Int) -> Bool {
+              return false
+            }
+          }
+          """,
+        expected: """
+          struct Test {
+            subscript(x: Int) -> Bool {
+              false
+            }
+          }
+          """)
+
+      XCTAssertFormatting(
+        OmitReturns.self,
+        input: """
+          struct Test {
+            subscript(x: Int) -> Bool {
+              get {
+                return false
+              }
+              set { }
+            }
+          }
+          """,
+        expected: """
+          struct Test {
+            subscript(x: Int) -> Bool {
+              get {
+                false
+              }
+              set { }
+            }
+          }
+          """)
+  }
 }