|
1 | 1 | use ide_db::helpers::{for_each_tail_expr, node_ext::walk_expr, FamousDefs};
|
| 2 | +use itertools::Itertools; |
2 | 3 | use syntax::{
|
3 | 4 | ast::{self, Expr},
|
4 |
| - match_ast, AstNode, |
| 5 | + match_ast, AstNode, TextRange, TextSize, |
5 | 6 | };
|
6 | 7 |
|
7 | 8 | use crate::{AssistContext, AssistId, AssistKind, Assists};
|
@@ -60,25 +61,55 @@ pub(crate) fn unwrap_result_return_type(acc: &mut Assists, ctx: &AssistContext)
|
60 | 61 | });
|
61 | 62 | for_each_tail_expr(&body, tail_cb);
|
62 | 63 |
|
63 |
| - for ret_expr_arg in exprs_to_unwrap { |
64 |
| - let new_ret_expr = ret_expr_arg.to_string(); |
65 |
| - let new_ret_expr = |
66 |
| - new_ret_expr.trim_start_matches("Ok(").trim_start_matches("Err("); |
67 |
| - builder.replace( |
68 |
| - ret_expr_arg.syntax().text_range(), |
69 |
| - new_ret_expr.strip_suffix(')').unwrap_or(new_ret_expr), |
70 |
| - ) |
71 |
| - } |
72 |
| - |
| 64 | + let mut is_unit_type = false; |
73 | 65 | if let Some((_, inner_type)) = type_ref.to_string().split_once('<') {
|
74 | 66 | let inner_type = match inner_type.split_once(',') {
|
75 | 67 | Some((success_inner_type, _)) => success_inner_type,
|
76 | 68 | None => inner_type,
|
77 | 69 | };
|
78 |
| - builder.replace( |
79 |
| - type_ref.syntax().text_range(), |
80 |
| - inner_type.strip_suffix('>').unwrap_or(inner_type), |
81 |
| - ) |
| 70 | + let new_ret_type = inner_type.strip_suffix('>').unwrap_or(inner_type); |
| 71 | + if new_ret_type == "()" { |
| 72 | + is_unit_type = true; |
| 73 | + let text_range = TextRange::new( |
| 74 | + ret_type.syntax().text_range().start(), |
| 75 | + ret_type.syntax().text_range().end() + TextSize::from(1u32), |
| 76 | + ); |
| 77 | + builder.delete(text_range) |
| 78 | + } else { |
| 79 | + builder.replace( |
| 80 | + type_ref.syntax().text_range(), |
| 81 | + inner_type.strip_suffix('>').unwrap_or(inner_type), |
| 82 | + ) |
| 83 | + } |
| 84 | + } |
| 85 | + |
| 86 | + for ret_expr_arg in exprs_to_unwrap { |
| 87 | + let ret_expr_str = ret_expr_arg.to_string(); |
| 88 | + if ret_expr_str.starts_with("Ok(") || ret_expr_str.starts_with("Err(") { |
| 89 | + let arg_list = ret_expr_arg.syntax().children().find_map(ast::ArgList::cast); |
| 90 | + if let Some(arg_list) = arg_list { |
| 91 | + if is_unit_type { |
| 92 | + match ret_expr_arg.syntax().prev_sibling_or_token() { |
| 93 | + // Useful to delete the entire line without leaving trailing whitespaces |
| 94 | + Some(whitespace) => { |
| 95 | + let new_range = TextRange::new( |
| 96 | + whitespace.text_range().start(), |
| 97 | + ret_expr_arg.syntax().text_range().end(), |
| 98 | + ); |
| 99 | + builder.delete(new_range); |
| 100 | + } |
| 101 | + None => { |
| 102 | + builder.delete(ret_expr_arg.syntax().text_range()); |
| 103 | + } |
| 104 | + } |
| 105 | + } else { |
| 106 | + builder.replace( |
| 107 | + ret_expr_arg.syntax().text_range(), |
| 108 | + arg_list.args().join(", "), |
| 109 | + ); |
| 110 | + } |
| 111 | + } |
| 112 | + } |
82 | 113 | }
|
83 | 114 | },
|
84 | 115 | )
|
@@ -126,6 +157,49 @@ fn foo() -> i32 {
|
126 | 157 | );
|
127 | 158 | }
|
128 | 159 |
|
| 160 | + #[test] |
| 161 | + fn unwrap_result_return_type_unit_type() { |
| 162 | + check_assist( |
| 163 | + unwrap_result_return_type, |
| 164 | + r#" |
| 165 | +//- minicore: result |
| 166 | +fn foo() -> Result<(), Box<dyn Error$0>> { |
| 167 | + Ok(()) |
| 168 | +} |
| 169 | +"#, |
| 170 | + r#" |
| 171 | +fn foo() { |
| 172 | +} |
| 173 | +"#, |
| 174 | + ); |
| 175 | + } |
| 176 | + |
| 177 | + #[test] |
| 178 | + fn unwrap_result_return_type_ending_with_parent() { |
| 179 | + check_assist( |
| 180 | + unwrap_result_return_type, |
| 181 | + r#" |
| 182 | +//- minicore: result |
| 183 | +fn foo() -> Result<i32, Box<dyn Error$0>> { |
| 184 | + if true { |
| 185 | + Ok(42) |
| 186 | + } else { |
| 187 | + foo() |
| 188 | + } |
| 189 | +} |
| 190 | +"#, |
| 191 | + r#" |
| 192 | +fn foo() -> i32 { |
| 193 | + if true { |
| 194 | + 42 |
| 195 | + } else { |
| 196 | + foo() |
| 197 | + } |
| 198 | +} |
| 199 | +"#, |
| 200 | + ); |
| 201 | + } |
| 202 | + |
129 | 203 | #[test]
|
130 | 204 | fn unwrap_return_type_break_split_tail() {
|
131 | 205 | check_assist(
|
|
0 commit comments