Avoid dropping inline comments in control_statement fixes

CHANGELOG.md

@@ -42,6 +42,11 @@
 
 ### Bug Fixes
 
+* Preserve inline comments when `control_statement` removes parentheses during
+  `swiftlint --fix`.
+  [theamodhshetty](https://github.com/theamodhshetty)
+  [#6207](https://github.com/realm/SwiftLint/issues/6207)
+
 * Add an `ignore_attributes` option to `implicit_optional_initialization` so
   wrappers/attributes that require explicit `= nil` can be excluded from
   style checks for both `style: always` and `style: never`.  

Source/SwiftLintBuiltInRules/Rules/Style/ControlStatementRule.swift

@@ -64,6 +64,15 @@ struct ControlStatementRule: Rule {
             Example("do {} ↓catch(let error as NSError) {}"): Example("do {} catch let error as NSError {}"),
             Example("↓if (max(a, b) < c) {}"): Example("if max(a, b) < c {}"),
             Example("""
+            ↓if (abc == 1/* && cdf == 2*/) {
+                print("Hello, world!")
+            }
+            """): Example("""
+                if abc == 1/* && cdf == 2*/ {
+                    print("Hello, world!")
+                }
+                """),
+            Example("""
             if (a),
                ( b == 1 ) {}
             """): Example("""
@@ -111,66 +120,105 @@ private extension ControlStatementRule {
 
     final class Rewriter: ViolationsSyntaxRewriter<ConfigurationType> {
         override func visit(_ node: CatchClauseSyntax) -> CatchClauseSyntax {
+            let node = node.with(\.body, super.visit(node.body))
             guard case let items = node.catchItems, items.containSuperfluousParens == true else {
-                return super.visit(node)
+                return node
             }
             numberOfCorrections += 1
-            let node = node
+            return node
                 .with(\.catchKeyword, node.catchKeyword.with(\.trailingTrivia, .space))
                 .with(\.catchItems, items.withoutParens)
-            return super.visit(node)
+                .with(\.body, node.body.withLeftBraceLeadingTrivia(items.trailingCommentTrivia))
         }
 
         override func visit(_ node: GuardStmtSyntax) -> StmtSyntax {
+            let node = node.with(\.body, super.visit(node.body))
             guard node.conditions.containSuperfluousParens else {
-                return super.visit(node)
+                return StmtSyntax(node)
             }
             numberOfCorrections += 1
-            let node = node
-                .with(\.guardKeyword, node.guardKeyword.with(\.trailingTrivia, .space))
-                .with(\.conditions, node.conditions.withoutParens)
-            return super.visit(node)
+            let elseKeyword = node.elseKeyword
+                .with(\.leadingTrivia, node.conditions.trailingCommentTrivia + node.elseKeyword.leadingTrivia)
+            return StmtSyntax(
+                node
+                    .with(\.guardKeyword, node.guardKeyword.with(\.trailingTrivia, .space))
+                    .with(\.conditions, node.conditions.withoutParens)
+                    .with(\.elseKeyword, elseKeyword)
+            )
         }
 
         override func visit(_ node: IfExprSyntax) -> ExprSyntax {
+            let node = node
+                .with(\.body, super.visit(node.body))
+                .with(\.elseBody, node.elseBody.map(rewriteElseBody))
             guard node.conditions.containSuperfluousParens else {
-                return super.visit(node)
+                return ExprSyntax(node)
             }
             numberOfCorrections += 1
-            let node = node
-                .with(\.ifKeyword, node.ifKeyword.with(\.trailingTrivia, .space))
-                .with(\.conditions, node.conditions.withoutParens)
-            return super.visit(node)
+            return ExprSyntax(
+                node
+                    .with(\.ifKeyword, node.ifKeyword.with(\.trailingTrivia, .space))
+                    .with(\.conditions, node.conditions.withoutParens)
+                    .with(\.body, node.body.withLeftBraceLeadingTrivia(node.conditions.trailingCommentTrivia))
+            )
         }
 
         override func visit(_ node: SwitchExprSyntax) -> ExprSyntax {
+            let node = node.with(\.cases, super.visit(node.cases))
             guard let tupleElement = node.subject.unwrapped else {
-                return super.visit(node)
+                return ExprSyntax(node)
             }
             numberOfCorrections += 1
-            let node = node
-                .with(\.switchKeyword, node.switchKeyword.with(\.trailingTrivia, .space))
-                .with(\.subject, tupleElement.with(\.trailingTrivia, .space))
-            return super.visit(node)
+            return ExprSyntax(
+                node
+                    .with(\.switchKeyword, node.switchKeyword.with(\.trailingTrivia, .space))
+                    .with(\.subject, tupleElement.with(\.trailingTrivia, tupleElement.trailingTrivia + .space))
+            )
         }
 
         override func visit(_ node: WhileStmtSyntax) -> StmtSyntax {
+            let node = node.with(\.body, super.visit(node.body))
             guard node.conditions.containSuperfluousParens else {
-                return super.visit(node)
+                return StmtSyntax(node)
             }
             numberOfCorrections += 1
-            let node = node
-                .with(\.whileKeyword, node.whileKeyword.with(\.trailingTrivia, .space))
-                .with(\.conditions, node.conditions.withoutParens)
-            return super.visit(node)
+            return StmtSyntax(
+                node
+                    .with(\.whileKeyword, node.whileKeyword.with(\.trailingTrivia, .space))
+                    .with(\.conditions, node.conditions.withoutParens)
+                    .with(\.body, node.body.withLeftBraceLeadingTrivia(node.conditions.trailingCommentTrivia))
+            )
+        }
+
+        private func rewriteElseBody(_ elseBody: IfExprSyntax.ElseBody) -> IfExprSyntax.ElseBody {
+            switch elseBody {
+            case .ifExpr(let ifExpr):
+                guard let rewritten = visit(ifExpr).as(IfExprSyntax.self) else {
+                    return elseBody
+                }
+                return .ifExpr(rewritten)
+            case .codeBlock(let codeBlock):
+                return .codeBlock(super.visit(codeBlock))
+            }
         }
     }
 }
 
 private extension ExprSyntax {
     var unwrapped: ExprSyntax? {
-        if let expr = `as`(TupleExprSyntax.self)?.elements.onlyElement?.expression {
-            return containsTrailingClosure(Syntax(expr)) ? nil : expr
+        if let tuple = `as`(TupleExprSyntax.self),
+           let element = tuple.elements.onlyElement {
+            var unwrapped = element.expression
+            if containsTrailingClosure(Syntax(unwrapped)) {
+                return nil
+            }
+            if tuple.leftParen.trailingTrivia.isNotEmpty {
+                unwrapped = unwrapped.with(\.leadingTrivia, tuple.leftParen.trailingTrivia + unwrapped.leadingTrivia)
+            }
+            if tuple.rightParen.leadingTrivia.containsComments {
+                unwrapped = unwrapped.with(\.trailingTrivia, unwrapped.trailingTrivia + tuple.rightParen.leadingTrivia)
+            }
+            return unwrapped
         }
         return nil
     }
@@ -199,16 +247,38 @@ private extension ConditionElementListSyntax {
     var withoutParens: Self {
         let conditions = map { (element: ConditionElementSyntax) -> ConditionElementSyntax in
             if let expression = element.condition.as(ExprSyntax.self)?.unwrapped {
-                return element
-                    .with(\.condition, .expression(expression))
+                let commentTrivia = expression.trailingTrivia.containsComments ? expression.trailingTrivia : []
+                let trailingTrivia = if commentTrivia.isNotEmpty, element.trailingComma == nil {
+                    Trivia()
+                } else {
+                    element.trailingTrivia
+                }
+                var updated = element
+                    .with(\.condition, .expression(expression.with(\.trailingTrivia, [])))
                     .with(\.leadingTrivia, element.leadingTrivia)
-                    .with(\.trailingTrivia, element.trailingTrivia)
+                    .with(\.trailingTrivia, trailingTrivia)
+                if let trailingComma = element.trailingComma, commentTrivia.isNotEmpty {
+                    updated = updated.with(
+                        \.trailingComma,
+                        trailingComma.with(\.leadingTrivia, commentTrivia + trailingComma.leadingTrivia)
+                    )
+                }
+                return updated
             }
             return element
         }
         return Self(conditions)
             .with(\.leadingTrivia, leadingTrivia)
-            .with(\.trailingTrivia, trailingTrivia)
+            .with(\.trailingTrivia, trailingCommentTrivia.isNotEmpty ? [] : trailingTrivia)
+    }
+
+    var trailingCommentTrivia: Trivia {
+        guard let last,
+              let expression = last.condition.as(ExprSyntax.self)?.unwrapped,
+              expression.trailingTrivia.containsComments else {
+            return []
+        }
+        return expression.trailingTrivia + last.trailingTrivia
     }
 }
 
@@ -220,16 +290,29 @@ private extension CatchItemListSyntax {
     var withoutParens: Self {
         let items = map { (item: CatchItemSyntax) -> CatchItemSyntax in
             if let expression = item.unwrapped {
+                let commentTrivia = expression.trailingTrivia.containsComments ? expression.trailingTrivia : []
+                let pattern = PatternSyntax(
+                    ExpressionPatternSyntax(expression: expression.with(\.trailingTrivia, []))
+                )
                 return item
-                    .with(\.pattern, PatternSyntax(ExpressionPatternSyntax(expression: expression)))
+                    .with(\.pattern, pattern)
                     .with(\.leadingTrivia, item.leadingTrivia)
-                    .with(\.trailingTrivia, item.trailingTrivia)
+                    .with(\.trailingTrivia, commentTrivia + item.trailingTrivia)
             }
             return item
         }
         return Self(items)
             .with(\.leadingTrivia, leadingTrivia)
-            .with(\.trailingTrivia, trailingTrivia)
+            .with(\.trailingTrivia, trailingCommentTrivia.isNotEmpty ? [] : trailingTrivia)
+    }
+
+    var trailingCommentTrivia: Trivia {
+        guard let last,
+              let expression = last.unwrapped,
+              expression.trailingTrivia.containsComments else {
+            return []
+        }
+        return expression.trailingTrivia
     }
 }
 
@@ -238,3 +321,12 @@ private extension CatchItemSyntax {
         pattern?.as(ExpressionPatternSyntax.self)?.expression.unwrapped
     }
 }
+
+private extension CodeBlockSyntax {
+    func withLeftBraceLeadingTrivia(_ trivia: Trivia) -> Self {
+        guard trivia.isNotEmpty else {
+            return self
+        }
+        return with(\.leftBrace, leftBrace.with(\.leadingTrivia, trivia + leftBrace.leadingTrivia))
+    }
+}