fix tests

Author: cameron1024

crates/hir/src/lib.rs

@@ -2857,6 +2857,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<Type> {
         match &self.ty.kind(Interner) {
             TyKind::Ref(.., ty) => Some(self.derived(ty.clone())),

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

@@ -1,10 +1,22 @@
 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<()> {
+// 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)
 }
 

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

@@ -1,6 +1,21 @@
 use crate::assist_context::{AssistContext, Assists};
 use crate::utils::{replace_arith, ArithKind};
 
+// 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<'_>,

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

@@ -1,6 +1,21 @@
 use crate::assist_context::{AssistContext, Assists};
 use crate::utils::{replace_arith, ArithKind};
 
+// 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<'_>,

crates/ide-assists/src/tests/generated.rs

@@ -1978,6 +1978,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(

crates/ide-assists/src/utils.rs

@@ -722,8 +722,8 @@ 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",
+            ArithKind::Checked => "replace_arith_with_checked",
+            ArithKind::Wrapping => "replace_arith_with_wrapping",
         };
 
         AssistId(s, AssistKind::RefactorRewrite)
@@ -771,6 +771,10 @@ pub(crate) fn replace_arith(
 ) -> 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);
@@ -782,6 +786,13 @@ pub(crate) fn replace_arith(
     })
 }
 
+fn is_primitive_int(ctx: &AssistContext<'_>, expr: &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<(Expr, ArithOp, Expr)> {
     let expr = ctx.find_node_at_offset::<BinExpr>()?;