1
+ use ide_db:: helpers:: FamousDefs ;
2
+ use stdx:: format_to;
3
+ use syntax:: { AstNode , ast:: { self , ArgListOwner } } ;
4
+
5
+ use crate :: { AssistContext , AssistId , AssistKind , Assists } ;
6
+
7
+ /// Assist: convert_iter_for_each_to_for
8
+ //
9
+ /// Converts an Iterator::for_each function into a for loop.
10
+ ///
11
+ /// ```rust
12
+ /// fn main() {
13
+ /// let vec = vec![(1, 2), (2, 3), (3, 4)];
14
+ /// x.iter().for_each(|(x, y)| {
15
+ /// println!("x: {}, y: {}", x, y);
16
+ /// })
17
+ /// }
18
+ /// ```
19
+ /// ->
20
+ /// ```rust
21
+ /// fn main() {
22
+ /// let vec = vec![(1, 2), (2, 3), (3, 4)];
23
+ /// for (x, y) in x.iter() {
24
+ /// println!("x: {}, y: {}", x, y);
25
+ /// });
26
+ /// }
27
+ /// ```
28
+ pub ( crate ) fn convert_iter_for_each_to_for ( acc : & mut Assists , ctx : & AssistContext ) -> Option < ( ) > {
29
+ let closure;
30
+
31
+ let total_expr = match ctx. find_node_at_offset :: < ast:: Expr > ( ) ? {
32
+ ast:: Expr :: MethodCallExpr ( expr) => {
33
+ closure = match expr. arg_list ( ) ?. args ( ) . next ( ) ? {
34
+ ast:: Expr :: ClosureExpr ( expr) => expr,
35
+ _ => { return None ; }
36
+ } ;
37
+
38
+ expr
39
+ } ,
40
+ ast:: Expr :: ClosureExpr ( expr) => {
41
+ closure = expr;
42
+ ast:: MethodCallExpr :: cast ( closure. syntax ( ) . ancestors ( ) . nth ( 2 ) ?) ?
43
+ } ,
44
+ _ => { return None ; }
45
+ } ;
46
+
47
+ let ( total_expr, parent) = validate_method_call_expr ( & ctx. sema , total_expr) ?;
48
+
49
+ let param_list = closure. param_list ( ) ?;
50
+ let param = param_list. params ( ) . next ( ) ?;
51
+ let body = closure. body ( ) ?;
52
+
53
+ acc. add (
54
+ AssistId ( "convert_iter_for_each_to_for" , AssistKind :: RefactorRewrite ) ,
55
+ "Replace this `Iterator::for_each` with a for loop" ,
56
+ total_expr. syntax ( ) . text_range ( ) ,
57
+ |builder| {
58
+ let mut buf = String :: new ( ) ;
59
+
60
+ format_to ! ( buf, "for {} in {} " , param, parent) ;
61
+
62
+ match body {
63
+ ast:: Expr :: BlockExpr ( body) => format_to ! ( buf, "{}" , body) ,
64
+ _ => format_to ! ( buf, "{{\n {}\n }}" , body)
65
+ }
66
+
67
+ builder. replace ( total_expr. syntax ( ) . text_range ( ) , buf)
68
+ } ,
69
+ )
70
+ }
71
+
72
+ fn validate_method_call_expr (
73
+ sema : & hir:: Semantics < ide_db:: RootDatabase > ,
74
+ expr : ast:: MethodCallExpr ,
75
+ ) -> Option < ( ast:: Expr , ast:: Expr ) > {
76
+ if expr. name_ref ( ) ?. text ( ) != "for_each" {
77
+ return None ;
78
+ }
79
+
80
+ let expr = ast:: Expr :: MethodCallExpr ( expr) ;
81
+ let parent = ast:: Expr :: cast ( expr. syntax ( ) . first_child ( ) ?) ?;
82
+
83
+ let it_type = sema. type_of_expr ( & parent) ?;
84
+ let module = sema. scope ( parent. syntax ( ) ) . module ( ) ?;
85
+ let krate = module. krate ( ) ;
86
+
87
+ let iter_trait = FamousDefs ( sema, Some ( krate) ) . core_iter_Iterator ( ) ?;
88
+ it_type. impls_trait ( sema. db , iter_trait, & [ ] ) . then ( || ( expr, parent) )
89
+ }
90
+
91
+ #[ cfg( test) ]
92
+ mod tests {
93
+ use crate :: tests:: check_assist;
94
+
95
+ use super :: * ;
96
+
97
+ #[ test]
98
+ fn test_for_each_in_method ( ) {
99
+ check_assist (
100
+ convert_iter_for_each_to_for,
101
+ r"
102
+ fn main() {
103
+ let x = vec![(1, 1), (2, 2), (3, 3), (4, 4)];
104
+ x.iter().$0for_each(|(x, y)| {
105
+ dbg!(x, y)
106
+ });
107
+ }" ,
108
+ r"
109
+ fn main() {
110
+ let x = vec![(1, 1), (2, 2), (3, 3), (4, 4)];
111
+ for (x, y) in x.iter() {
112
+ dbg!(x, y)
113
+ };
114
+ }" ,
115
+ )
116
+ }
117
+
118
+ #[ test]
119
+ fn test_for_each_in_closure ( ) {
120
+ check_assist (
121
+ convert_iter_for_each_to_for,
122
+ r"
123
+ fn main() {
124
+ let x = vec![(1, 1), (2, 2), (3, 3), (4, 4)];
125
+ x.iter().for_each($0|(x, y)| {
126
+ dbg!(x, y)
127
+ });
128
+ }" ,
129
+ r"
130
+ fn main() {
131
+ let x = vec![(1, 1), (2, 2), (3, 3), (4, 4)];
132
+ for (x, y) in x.iter() {
133
+ dbg!(x, y)
134
+ };
135
+ }" ,
136
+ )
137
+ }
138
+ }
0 commit comments