Ignore end-of-line comments when dirtying if-with-same-arms branches

crates/ruff/resources/test/fixtures/flake8_simplify/SIM114.py

@@ -101,3 +101,16 @@ def foo(p):
     x = 1
 elif c:
     x = 1
+
+
+def foo():
+    a = True
+    b = False
+    if a > b:  # end-of-line
+        return 3
+    elif a == b:
+        return 3
+    elif a < b:  # end-of-line
+        return 4
+    elif b is None:
+        return 4

crates/ruff/src/rules/flake8_simplify/rules/ast_if.rs

@@ -10,7 +10,7 @@ use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::comparable::{ComparableConstant, ComparableExpr, ComparableStmt};
 use ruff_python_ast::helpers::{any_over_expr, contains_effect, first_colon_range, has_comments};
 use ruff_python_ast::source_code::Locator;
-use ruff_python_ast::stmt_if::if_elif_branches;
+use ruff_python_ast::stmt_if::{if_elif_branches, IfElifBranch};
 use ruff_python_semantic::SemanticModel;
 use ruff_python_trivia::UniversalNewlines;
 
@@ -676,6 +676,15 @@ pub(crate) fn use_ternary_operator(checker: &mut Checker, stmt: &Stmt) {
     checker.diagnostics.push(diagnostic);
 }
 
+/// Return the [`TextRange`] of an [`IfElifBranch`]'s body (from the end of the test to the end of
+/// the body).
+fn body_range(branch: &IfElifBranch, locator: &Locator) -> TextRange {
+    TextRange::new(
+        locator.line_end(branch.test.end()),
+        locator.line_end(branch.range.end()),
+    )
+}
+
 /// SIM114
 pub(crate) fn if_with_same_arms(checker: &mut Checker, locator: &Locator, stmt_if: &StmtIf) {
     let mut branches_iter = if_elif_branches(stmt_if).peekable();
@@ -698,24 +707,19 @@ pub(crate) fn if_with_same_arms(checker: &mut Checker, locator: &Locator, stmt_i
         }
 
         // ...and the same comments
-        let first_comments: Vec<_> = checker
+        let first_comments = checker
             .indexer()
-            .comments_in_range(current_branch.range, locator)
-            .collect();
-        let second_comments: Vec<_> = checker
+            .comments_in_range(body_range(&current_branch, locator), locator);
+        let second_comments = checker
             .indexer()
-            .comments_in_range(following_branch.range, locator)
-            .collect();
-        if first_comments != second_comments {
+            .comments_in_range(body_range(following_branch, locator), locator);
+        if !first_comments.eq(second_comments) {
             continue;
         }
 
         checker.diagnostics.push(Diagnostic::new(
             IfWithSameArms,
-            TextRange::new(
-                current_branch.range.start(),
-                following_branch.body.last().unwrap().end(),
-            ),
+            TextRange::new(current_branch.range.start(), following_branch.range.end()),
         ));
     }
 }

crates/ruff/src/rules/flake8_simplify/snapshots/ruff__rules__flake8_simplify__tests__SIM114_SIM114.py.snap

@@ -138,4 +138,30 @@ SIM114.py:62:1: SIM114 Combine `if` branches using logical `or` operator
    | |______________^ SIM114
    |
 
+SIM114.py:109:5: SIM114 Combine `if` branches using logical `or` operator
+    |
+107 |       a = True
+108 |       b = False
+109 |       if a > b:  # end-of-line
+    |  _____^
+110 | |         return 3
+111 | |     elif a == b:
+112 | |         return 3
+    | |________________^ SIM114
+113 |       elif a < b:  # end-of-line
+114 |           return 4
+    |
+
+SIM114.py:113:5: SIM114 Combine `if` branches using logical `or` operator
+    |
+111 |       elif a == b:
+112 |           return 3
+113 |       elif a < b:  # end-of-line
+    |  _____^
+114 | |         return 4
+115 | |     elif b is None:
+116 | |         return 4
+    | |________________^ SIM114
+    |
+