]>
Commit | Line | Data |
---|---|---|
9c376795 FG |
1 | use ide_db::assists::{AssistId, AssistKind, GroupLabel}; |
2 | use syntax::{ | |
3 | ast::{self, ArithOp, BinaryOp}, | |
4 | AstNode, TextRange, | |
5 | }; | |
6 | ||
7 | use crate::assist_context::{AssistContext, Assists}; | |
8 | ||
9 | // Assist: replace_arith_with_checked | |
10 | // | |
11 | // Replaces arithmetic on integers with the `checked_*` equivalent. | |
12 | // | |
13 | // ``` | |
14 | // fn main() { | |
15 | // let x = 1 $0+ 2; | |
16 | // } | |
17 | // ``` | |
18 | // -> | |
19 | // ``` | |
20 | // fn main() { | |
21 | // let x = 1.checked_add(2); | |
22 | // } | |
23 | // ``` | |
24 | pub(crate) fn replace_arith_with_checked(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> { | |
25 | replace_arith(acc, ctx, ArithKind::Checked) | |
26 | } | |
27 | ||
28 | // Assist: replace_arith_with_saturating | |
29 | // | |
30 | // Replaces arithmetic on integers with the `saturating_*` equivalent. | |
31 | // | |
32 | // ``` | |
33 | // fn main() { | |
34 | // let x = 1 $0+ 2; | |
35 | // } | |
36 | // ``` | |
37 | // -> | |
38 | // ``` | |
39 | // fn main() { | |
40 | // let x = 1.saturating_add(2); | |
41 | // } | |
42 | // ``` | |
43 | pub(crate) fn replace_arith_with_saturating( | |
44 | acc: &mut Assists, | |
45 | ctx: &AssistContext<'_>, | |
46 | ) -> Option<()> { | |
47 | replace_arith(acc, ctx, ArithKind::Saturating) | |
48 | } | |
49 | ||
50 | // Assist: replace_arith_with_wrapping | |
51 | // | |
52 | // Replaces arithmetic on integers with the `wrapping_*` equivalent. | |
53 | // | |
54 | // ``` | |
55 | // fn main() { | |
56 | // let x = 1 $0+ 2; | |
57 | // } | |
58 | // ``` | |
59 | // -> | |
60 | // ``` | |
61 | // fn main() { | |
62 | // let x = 1.wrapping_add(2); | |
63 | // } | |
64 | // ``` | |
65 | pub(crate) fn replace_arith_with_wrapping( | |
66 | acc: &mut Assists, | |
67 | ctx: &AssistContext<'_>, | |
68 | ) -> Option<()> { | |
69 | replace_arith(acc, ctx, ArithKind::Wrapping) | |
70 | } | |
71 | ||
72 | fn replace_arith(acc: &mut Assists, ctx: &AssistContext<'_>, kind: ArithKind) -> Option<()> { | |
73 | let (lhs, op, rhs) = parse_binary_op(ctx)?; | |
74 | ||
75 | if !is_primitive_int(ctx, &lhs) || !is_primitive_int(ctx, &rhs) { | |
76 | return None; | |
77 | } | |
78 | ||
79 | let start = lhs.syntax().text_range().start(); | |
80 | let end = rhs.syntax().text_range().end(); | |
81 | let range = TextRange::new(start, end); | |
82 | ||
83 | acc.add_group( | |
9ffffee4 | 84 | &GroupLabel("Replace arithmetic...".into()), |
9c376795 FG |
85 | kind.assist_id(), |
86 | kind.label(), | |
87 | range, | |
88 | |builder| { | |
89 | let method_name = kind.method_name(op); | |
90 | ||
91 | builder.replace(range, format!("{lhs}.{method_name}({rhs})")) | |
92 | }, | |
93 | ) | |
94 | } | |
95 | ||
96 | fn is_primitive_int(ctx: &AssistContext<'_>, expr: &ast::Expr) -> bool { | |
97 | match ctx.sema.type_of_expr(expr) { | |
98 | Some(ty) => ty.adjusted().is_int_or_uint(), | |
99 | _ => false, | |
100 | } | |
101 | } | |
102 | ||
103 | /// Extract the operands of an arithmetic expression (e.g. `1 + 2` or `1.checked_add(2)`) | |
104 | fn parse_binary_op(ctx: &AssistContext<'_>) -> Option<(ast::Expr, ArithOp, ast::Expr)> { | |
105 | let expr = ctx.find_node_at_offset::<ast::BinExpr>()?; | |
106 | ||
107 | let op = match expr.op_kind() { | |
108 | Some(BinaryOp::ArithOp(ArithOp::Add)) => ArithOp::Add, | |
109 | Some(BinaryOp::ArithOp(ArithOp::Sub)) => ArithOp::Sub, | |
110 | Some(BinaryOp::ArithOp(ArithOp::Mul)) => ArithOp::Mul, | |
111 | Some(BinaryOp::ArithOp(ArithOp::Div)) => ArithOp::Div, | |
112 | _ => return None, | |
113 | }; | |
114 | ||
115 | let lhs = expr.lhs()?; | |
116 | let rhs = expr.rhs()?; | |
117 | ||
118 | Some((lhs, op, rhs)) | |
119 | } | |
120 | ||
121 | pub(crate) enum ArithKind { | |
122 | Saturating, | |
123 | Wrapping, | |
124 | Checked, | |
125 | } | |
126 | ||
127 | impl ArithKind { | |
128 | fn assist_id(&self) -> AssistId { | |
129 | let s = match self { | |
130 | ArithKind::Saturating => "replace_arith_with_saturating", | |
131 | ArithKind::Checked => "replace_arith_with_checked", | |
132 | ArithKind::Wrapping => "replace_arith_with_wrapping", | |
133 | }; | |
134 | ||
135 | AssistId(s, AssistKind::RefactorRewrite) | |
136 | } | |
137 | ||
138 | fn label(&self) -> &'static str { | |
139 | match self { | |
140 | ArithKind::Saturating => "Replace arithmetic with call to saturating_*", | |
141 | ArithKind::Checked => "Replace arithmetic with call to checked_*", | |
142 | ArithKind::Wrapping => "Replace arithmetic with call to wrapping_*", | |
143 | } | |
144 | } | |
145 | ||
146 | fn method_name(&self, op: ArithOp) -> String { | |
147 | let prefix = match self { | |
148 | ArithKind::Checked => "checked_", | |
149 | ArithKind::Wrapping => "wrapping_", | |
150 | ArithKind::Saturating => "saturating_", | |
151 | }; | |
152 | ||
153 | let suffix = match op { | |
154 | ArithOp::Add => "add", | |
155 | ArithOp::Sub => "sub", | |
156 | ArithOp::Mul => "mul", | |
157 | ArithOp::Div => "div", | |
158 | _ => unreachable!("this function should only be called with +, -, / or *"), | |
159 | }; | |
160 | format!("{prefix}{suffix}") | |
161 | } | |
162 | } | |
163 | ||
164 | #[cfg(test)] | |
165 | mod tests { | |
166 | use crate::tests::check_assist; | |
167 | ||
168 | use super::*; | |
169 | ||
170 | #[test] | |
171 | fn arith_kind_method_name() { | |
172 | assert_eq!(ArithKind::Saturating.method_name(ArithOp::Add), "saturating_add"); | |
173 | assert_eq!(ArithKind::Checked.method_name(ArithOp::Sub), "checked_sub"); | |
174 | } | |
175 | ||
176 | #[test] | |
177 | fn replace_arith_with_checked_add() { | |
178 | check_assist( | |
179 | replace_arith_with_checked, | |
180 | r#" | |
181 | fn main() { | |
182 | let x = 1 $0+ 2; | |
183 | } | |
184 | "#, | |
185 | r#" | |
186 | fn main() { | |
187 | let x = 1.checked_add(2); | |
188 | } | |
189 | "#, | |
190 | ) | |
191 | } | |
192 | ||
193 | #[test] | |
194 | fn replace_arith_with_saturating_add() { | |
195 | check_assist( | |
196 | replace_arith_with_saturating, | |
197 | r#" | |
198 | fn main() { | |
199 | let x = 1 $0+ 2; | |
200 | } | |
201 | "#, | |
202 | r#" | |
203 | fn main() { | |
204 | let x = 1.saturating_add(2); | |
205 | } | |
206 | "#, | |
207 | ) | |
208 | } | |
209 | ||
210 | #[test] | |
211 | fn replace_arith_with_wrapping_add() { | |
212 | check_assist( | |
213 | replace_arith_with_wrapping, | |
214 | r#" | |
215 | fn main() { | |
216 | let x = 1 $0+ 2; | |
217 | } | |
218 | "#, | |
219 | r#" | |
220 | fn main() { | |
221 | let x = 1.wrapping_add(2); | |
222 | } | |
223 | "#, | |
224 | ) | |
225 | } | |
226 | } |