Unnamed repository; edit this file 'description' to name the repository.
Merge pull request #22180 from A4-Tacks/arith-assign
feat: offer on compound assign for replace_arith_op
| -rw-r--r-- | crates/ide-assists/src/handlers/replace_arith_op.rs | 53 |
1 files changed, 39 insertions, 14 deletions
diff --git a/crates/ide-assists/src/handlers/replace_arith_op.rs b/crates/ide-assists/src/handlers/replace_arith_op.rs index 5ad5efac05..a04f6f3833 100644 --- a/crates/ide-assists/src/handlers/replace_arith_op.rs +++ b/crates/ide-assists/src/handlers/replace_arith_op.rs @@ -4,7 +4,10 @@ use syntax::{ ast::{self, ArithOp, BinaryOp}, }; -use crate::assist_context::{AssistContext, Assists}; +use crate::{ + assist_context::{AssistContext, Assists}, + utils::wrap_paren, +}; // Assist: replace_arith_with_checked // @@ -70,7 +73,7 @@ pub(crate) fn replace_arith_with_wrapping( } fn replace_arith(acc: &mut Assists, ctx: &AssistContext<'_>, kind: ArithKind) -> Option<()> { - let (lhs, op, rhs) = parse_binary_op(ctx)?; + let (lhs, op, is_assign, rhs) = parse_binary_op(ctx)?; let op_expr = lhs.syntax().parent()?; if !is_primitive_int(ctx, &lhs) || !is_primitive_int(ctx, &rhs) { @@ -87,11 +90,13 @@ fn replace_arith(acc: &mut Assists, ctx: &AssistContext<'_>, kind: ArithKind) -> let make = editor.make(); let method_name = kind.method_name(op); - let needs_parentheses = - lhs.precedence().needs_parentheses_in(ast::prec::ExprPrecedence::Postfix); - let receiver = if needs_parentheses { make.expr_paren(lhs).into() } else { lhs }; - let arith_expr = - make.expr_method_call(receiver, make.name_ref(&method_name), make.arg_list([rhs])); + let receiver = wrap_paren(lhs.clone(), make, ast::prec::ExprPrecedence::Postfix); + let mut arith_expr = make + .expr_method_call(receiver, make.name_ref(&method_name), make.arg_list([rhs])) + .into(); + if is_assign { + arith_expr = make.expr_assignment(lhs, arith_expr).into(); + } editor.replace(op_expr, arith_expr.syntax()); builder.add_file_edits(ctx.vfs_file_id(), editor); }, @@ -106,24 +111,25 @@ fn is_primitive_int(ctx: &AssistContext<'_>, expr: &ast::Expr) -> bool { } /// Extract the operands of an arithmetic expression (e.g. `1 + 2` or `1.checked_add(2)`) -fn parse_binary_op(ctx: &AssistContext<'_>) -> Option<(ast::Expr, ArithOp, ast::Expr)> { +fn parse_binary_op(ctx: &AssistContext<'_>) -> Option<(ast::Expr, ArithOp, bool, ast::Expr)> { if !ctx.has_empty_selection() { return None; } let expr = ctx.find_node_at_offset::<ast::BinExpr>()?; - let op = match expr.op_kind() { - Some(BinaryOp::ArithOp(ArithOp::Add)) => ArithOp::Add, - Some(BinaryOp::ArithOp(ArithOp::Sub)) => ArithOp::Sub, - Some(BinaryOp::ArithOp(ArithOp::Mul)) => ArithOp::Mul, - Some(BinaryOp::ArithOp(ArithOp::Div)) => ArithOp::Div, + let (op, is_assign) = match expr.op_kind()? { + BinaryOp::ArithOp(arith_op) => (arith_op, false), + BinaryOp::Assignment { op: Some(op) } => (op, true), _ => return None, }; + if !matches!(op, ArithOp::Add | ArithOp::Sub | ArithOp::Mul | ArithOp::Div) { + return None; + } let lhs = expr.lhs()?; let rhs = expr.rhs()?; - Some((lhs, op, rhs)) + Some((lhs, op, is_assign, rhs)) } pub(crate) enum ArithKind { @@ -250,6 +256,25 @@ fn main() { } #[test] + fn replace_arith_with_wrapping_add_assign() { + check_assist( + replace_arith_with_wrapping, + r#" +fn main() { + let mut x = 1; + x $0+= 2; +} +"#, + r#" +fn main() { + let mut x = 1; + x = x.wrapping_add(2); +} +"#, + ) + } + + #[test] fn replace_arith_not_applicable_with_non_empty_selection() { check_assist_not_applicable( replace_arith_with_checked, |