fix: Enhance remove_parentheses assist to handle return expressions

Author: asuto15

crates/ide-assists/src/handlers/remove_parentheses.rs

@@ -163,17 +163,18 @@ mod tests {
     }
 
     #[test]
-    fn remove_parens_prefix_with_return_no_value() {
-        // removing `()` from !(return) would make `!return` which is invalid syntax
-        check_assist_not_applicable(
+    fn remove_parens_prefix_with_ret_like_prefix() {
+        check_assist(remove_parentheses, r#"fn f() { !$0(return) }"#, r#"fn f() { !return }"#);
+        // `break`, `continue` behave the same under prefix operators
+        check_assist(remove_parentheses, r#"fn f() { !$0(break) }"#, r#"fn f() { !break }"#);
+        check_assist(remove_parentheses, r#"fn f() { !$0(continue) }"#, r#"fn f() { !continue }"#);
+        check_assist(
             remove_parentheses,
-            r#"fn main() { let _x = true && !$0(return) || true; }"#,
+            r#"fn f() { !$0(return false) }"#,
+            r#"fn f() { !return false }"#,
         );
-        check_assist_not_applicable(remove_parentheses, r#"fn f() { !$0(return) }"#);
-        check_assist_not_applicable(remove_parentheses, r#"fn f() { !$0(break) }"#);
-        check_assist_not_applicable(remove_parentheses, r#"fn f() { !$0(continue) }"#);
 
-        // Binary operators should still allow removal
+        // Binary operators should still allow removal unless a ret-like expression is immediately followed by `||` or `&&`.
         check_assist(
             remove_parentheses,
             r#"fn f() { true || $0(return) }"#,
@@ -247,6 +248,79 @@ mod tests {
         );
     }
 
+    #[test]
+    fn remove_parens_return_in_unary_not() {
+        check_assist(
+            remove_parentheses,
+            r#"fn f() { cond && !$0(return) }"#,
+            r#"fn f() { cond && !return }"#,
+        );
+        check_assist(
+            remove_parentheses,
+            r#"fn f() { cond && !$0(return false) }"#,
+            r#"fn f() { cond && !return false }"#,
+        );
+    }
+
+    #[test]
+    fn remove_parens_return_in_disjunction_with_closure_risk() {
+        // `return` may only be blocked when it would form `return ||` or `return &&`
+        check_assist_not_applicable(
+            remove_parentheses,
+            r#"fn f() { let _x = true && $0(return) || true; }"#,
+        );
+        check_assist_not_applicable(
+            remove_parentheses,
+            r#"fn f() { let _x = true && !$0(return) || true; }"#,
+        );
+        check_assist_not_applicable(
+            remove_parentheses,
+            r#"fn f() { let _x = true && $0(return false) || true; }"#,
+        );
+        check_assist_not_applicable(
+            remove_parentheses,
+            r#"fn f() { let _x = true && !$0(return false) || true; }"#,
+        );
+        check_assist_not_applicable(
+            remove_parentheses,
+            r#"fn f() { let _x = true && $0(return) && true; }"#,
+        );
+        check_assist_not_applicable(
+            remove_parentheses,
+            r#"fn f() { let _x = true && !$0(return) && true; }"#,
+        );
+        check_assist_not_applicable(
+            remove_parentheses,
+            r#"fn f() { let _x = true && $0(return false) && true; }"#,
+        );
+        check_assist_not_applicable(
+            remove_parentheses,
+            r#"fn f() { let _x = true && !$0(return false) && true; }"#,
+        );
+        check_assist_not_applicable(
+            remove_parentheses,
+            r#"fn f() { let _x = $0(return) || true; }"#,
+        );
+        check_assist_not_applicable(
+            remove_parentheses,
+            r#"fn f() { let _x = $0(return) && true; }"#,
+        );
+    }
+
+    #[test]
+    fn remove_parens_return_in_disjunction_is_ok() {
+        check_assist(
+            remove_parentheses,
+            r#"fn f() { let _x = true || $0(return); }"#,
+            r#"fn f() { let _x = true || return; }"#,
+        );
+        check_assist(
+            remove_parentheses,
+            r#"fn f() { let _x = true && $0(return); }"#,
+            r#"fn f() { let _x = true && return; }"#,
+        );
+    }
+
     #[test]
     fn remove_parens_double_paren_stmt() {
         check_assist(

crates/syntax/src/ast/prec.rs

@@ -3,14 +3,15 @@
 use stdx::always;
 
 use crate::{
-    AstNode, SyntaxNode,
+    AstNode, Direction, SyntaxNode, T,
+    algo::skip_trivia_token,
     ast::{self, BinaryOp, Expr, HasArgList, RangeItem},
     match_ast,
 };
 
 #[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
 pub enum ExprPrecedence {
-    // return val, break val, yield val, closures
+    // return, break, continue, yield, yeet, become (with or without value)
     Jump,
     // = += -= *= /= %= &= |= ^= <<= >>=
     Assign,
@@ -76,18 +77,12 @@ pub fn precedence(expr: &ast::Expr) -> ExprPrecedence {
             Some(_) => ExprPrecedence::Unambiguous,
         },
 
-        Expr::BreakExpr(e) if e.expr().is_some() => ExprPrecedence::Jump,
-        Expr::BecomeExpr(e) if e.expr().is_some() => ExprPrecedence::Jump,
-        Expr::ReturnExpr(e) if e.expr().is_some() => ExprPrecedence::Jump,
-        Expr::YeetExpr(e) if e.expr().is_some() => ExprPrecedence::Jump,
-        Expr::YieldExpr(e) if e.expr().is_some() => ExprPrecedence::Jump,
-
         Expr::BreakExpr(_)
         | Expr::BecomeExpr(_)
         | Expr::ReturnExpr(_)
         | Expr::YeetExpr(_)
         | Expr::YieldExpr(_)
-        | Expr::ContinueExpr(_) => ExprPrecedence::Unambiguous,
+        | Expr::ContinueExpr(_) => ExprPrecedence::Jump,
 
         Expr::RangeExpr(..) => ExprPrecedence::Range,
 
@@ -226,9 +221,17 @@ impl Expr {
             return false;
         }
 
-        // Special-case prefix operators with return/break/etc without value
-        // e.g., `!(return)` - parentheses are necessary
-        if self.is_ret_like_with_no_value() && parent.is_prefix() {
+        // Keep parens when a ret-like expr is followed by `||` or `&&`.
+        // For `||`, removing parens could reparse as `<ret-like> || <closure>`.
+        // For `&&`, we avoid introducing `<ret-like> && <expr>` into a binary chain.
+
+        if self.precedence() == ExprPrecedence::Jump
+            && let Some(node) =
+                place_of.ancestors().find(|it| it.parent().is_none_or(|p| &p == parent.syntax()))
+            && let Some(next) =
+                node.last_token().and_then(|t| skip_trivia_token(t.next_token()?, Direction::Next))
+            && matches!(next.kind(), T![||] | T![&&])
+        {
             return true;
         }