@@ -28,50 +28,54 @@ use crate::{AssistContext, AssistId, AssistKind, Assists};
28
28
/// }
29
29
/// ```
30
30
pub ( crate ) fn convert_iter_for_each_to_for ( acc : & mut Assists , ctx : & AssistContext ) -> Option < ( ) > {
31
- let closure;
32
-
33
- let total_expr = match ctx. find_node_at_offset :: < ast:: Expr > ( ) ? {
34
- ast:: Expr :: MethodCallExpr ( expr) => {
35
- closure = match expr. arg_list ( ) ?. args ( ) . next ( ) ? {
36
- ast:: Expr :: ClosureExpr ( expr) => expr,
37
- _ => {
38
- return None ;
39
- }
40
- } ;
41
-
42
- expr
43
- }
44
- ast:: Expr :: ClosureExpr ( expr) => {
45
- closure = expr;
46
- ast:: MethodCallExpr :: cast ( closure. syntax ( ) . ancestors ( ) . nth ( 2 ) ?) ?
47
- }
31
+ let method;
32
+
33
+ let stmt = if let Some ( stmt) = ctx. find_node_at_offset :: < ast:: ExprStmt > ( ) {
34
+ method = ast:: MethodCallExpr :: cast ( stmt. syntax ( ) . first_child ( ) ?) ?;
35
+ Some ( stmt)
36
+ } else {
37
+ method = match ctx. find_node_at_offset :: < ast:: Expr > ( ) ? {
38
+ ast:: Expr :: MethodCallExpr ( expr) => expr,
39
+ ast:: Expr :: ClosureExpr ( expr) => {
40
+ ast:: MethodCallExpr :: cast ( expr. syntax ( ) . ancestors ( ) . nth ( 2 ) ?) ?
41
+ }
42
+ _ => {
43
+ return None ;
44
+ }
45
+ } ;
46
+ None
47
+ } ;
48
+
49
+ let closure = match method. arg_list ( ) ?. args ( ) . next ( ) ? {
50
+ ast:: Expr :: ClosureExpr ( expr) => expr,
48
51
_ => {
49
52
return None ;
50
53
}
51
54
} ;
52
55
53
- let ( total_expr , parent) = validate_method_call_expr ( & ctx. sema , total_expr ) ?;
56
+ let ( method , parent) = validate_method_call_expr ( & ctx. sema , method ) ?;
54
57
55
58
let param_list = closure. param_list ( ) ?;
56
59
let param = param_list. params ( ) . next ( ) ?. pat ( ) ?;
57
60
let body = closure. body ( ) ?;
58
61
62
+ let indent = stmt. as_ref ( ) . map_or ( method. indent_level ( ) , |stmt| stmt. indent_level ( ) ) ;
63
+ let syntax = stmt. as_ref ( ) . map_or ( method. syntax ( ) , |stmt| stmt. syntax ( ) ) ;
64
+
59
65
acc. add (
60
66
AssistId ( "convert_iter_for_each_to_for" , AssistKind :: RefactorRewrite ) ,
61
67
"Replace this `Iterator::for_each` with a for loop" ,
62
- total_expr . syntax ( ) . text_range ( ) ,
68
+ syntax. text_range ( ) ,
63
69
|builder| {
64
- let original_indentation = total_expr. indent_level ( ) ;
65
-
66
70
let block = match body {
67
71
ast:: Expr :: BlockExpr ( block) => block,
68
72
_ => make:: block_expr ( Vec :: new ( ) , Some ( body) ) ,
69
73
}
70
74
. reset_indent ( )
71
- . indent ( original_indentation ) ;
75
+ . indent ( indent ) ;
72
76
73
77
let expr_for_loop = make:: expr_for_loop ( param, parent, block) ;
74
- builder. replace_ast ( total_expr , expr_for_loop)
78
+ builder. replace ( syntax . text_range ( ) , expr_for_loop. syntax ( ) . text ( ) )
75
79
} ,
76
80
)
77
81
}
@@ -125,7 +129,7 @@ impl Empty {
125
129
}
126
130
127
131
#[ test]
128
- fn test_for_each_in_method ( ) {
132
+ fn test_for_each_in_method_stmt ( ) {
129
133
check_assist_with_fixtures (
130
134
r#"
131
135
use empty_iter::*;
@@ -141,14 +145,37 @@ fn main() {
141
145
let x = Empty;
142
146
for (x, y) in x.iter() {
143
147
println!("x: {}, y: {}", x, y);
144
- };
148
+ }
145
149
}
146
150
"# ,
147
151
)
148
152
}
149
153
150
154
#[ test]
151
- fn test_for_each_without_braces ( ) {
155
+ fn test_for_each_in_method ( ) {
156
+ check_assist_with_fixtures (
157
+ r#"
158
+ use empty_iter::*;
159
+ fn main() {
160
+ let x = Empty;
161
+ x.iter().$0for_each(|(x, y)| {
162
+ println!("x: {}, y: {}", x, y);
163
+ })
164
+ }"# ,
165
+ r#"
166
+ use empty_iter::*;
167
+ fn main() {
168
+ let x = Empty;
169
+ for (x, y) in x.iter() {
170
+ println!("x: {}, y: {}", x, y);
171
+ }
172
+ }
173
+ "# ,
174
+ )
175
+ }
176
+
177
+ #[ test]
178
+ fn test_for_each_without_braces_stmt ( ) {
152
179
check_assist_with_fixtures (
153
180
r#"
154
181
use empty_iter::*;
@@ -162,14 +189,14 @@ fn main() {
162
189
let x = Empty;
163
190
for (x, y) in x.iter() {
164
191
println!("x: {}, y: {}", x, y)
165
- };
192
+ }
166
193
}
167
194
"# ,
168
195
)
169
196
}
170
197
171
198
#[ test]
172
- fn test_for_each_in_closure ( ) {
199
+ fn test_for_each_in_closure_stmt ( ) {
173
200
check_assist_with_fixtures (
174
201
r#"
175
202
use empty_iter::*;
@@ -183,7 +210,7 @@ fn main() {
183
210
let x = Empty;
184
211
for (x, y) in x.iter() {
185
212
println!("x: {}, y: {}", x, y)
186
- };
213
+ }
187
214
}
188
215
"# ,
189
216
)
0 commit comments