Merge pull request swiftlang#154 from google/fix-early-exits

Clean up UseEarlyExits and fix bug sometimes causing code loss.

Author: dabelknap

Sources/SwiftFormatRules/UseEarlyExits.swift

@@ -16,73 +16,95 @@ import SwiftSyntax
 
 /// Early exits should be used whenever possible.
 ///
-/// Practically, this means that `if ... else return/throw/break` constructs should be replaced by
-/// `guard ... else { return/throw/break }` constructs in order to keep indentation levels low.
+/// This means that `if ... else { return/throw/break/continue }` constructs should be replaced by
+/// `guard ... else { return/throw/break/continue }` constructs in order to keep indentation levels
+/// low. Specifically, code of the following form:
 ///
-/// Lint: `if ... else return/throw/break` constructs will yield a lint error.
+/// ```swift
+/// if condition {
+///   trueBlock
+/// } else {
+///   falseBlock
+///   return/throw/break/continue
+/// }
+/// ```
 ///
-/// Format: `if ... else return/throw/break` constructs will be replaced with equivalent
-///         `guard ... else { return/throw/break }` constructs.
-///         TODO(abl): replace implicit guards as well?
+/// will be transformed into:
+///
+/// ```swift
+/// guard condition else {
+///   falseBlock
+///   return/throw/break/continue
+/// }
+/// trueBlock
+/// ```
+///
+/// Lint: `if ... else { return/throw/break/continue }` constructs will yield a lint error.
+///
+/// Format: `if ... else { return/throw/break/continue }` constructs will be replaced with
+///         equivalent `guard ... else { return/throw/break/continue }` constructs.
 ///
 /// - SeeAlso: https://google.github.io/swift#guards-for-early-exits
 public final class UseEarlyExits: SyntaxFormatRule {
 
-  public override func visit(_ node: CodeBlockSyntax) -> Syntax {
-    
-    var newItems: [CodeBlockItemSyntax] = []
-    for statement in node.statements {
-      guard let ifStmt = statement.item as? IfStmtSyntax else { return node }
-      guard let elseStmt = ifStmt.elseBody else { return node }
-      guard let elseBody = elseStmt as? CodeBlockSyntax else { return node }
+  public override func visit(_ node: CodeBlockItemListSyntax) -> Syntax {
+    // Continue recursing down the tree first, so that any nested/child nodes get transformed first.
+    let nodeAfterTransformingChildren = super.visit(node)
+    guard let codeBlockItems = nodeAfterTransformingChildren as? CodeBlockItemListSyntax else {
+      return nodeAfterTransformingChildren
+    }
+
+    return SyntaxFactory.makeCodeBlockItemList(
+      codeBlockItems.flatMap { (codeBlockItem: CodeBlockItemSyntax) -> [CodeBlockItemSyntax] in
+        // The `elseBody` of an `IfStmtSyntax` will be a `CodeBlockSyntax` if it's an `else` block,
+        // or another `IfStmtSyntax` if it's an `else if` block. We only want to handle the former.
+        guard let ifStatement = codeBlockItem.item as? IfStmtSyntax,
+          let elseBody = ifStatement.elseBody as? CodeBlockSyntax,
+          codeBlockEndsWithEarlyExit(elseBody)
+        else {
+          return [codeBlockItem]
+        }
 
-      if elseContainsControlStmt(elseStmt: elseStmt) {
-        diagnose(.useGuardStmt, on: ifStmt)
-        guard let moveDeletedIfCode = visit(
-          ifStmt.body.withLeftBrace(nil).withRightBrace(nil)) as? CodeBlockSyntax else { continue }
-        guard let moveElseBody = visit(elseBody) as? CodeBlockSyntax else { continue }
-        
-        let ifConditions = ifStmt.conditions
-        let formattedGuardKeyword = SyntaxFactory.makeGuardKeyword(
-          leadingTrivia: ifStmt.ifKeyword.leadingTrivia,
+        diagnose(.useGuardStatement, on: ifStatement.elseKeyword)
+
+        let trueBlock = ifStatement.body.withLeftBrace(nil).withRightBrace(nil)
+
+        let guardKeyword = SyntaxFactory.makeGuardKeyword(
+          leadingTrivia: ifStatement.ifKeyword.leadingTrivia,
           trailingTrivia: .spaces(1))
-        let newGuardStmt = SyntaxFactory.makeGuardStmt(
-          guardKeyword: formattedGuardKeyword,
-          conditions: ifConditions,
+        let guardStatement = SyntaxFactory.makeGuardStmt(
+          guardKeyword: guardKeyword,
+          conditions: ifStatement.conditions,
           elseKeyword: SyntaxFactory.makeElseKeyword(trailingTrivia: .spaces(1)),
-          body: moveElseBody)
-        newItems.append(
-          SyntaxFactory.makeCodeBlockItem(item: newGuardStmt,
-                                          semicolon: nil)
-        )
-        newItems.append(
-          SyntaxFactory.makeCodeBlockItem(item: moveDeletedIfCode,
-                                          semicolon: nil)
-        )
-      }
-    }
+          body: elseBody)
 
-    let newNode = node.withStatements(SyntaxFactory.makeCodeBlockItemList(newItems))
-    return super.visit(newNode)
+        return [
+          SyntaxFactory.makeCodeBlockItem(item: guardStatement, semicolon: nil),
+          SyntaxFactory.makeCodeBlockItem(item: trueBlock, semicolon: nil),
+        ]
+      })
   }
 
-  func elseContainsControlStmt(elseStmt: Syntax) -> Bool {
-    for child in elseStmt.children {
-      guard let codeBlockList = child as? CodeBlockItemListSyntax else { continue }
-      guard let last = codeBlockList.child(at: codeBlockList.count - 1) as?
-        CodeBlockItemSyntax else { continue }
+  /// Returns true if the last statement in the given code block is one that will cause an early
+  /// exit from the control flow construct or function.
+  private func codeBlockEndsWithEarlyExit(_ codeBlock: CodeBlockSyntax) -> Bool {
+    let statements = codeBlock.statements
+    guard statements.count > 0 else { return false }
+    let lastStatement = statements[statements.count - 1]
 
-      switch last.item {
-      case is ReturnStmtSyntax, is ThrowStmtSyntax, is BreakStmtSyntax, is ContinueStmtSyntax:
-        return true
-      default:
-        continue
-      }
+    switch lastStatement.item {
+    case is ReturnStmtSyntax, is ThrowStmtSyntax, is BreakStmtSyntax, is ContinueStmtSyntax:
+      return true
+    default:
+      return false
     }
-    return false
   }
 }
 
 extension Diagnostic.Message {
-  static let useGuardStmt = Diagnostic.Message(.warning, "replace with guard statement")
+
+  static let useGuardStatement = Diagnostic.Message(
+    .warning,
+    "replace the `if/else` block with a `guard` statement containing the early exit"
+  )
 }

Sources/swift-format/PopulatePipeline.swift

@@ -201,7 +201,7 @@ func populate(_ pipeline: Pipeline) {
   pipeline.addFormatter(
     UseEarlyExits.self,
     for:
-      CodeBlockSyntax.self
+      CodeBlockItemListSyntax.self
   )
 
   pipeline.addFormatter(

Tests/SwiftFormatRulesTests/UseEarlyExitsTests.swift

@@ -5,68 +5,142 @@ import XCTest
 @testable import SwiftFormatRules
 
 public class UseEarlyExitsTests: DiagnosingTestCase {
-  public func testIfToGuardStmtSwitch() {
-    // The indentation in the expected output is explicitly incorrect because this formatting rule
-    // does not fix it up with the assumption that the pretty-printer will handle it.
+
+  public func testBasicIfElse() {
+    // In this and other tests, the indentation of the true block in the expected output is
+    // explicitly incorrect because this formatting rule does not fix it up with the assumption that
+    // the pretty-printer will handle it.
+    XCTAssertFormatting(
+      UseEarlyExits.self,
+      input: """
+        if condition {
+          trueBlock()
+        } else {
+          falseBlock()
+          return
+        }
+        """,
+      expected: """
+        guard condition else {
+          falseBlock()
+          return
+        }
+          trueBlock()
+        """)
+  }
+
+  public func testIfElseWithBothEarlyExiting() {
+    XCTAssertFormatting(
+      UseEarlyExits.self,
+      input: """
+        if condition {
+          trueBlock()
+          return
+        } else {
+          falseBlock()
+          return
+        }
+        """,
+      expected: """
+        guard condition else {
+          falseBlock()
+          return
+        }
+          trueBlock()
+          return
+        """)
+  }
+
+  public func testElseIfsDoNotChange() {
+    let input = """
+      if condition {
+        trueBlock()
+      } else if otherCondition {
+        otherBlock()
+        return
+      }
+      """
+    XCTAssertFormatting(UseEarlyExits.self, input: input, expected: input)
+  }
+
+  public func testElsesAtEndOfElseIfsDoNotChange() {
+    let input = """
+      if condition {
+        trueBlock()
+      } else if otherCondition {
+        otherBlock()
+        return
+      } else {
+        falseBlock()
+        return
+      }
+      """
+    XCTAssertFormatting(UseEarlyExits.self, input: input, expected: input)
+  }
+
+  public func testComplex() {
     XCTAssertFormatting(
       UseEarlyExits.self,
       input: """
-             func discombobulate(_ values: [Int]) throws -> Int {
+        func discombobulate(_ values: [Int]) throws -> Int {
 
-               // Comment 1
+          // Comment 1
 
-               /*Comment 2*/ if let first = values.first {
-                 // Comment 3
+          /*Comment 2*/ if let first = values.first {
+            // Comment 3
 
-                 /// Doc comment
-                 if first >= 0 {
-                   // Comment 4
-                   var result = 0
-                   for value in values {
-                     result += invertedCombobulatoryFactor(of: value)
-                   }
-                   return result
-                 } else {
-                   print("Can't have negative energy")
-                   throw DiscombobulationError.negativeEnergy
-                 }
-               } else {
-                 print("The array was empty")
-                 throw DiscombobulationError.arrayWasEmpty
-               }
-             }
-             """,
+            /// Doc comment
+            if first >= 0 {
+              // Comment 4
+              var result = 0
+              for value in values {
+                result += invertedCombobulatoryFactor(of: value)
+              }
+              return result
+            } else {
+              print("Can't have negative energy")
+              throw DiscombobulationError.negativeEnergy
+            }
+          } else {
+            print("The array was empty")
+            throw DiscombobulationError.arrayWasEmpty
+          }
+        }
+        """,
       expected: """
-                func discombobulate(_ values: [Int]) throws -> Int {
+        func discombobulate(_ values: [Int]) throws -> Int {
 
-                  // Comment 1
+          // Comment 1
 
-                  /*Comment 2*/ guard let first = values.first else {
-                    print("The array was empty")
-                    throw DiscombobulationError.arrayWasEmpty
-                  }
-                    // Comment 3
+          /*Comment 2*/ guard let first = values.first else {
+            print("The array was empty")
+            throw DiscombobulationError.arrayWasEmpty
+          }
+            // Comment 3
 
-                    /// Doc comment
-                    guard first >= 0 else {
-                      print("Can't have negative energy")
-                      throw DiscombobulationError.negativeEnergy
-                    }
-                      // Comment 4
-                      var result = 0
-                      for value in values {
-                        result += invertedCombobulatoryFactor(of: value)
-                      }
-                      return result
-                }
-                """)
-    
+            /// Doc comment
+            guard first >= 0 else {
+              print("Can't have negative energy")
+              throw DiscombobulationError.negativeEnergy
+            }
+              // Comment 4
+              var result = 0
+              for value in values {
+                result += invertedCombobulatoryFactor(of: value)
+              }
+              return result
+        }
+        """)
   }
 
   #if !os(macOS)
   static let allTests = [
-    UseEarlyExitsTests.testSwitchStmts,
-    ]
+    UseEarlyExitsTests.testBasicIfElse,
+    UseEarlyExitsTests.testIfElseWithBothEarlyExiting,
+    UseEarlyExitsTests.testElseIfsDoNotChange,
+    UseEarlyExitsTests.testElsesAtEndOfElseIfsDoNotChange,
+    UseEarlyExitsTests.testComplex,
+  ]
   #endif
 
 }