add wrapping/checked/saturating assist

5 files changed, 215 insertions, 4 deletions

diff --git a/crates/ide-assists/src/handlers/replace_arith_with_checked.rs b/crates/ide-assists/src/handlers/replace_arith_with_checked.rs
new file mode 100644
index 00000000000..419b6febf1c
--- /dev/null
+++ b/crates/ide-assists/src/handlers/replace_arith_with_checked.rs
@@ -0,0 +1,33 @@
+use crate::assist_context::{AssistContext, Assists};
+use crate::utils::{replace_arith, ArithKind};
+
+pub(crate) fn replace_arith_with_checked(
+    acc: &mut Assists,
+    ctx: &AssistContext<'_>,
+) -> Option<()> {
+    replace_arith(acc, ctx, ArithKind::Checked)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::check_assist;
+
+    use super::*;
+
+    #[test]
+    fn replace_arith_with_saturating_add() {
+        check_assist(
+            replace_arith_with_checked,
+            r#"
+fn main() {
+    let x = 1 $0+ 2;
+}
+"#,
+            r#"
+fn main() {
+    let x = 1.checked_add(2);
+}
+"#,
+        )
+    }
+}
diff --git a/crates/ide-assists/src/handlers/replace_arith_with_saturating.rs b/crates/ide-assists/src/handlers/replace_arith_with_saturating.rs
new file mode 100644
index 00000000000..8c620ad93f0
--- /dev/null
+++ b/crates/ide-assists/src/handlers/replace_arith_with_saturating.rs
@@ -0,0 +1,33 @@
+use crate::assist_context::{AssistContext, Assists};
+use crate::utils::{replace_arith, ArithKind};
+
+pub(crate) fn replace_arith_with_saturating(
+    acc: &mut Assists,
+    ctx: &AssistContext<'_>,
+) -> Option<()> {
+    replace_arith(acc, ctx, ArithKind::Saturating)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::check_assist;
+
+    use super::*;
+
+    #[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);
+}
+"#,
+        )
+    }
+}
diff --git a/crates/ide-assists/src/handlers/replace_arith_with_wrapping.rs b/crates/ide-assists/src/handlers/replace_arith_with_wrapping.rs
new file mode 100644
index 00000000000..0054f580263
--- /dev/null
+++ b/crates/ide-assists/src/handlers/replace_arith_with_wrapping.rs
@@ -0,0 +1,33 @@
+use crate::assist_context::{AssistContext, Assists};
+use crate::utils::{replace_arith, ArithKind};
+
+pub(crate) fn replace_arith_with_wrapping(
+    acc: &mut Assists,
+    ctx: &AssistContext<'_>,
+) -> Option<()> {
+    replace_arith(acc, ctx, ArithKind::Wrapping)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::tests::check_assist;
+
+    use super::*;
+
+    #[test]
+    fn replace_arith_with_saturating_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 a07318cefad..5728730c913 100644
--- a/crates/ide-assists/src/lib.rs
+++ b/crates/ide-assists/src/lib.rs
@@ -183,6 +183,10 @@ mod handlers {
     mod replace_derive_with_manual_impl;
     mod replace_if_let_with_match;
     mod replace_or_with_or_else;
+    mod replace_arith_with_saturating;
+    mod replace_arith_with_checked;
+    mod replace_arith_with_wrapping;
+    mod replace_arith_with_saturating;
     mod introduce_named_generic;
     mod replace_let_with_if_let;
     mod replace_qualified_name_with_use;
@@ -286,6 +290,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_with_wrapping::replace_arith_with_wrapping,
+            replace_arith_with_checked::replace_arith_with_checked,
+            replace_arith_with_saturating::replace_arith_with_saturating,
             sort_items::sort_items,
             split_import::split_import,
             toggle_ignore::toggle_ignore,
diff --git a/crates/ide-assists/src/utils.rs b/crates/ide-assists/src/utils.rs
index db32e7182c4..dd2b6785c4b 100644
--- a/crates/ide-assists/src/utils.rs
+++ b/crates/ide-assists/src/utils.rs
@@ -4,21 +4,27 @@ use std::ops;
 
 pub(crate) use gen_trait_fn_body::gen_trait_fn_body;
 use hir::{db::HirDatabase, HirDisplay, Semantics};
-use ide_db::{famous_defs::FamousDefs, path_transform::PathTransform, RootDatabase, SnippetCap};
+use ide_db::{
+    assists::{AssistId, AssistKind},
+    famous_defs::FamousDefs,
+    path_transform::PathTransform,
+    RootDatabase, SnippetCap,
+};
 use stdx::format_to;
 use syntax::{
     ast::{
         self,
         edit::{self, AstNodeEdit},
         edit_in_place::{AttrsOwnerEdit, Removable},
-        make, HasArgList, HasAttrs, HasGenericParams, HasName, HasTypeBounds, Whitespace,
+        make, ArithOp, BinExpr, BinaryOp, Expr, HasArgList, HasAttrs, HasGenericParams, HasName,
+        HasTypeBounds, Whitespace,
     },
-    ted, AstNode, AstToken, Direction, SourceFile,
+    ted, AstNode, AstToken, Direction, SmolStr, SourceFile,
     SyntaxKind::*,
     SyntaxNode, TextRange, TextSize, T,
 };
 
-use crate::assist_context::{AssistContext, SourceChangeBuilder};
+use crate::assist_context::{AssistContext, Assists, SourceChangeBuilder};
 
 pub(crate) mod suggest_name;
 mod gen_trait_fn_body;
@@ -705,3 +711,102 @@ pub(crate) fn convert_param_list_to_arg_list(list: ast::ParamList) -> ast::ArgLi
     }
     make::arg_list(args)
 }
+
+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_saturating",
+            ArithKind::Wrapping => "replace_arith_with_saturating",
+        };
+
+        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) -> SmolStr {
+        // is this too much effort to avoid an allocation? is there a better way?
+        let mut bytes = [0u8; 14];
+        let prefix = match self {
+            ArithKind::Checked => "checked_",
+            ArithKind::Wrapping => "wrapping_",
+            ArithKind::Saturating => "saturating_",
+        };
+
+        bytes[0..(prefix.len())].copy_from_slice(prefix.as_bytes());
+
+        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 *"),
+        };
+
+        bytes[(prefix.len())..(prefix.len() + suffix.len())].copy_from_slice(suffix.as_bytes());
+
+        let len = prefix.len() + suffix.len();
+        let s = core::str::from_utf8(&bytes[0..len]).unwrap();
+        SmolStr::from(s)
+    }
+}
+
+pub(crate) fn replace_arith(
+    acc: &mut Assists,
+    ctx: &AssistContext<'_>,
+    kind: ArithKind,
+) -> Option<()> {
+    let (lhs, op, rhs) = parse_binary_op(ctx)?;
+
+    let start = lhs.syntax().text_range().start();
+    let end = rhs.syntax().text_range().end();
+    let range = TextRange::new(start, end);
+
+    acc.add(kind.assist_id(), kind.label(), range, |builder| {
+        let method_name = kind.method_name(op);
+
+        builder.replace(range, format!("{lhs}.{method_name}({rhs})"))
+    })
+}
+
+/// Extract the operands of an arithmetic expression (e.g. `1 + 2` or `1.checked_add(2)`)
+fn parse_binary_op(ctx: &AssistContext<'_>) -> Option<(Expr, ArithOp, Expr)> {
+    let expr = ctx.find_node_at_offset::<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,
+        _ => return None,
+    };
+
+    let lhs = expr.lhs()?;
+    let rhs = expr.rhs()?;
+
+    Some((lhs, op, rhs))
+}
+
+#[cfg(test)]
+mod tests {
+    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");
+    }
+}