diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs index 2b24b5c31b..08fd4453df 100644 --- a/crates/hir/src/lib.rs +++ b/crates/hir/src/lib.rs @@ -2926,6 +2926,13 @@ impl Type { matches!(self.ty.kind(Interner), TyKind::Scalar(Scalar::Uint(UintTy::Usize))) } + pub fn is_int_or_uint(&self) -> bool { + match self.ty.kind(Interner) { + TyKind::Scalar(Scalar::Int(_) | Scalar::Uint(_)) => true, + _ => false, + } + } + pub fn remove_ref(&self) -> Option { match &self.ty.kind(Interner) { TyKind::Ref(.., ty) => Some(self.derived(ty.clone())), diff --git a/crates/ide-assists/src/handlers/replace_arith_op.rs b/crates/ide-assists/src/handlers/replace_arith_op.rs new file mode 100644 index 0000000000..f1ca35cafc --- /dev/null +++ b/crates/ide-assists/src/handlers/replace_arith_op.rs @@ -0,0 +1,226 @@ +use ide_db::assists::{AssistId, AssistKind, GroupLabel}; +use syntax::{ + ast::{self, ArithOp, BinaryOp}, + AstNode, TextRange, +}; + +use crate::assist_context::{AssistContext, Assists}; + +// Assist: replace_arith_with_checked +// +// Replaces arithmetic on integers with the `checked_*` equivalent. +// +// ``` +// fn main() { +// let x = 1 $0+ 2; +// } +// ``` +// -> +// ``` +// fn main() { +// let x = 1.checked_add(2); +// } +// ``` +pub(crate) fn replace_arith_with_checked(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> { + replace_arith(acc, ctx, ArithKind::Checked) +} + +// Assist: replace_arith_with_saturating +// +// Replaces arithmetic on integers with the `saturating_*` equivalent. +// +// ``` +// fn main() { +// let x = 1 $0+ 2; +// } +// ``` +// -> +// ``` +// fn main() { +// let x = 1.saturating_add(2); +// } +// ``` +pub(crate) fn replace_arith_with_saturating( + acc: &mut Assists, + ctx: &AssistContext<'_>, +) -> Option<()> { + replace_arith(acc, ctx, ArithKind::Saturating) +} + +// Assist: replace_arith_with_wrapping +// +// Replaces arithmetic on integers with the `wrapping_*` equivalent. +// +// ``` +// fn main() { +// let x = 1 $0+ 2; +// } +// ``` +// -> +// ``` +// fn main() { +// let x = 1.wrapping_add(2); +// } +// ``` +pub(crate) fn replace_arith_with_wrapping( + acc: &mut Assists, + ctx: &AssistContext<'_>, +) -> Option<()> { + replace_arith(acc, ctx, ArithKind::Wrapping) +} + +fn replace_arith(acc: &mut Assists, ctx: &AssistContext<'_>, kind: ArithKind) -> Option<()> { + let (lhs, op, rhs) = parse_binary_op(ctx)?; + + if !is_primitive_int(ctx, &lhs) || !is_primitive_int(ctx, &rhs) { + return None; + } + + let start = lhs.syntax().text_range().start(); + let end = rhs.syntax().text_range().end(); + let range = TextRange::new(start, end); + + acc.add_group( + &GroupLabel("replace_arith".into()), + kind.assist_id(), + kind.label(), + range, + |builder| { + let method_name = kind.method_name(op); + + builder.replace(range, format!("{lhs}.{method_name}({rhs})")) + }, + ) +} + +fn is_primitive_int(ctx: &AssistContext<'_>, expr: &ast::Expr) -> bool { + match ctx.sema.type_of_expr(expr) { + Some(ty) => ty.adjusted().is_int_or_uint(), + _ => false, + } +} + +/// 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)> { + let expr = ctx.find_node_at_offset::()?; + + 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, + _ => return None, + }; + + let lhs = expr.lhs()?; + let rhs = expr.rhs()?; + + Some((lhs, op, rhs)) +} + +pub(crate) enum ArithKind { + Saturating, + Wrapping, + Checked, +} + +impl ArithKind { + fn assist_id(&self) -> AssistId { + let s = match self { + ArithKind::Saturating => "replace_arith_with_saturating", + ArithKind::Checked => "replace_arith_with_checked", + ArithKind::Wrapping => "replace_arith_with_wrapping", + }; + + AssistId(s, AssistKind::RefactorRewrite) + } + + fn label(&self) -> &'static str { + match self { + ArithKind::Saturating => "Replace arithmetic with call to saturating_*", + ArithKind::Checked => "Replace arithmetic with call to checked_*", + ArithKind::Wrapping => "Replace arithmetic with call to wrapping_*", + } + } + + fn method_name(&self, op: ArithOp) -> String { + let prefix = match self { + ArithKind::Checked => "checked_", + ArithKind::Wrapping => "wrapping_", + ArithKind::Saturating => "saturating_", + }; + + let suffix = match op { + ArithOp::Add => "add", + ArithOp::Sub => "sub", + ArithOp::Mul => "mul", + ArithOp::Div => "div", + _ => unreachable!("this function should only be called with +, -, / or *"), + }; + format!("{prefix}{suffix}") + } +} + +#[cfg(test)] +mod tests { + use crate::tests::check_assist; + + use super::*; + + #[test] + fn arith_kind_method_name() { + assert_eq!(ArithKind::Saturating.method_name(ArithOp::Add), "saturating_add"); + assert_eq!(ArithKind::Checked.method_name(ArithOp::Sub), "checked_sub"); + } + + #[test] + fn replace_arith_with_checked_add() { + check_assist( + replace_arith_with_checked, + r#" +fn main() { + let x = 1 $0+ 2; +} +"#, + r#" +fn main() { + let x = 1.checked_add(2); +} +"#, + ) + } + + #[test] + fn replace_arith_with_saturating_add() { + check_assist( + replace_arith_with_saturating, + r#" +fn main() { + let x = 1 $0+ 2; +} +"#, + r#" +fn main() { + let x = 1.saturating_add(2); +} +"#, + ) + } + + #[test] + fn replace_arith_with_wrapping_add() { + check_assist( + replace_arith_with_wrapping, + r#" +fn main() { + let x = 1 $0+ 2; +} +"#, + r#" +fn main() { + let x = 1.wrapping_add(2); +} +"#, + ) + } +} diff --git a/crates/ide-assists/src/lib.rs b/crates/ide-assists/src/lib.rs index 8b1247c640..06dd0efa2b 100644 --- a/crates/ide-assists/src/lib.rs +++ b/crates/ide-assists/src/lib.rs @@ -186,6 +186,7 @@ mod handlers { mod replace_derive_with_manual_impl; mod replace_if_let_with_match; mod replace_or_with_or_else; + mod replace_arith_op; mod introduce_named_generic; mod replace_let_with_if_let; mod replace_qualified_name_with_use; @@ -293,6 +294,9 @@ mod handlers { replace_or_with_or_else::replace_or_with_or_else, replace_turbofish_with_explicit_type::replace_turbofish_with_explicit_type, replace_qualified_name_with_use::replace_qualified_name_with_use, + replace_arith_op::replace_arith_with_wrapping, + replace_arith_op::replace_arith_with_checked, + replace_arith_op::replace_arith_with_saturating, sort_items::sort_items, split_import::split_import, toggle_ignore::toggle_ignore, diff --git a/crates/ide-assists/src/tests/generated.rs b/crates/ide-assists/src/tests/generated.rs index 80b8c27c7c..6b340c79c8 100644 --- a/crates/ide-assists/src/tests/generated.rs +++ b/crates/ide-assists/src/tests/generated.rs @@ -2066,6 +2066,57 @@ impl Foo for Bar { ) } +#[test] +fn doctest_replace_arith_with_checked() { + check_doc_test( + "replace_arith_with_checked", + r#####" +fn main() { + let x = 1 $0+ 2; +} +"#####, + r#####" +fn main() { + let x = 1.checked_add(2); +} +"#####, + ) +} + +#[test] +fn doctest_replace_arith_with_saturating() { + check_doc_test( + "replace_arith_with_saturating", + r#####" +fn main() { + let x = 1 $0+ 2; +} +"#####, + r#####" +fn main() { + let x = 1.saturating_add(2); +} +"#####, + ) +} + +#[test] +fn doctest_replace_arith_with_wrapping() { + check_doc_test( + "replace_arith_with_wrapping", + r#####" +fn main() { + let x = 1 $0+ 2; +} +"#####, + r#####" +fn main() { + let x = 1.wrapping_add(2); +} +"#####, + ) +} + #[test] fn doctest_replace_char_with_string() { check_doc_test(