11use crate :: { AssistContext , AssistId , Assists } ;
22
3- use ast:: LoopBodyOwner ;
3+ use ast:: { ElseBranch , Expr , LoopBodyOwner } ;
44use ra_fmt:: unwrap_trivial_block;
55use ra_syntax:: { ast, match_ast, AstNode , TextRange , T } ;
66
@@ -25,19 +25,11 @@ pub(crate) fn unwrap_block(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
2525 let l_curly_token = ctx. find_token_at_offset ( T ! [ '{' ] ) ?;
2626 let block = ast:: BlockExpr :: cast ( l_curly_token. parent ( ) ) ?;
2727 let parent = block. syntax ( ) . parent ( ) ?;
28+ let assist_id = AssistId ( "unwrap_block" ) ;
29+ let assist_label = "Unwrap block" ;
30+
2831 let ( expr, expr_to_unwrap) = match_ast ! {
2932 match parent {
30- ast:: IfExpr ( if_expr) => {
31- let expr_to_unwrap = if_expr. blocks( ) . find_map( |expr| extract_expr( ctx. frange. range, expr) ) ;
32- let expr_to_unwrap = expr_to_unwrap?;
33- // Find if we are in a else if block
34- let ancestor = if_expr. syntax( ) . parent( ) . and_then( ast:: IfExpr :: cast) ;
35-
36- match ancestor {
37- None => ( ast:: Expr :: IfExpr ( if_expr) , expr_to_unwrap) ,
38- Some ( ancestor) => ( ast:: Expr :: IfExpr ( ancestor) , expr_to_unwrap) ,
39- }
40- } ,
4133 ast:: ForExpr ( for_expr) => {
4234 let block_expr = for_expr. loop_body( ) ?;
4335 let expr_to_unwrap = extract_expr( ctx. frange. range, block_expr) ?;
@@ -53,27 +45,62 @@ pub(crate) fn unwrap_block(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
5345 let expr_to_unwrap = extract_expr( ctx. frange. range, block_expr) ?;
5446 ( ast:: Expr :: LoopExpr ( loop_expr) , expr_to_unwrap)
5547 } ,
48+ ast:: IfExpr ( if_expr) => {
49+ let mut resp = None ;
50+
51+ let then_branch = if_expr. then_branch( ) ?;
52+ if then_branch. l_curly_token( ) ?. text_range( ) . contains_range( ctx. frange. range) {
53+ if let Some ( ancestor) = if_expr. syntax( ) . parent( ) . and_then( ast:: IfExpr :: cast) {
54+ // For `else if` blocks
55+ let ancestor_then_branch = ancestor. then_branch( ) ?;
56+ let l_curly_token = then_branch. l_curly_token( ) ?;
57+
58+ let target = then_branch. syntax( ) . text_range( ) ;
59+ return acc. add( assist_id, assist_label, target, |edit| {
60+ let range_to_del_else_if = TextRange :: new( ancestor_then_branch. syntax( ) . text_range( ) . end( ) , l_curly_token. text_range( ) . start( ) ) ;
61+ let range_to_del_rest = TextRange :: new( then_branch. syntax( ) . text_range( ) . end( ) , if_expr. syntax( ) . text_range( ) . end( ) ) ;
62+
63+ edit. set_cursor( ancestor_then_branch. syntax( ) . text_range( ) . end( ) ) ;
64+ edit. delete( range_to_del_rest) ;
65+ edit. delete( range_to_del_else_if) ;
66+ edit. replace( target, update_expr_string( then_branch. to_string( ) , & [ ' ' , '{' ] ) ) ;
67+ } ) ;
68+ } else {
69+ resp = Some ( ( ast:: Expr :: IfExpr ( if_expr. clone( ) ) , Expr :: BlockExpr ( then_branch) ) ) ;
70+ }
71+ } else if let Some ( else_branch) = if_expr. else_branch( ) {
72+ match else_branch {
73+ ElseBranch :: Block ( else_block) => {
74+ let l_curly_token = else_block. l_curly_token( ) ?;
75+ if l_curly_token. text_range( ) . contains_range( ctx. frange. range) {
76+ let target = else_block. syntax( ) . text_range( ) ;
77+ return acc. add( assist_id, assist_label, target, |edit| {
78+ let range_to_del = TextRange :: new( then_branch. syntax( ) . text_range( ) . end( ) , l_curly_token. text_range( ) . start( ) ) ;
79+
80+ edit. set_cursor( then_branch. syntax( ) . text_range( ) . end( ) ) ;
81+ edit. delete( range_to_del) ;
82+ edit. replace( target, update_expr_string( else_block. to_string( ) , & [ ' ' , '{' ] ) ) ;
83+ } ) ;
84+ }
85+ } ,
86+ ElseBranch :: IfExpr ( _) => { } ,
87+ }
88+ }
89+
90+ resp?
91+ } ,
5692 _ => return None ,
5793 }
5894 } ;
5995
6096 let target = expr_to_unwrap. syntax ( ) . text_range ( ) ;
61- acc. add ( AssistId ( "unwrap_block" ) , "Unwrap block" , target, |edit| {
97+ acc. add ( assist_id , assist_label , target, |edit| {
6298 edit. set_cursor ( expr. syntax ( ) . text_range ( ) . start ( ) ) ;
6399
64- let pat_start: & [ _ ] = & [ ' ' , '{' , '\n' ] ;
65- let expr_to_unwrap = expr_to_unwrap. to_string ( ) ;
66- let expr_string = expr_to_unwrap. trim_start_matches ( pat_start) ;
67- let mut expr_string_lines: Vec < & str > = expr_string. lines ( ) . collect ( ) ;
68- expr_string_lines. pop ( ) ; // Delete last line
69-
70- let expr_string = expr_string_lines
71- . into_iter ( )
72- . map ( |line| line. replacen ( " " , "" , 1 ) ) // Delete indentation
73- . collect :: < Vec < String > > ( )
74- . join ( "\n " ) ;
75-
76- edit. replace ( expr. syntax ( ) . text_range ( ) , expr_string) ;
100+ edit. replace (
101+ expr. syntax ( ) . text_range ( ) ,
102+ update_expr_string ( expr_to_unwrap. to_string ( ) , & [ ' ' , '{' , '\n' ] ) ,
103+ ) ;
77104 } )
78105}
79106
@@ -87,6 +114,18 @@ fn extract_expr(cursor_range: TextRange, block: ast::BlockExpr) -> Option<ast::E
87114 }
88115}
89116
117+ fn update_expr_string ( expr_str : String , trim_start_pat : & [ char ] ) -> String {
118+ let expr_string = expr_str. trim_start_matches ( trim_start_pat) ;
119+ let mut expr_string_lines: Vec < & str > = expr_string. lines ( ) . collect ( ) ;
120+ expr_string_lines. pop ( ) ; // Delete last line
121+
122+ expr_string_lines
123+ . into_iter ( )
124+ . map ( |line| line. replacen ( " " , "" , 1 ) ) // Delete indentation
125+ . collect :: < Vec < String > > ( )
126+ . join ( "\n " )
127+ }
128+
90129#[ cfg( test) ]
91130mod tests {
92131 use crate :: tests:: { check_assist, check_assist_not_applicable} ;
@@ -142,7 +181,13 @@ mod tests {
142181 r#"
143182 fn main() {
144183 bar();
145- <|>println!("bar");
184+ if true {
185+ foo();
186+
187+ //comment
188+ bar();
189+ }<|>
190+ println!("bar");
146191 }
147192 "# ,
148193 ) ;
@@ -170,7 +215,127 @@ mod tests {
170215 r#"
171216 fn main() {
172217 //bar();
173- <|>println!("bar");
218+ if true {
219+ println!("true");
220+
221+ //comment
222+ //bar();
223+ }<|>
224+ println!("bar");
225+ }
226+ "# ,
227+ ) ;
228+ }
229+
230+ #[ test]
231+ fn simple_if_else_if_nested ( ) {
232+ check_assist (
233+ unwrap_block,
234+ r#"
235+ fn main() {
236+ //bar();
237+ if true {
238+ println!("true");
239+
240+ //comment
241+ //bar();
242+ } else if false {
243+ println!("bar");
244+ } else if true {<|>
245+ println!("foo");
246+ }
247+ }
248+ "# ,
249+ r#"
250+ fn main() {
251+ //bar();
252+ if true {
253+ println!("true");
254+
255+ //comment
256+ //bar();
257+ } else if false {
258+ println!("bar");
259+ }<|>
260+ println!("foo");
261+ }
262+ "# ,
263+ ) ;
264+ }
265+
266+ #[ test]
267+ fn simple_if_else_if_nested_else ( ) {
268+ check_assist (
269+ unwrap_block,
270+ r#"
271+ fn main() {
272+ //bar();
273+ if true {
274+ println!("true");
275+
276+ //comment
277+ //bar();
278+ } else if false {
279+ println!("bar");
280+ } else if true {
281+ println!("foo");
282+ } else {<|>
283+ println!("else");
284+ }
285+ }
286+ "# ,
287+ r#"
288+ fn main() {
289+ //bar();
290+ if true {
291+ println!("true");
292+
293+ //comment
294+ //bar();
295+ } else if false {
296+ println!("bar");
297+ } else if true {
298+ println!("foo");
299+ }<|>
300+ println!("else");
301+ }
302+ "# ,
303+ ) ;
304+ }
305+
306+ #[ test]
307+ fn simple_if_else_if_nested_middle ( ) {
308+ check_assist (
309+ unwrap_block,
310+ r#"
311+ fn main() {
312+ //bar();
313+ if true {
314+ println!("true");
315+
316+ //comment
317+ //bar();
318+ } else if false {
319+ println!("bar");
320+ } else if true {<|>
321+ println!("foo");
322+ } else {
323+ println!("else");
324+ }
325+ }
326+ "# ,
327+ r#"
328+ fn main() {
329+ //bar();
330+ if true {
331+ println!("true");
332+
333+ //comment
334+ //bar();
335+ } else if false {
336+ println!("bar");
337+ }<|>
338+ println!("foo");
174339 }
175340 "# ,
176341 ) ;
0 commit comments