@@ -5,7 +5,7 @@ use hir::Semantics;
55use  syntax:: { 
66    algo, 
77    ast:: { self ,  make,  AstNode ,  AttrsOwner ,  ModuleItemOwner ,  PathSegmentKind ,  VisibilityOwner } , 
8-     ted,  AstToken ,  Direction ,  NodeOrToken ,  SyntaxNode ,  SyntaxToken , 
8+     match_ast ,   ted,  AstToken ,  Direction ,  NodeOrToken ,  SyntaxNode ,  SyntaxToken , 
99} ; 
1010
1111use  crate :: { 
@@ -43,16 +43,32 @@ pub struct InsertUseConfig {
4343pub  enum  ImportScope  { 
4444    File ( ast:: SourceFile ) , 
4545    Module ( ast:: ItemList ) , 
46+     Block ( ast:: BlockExpr ) , 
4647} 
4748
4849impl  ImportScope  { 
49-     pub  fn  from ( syntax :  SyntaxNode )  -> Option < Self >  { 
50-         if  let  Some ( module)  = ast:: Module :: cast ( syntax. clone ( ) )  { 
51-             module. item_list ( ) . map ( ImportScope :: Module ) 
52-         }  else  if  let  this @ Some ( _)  = ast:: SourceFile :: cast ( syntax. clone ( ) )  { 
53-             this. map ( ImportScope :: File ) 
54-         }  else  { 
55-             ast:: ItemList :: cast ( syntax) . map ( ImportScope :: Module ) 
50+     fn  from ( syntax :  SyntaxNode )  -> Option < Self >  { 
51+         fn  contains_cfg_attr ( attrs :  & dyn  AttrsOwner )  -> bool  { 
52+             attrs
53+                 . attrs ( ) 
54+                 . any ( |attr| attr. as_simple_call ( ) . map_or ( false ,  |( ident,  _) | ident == "cfg" ) ) 
55+         } 
56+         match_ast !  { 
57+             match  syntax { 
58+                 ast:: Module ( module)  => module. item_list( ) . map( ImportScope :: Module ) , 
59+                 ast:: SourceFile ( file)  => Some ( ImportScope :: File ( file) ) , 
60+                 ast:: Fn ( func)  => contains_cfg_attr( & func) . then( || func. body( ) . map( ImportScope :: Block ) ) . flatten( ) , 
61+                 ast:: Const ( konst)  => contains_cfg_attr( & konst) . then( || match  konst. body( ) ? { 
62+                     ast:: Expr :: BlockExpr ( block)  => Some ( block) , 
63+                     _ => None , 
64+                 } ) . flatten( ) . map( ImportScope :: Block ) , 
65+                 ast:: Static ( statik)  => contains_cfg_attr( & statik) . then( || match  statik. body( ) ? { 
66+                     ast:: Expr :: BlockExpr ( block)  => Some ( block) , 
67+                     _ => None , 
68+                 } ) . flatten( ) . map( ImportScope :: Block ) , 
69+                 _ => None , 
70+ 
71+             } 
5672        } 
5773    } 
5874
@@ -73,13 +89,15 @@ impl ImportScope {
7389        match  self  { 
7490            ImportScope :: File ( file)  => file. syntax ( ) , 
7591            ImportScope :: Module ( item_list)  => item_list. syntax ( ) , 
92+             ImportScope :: Block ( block)  => block. syntax ( ) , 
7693        } 
7794    } 
7895
7996    pub  fn  clone_for_update ( & self )  -> Self  { 
8097        match  self  { 
8198            ImportScope :: File ( file)  => ImportScope :: File ( file. clone_for_update ( ) ) , 
8299            ImportScope :: Module ( item_list)  => ImportScope :: Module ( item_list. clone_for_update ( ) ) , 
100+             ImportScope :: Block ( block)  => ImportScope :: Block ( block. clone_for_update ( ) ) , 
83101        } 
84102    } 
85103
@@ -96,6 +114,7 @@ impl ImportScope {
96114        let  mut  use_stmts = match  self  { 
97115            ImportScope :: File ( f)  => f. items ( ) , 
98116            ImportScope :: Module ( m)  => m. items ( ) , 
117+             ImportScope :: Block ( b)  => b. items ( ) , 
99118        } 
100119        . filter_map ( use_stmt) ; 
101120        let  mut  res = ImportGranularityGuess :: Unknown ; 
@@ -319,28 +338,29 @@ fn insert_use_(
319338        ted:: insert ( ted:: Position :: after ( last_inner_element) ,  make:: tokens:: single_newline ( ) ) ; 
320339        return ; 
321340    } 
322-     match  scope { 
341+     let  l_curly =  match  scope { 
323342        ImportScope :: File ( _)  => { 
324343            cov_mark:: hit!( insert_group_empty_file) ; 
325344            ted:: insert ( ted:: Position :: first_child_of ( scope_syntax) ,  make:: tokens:: blank_line ( ) ) ; 
326-             ted:: insert ( ted:: Position :: first_child_of ( scope_syntax) ,  use_item. syntax ( ) ) 
345+             ted:: insert ( ted:: Position :: first_child_of ( scope_syntax) ,  use_item. syntax ( ) ) ; 
346+             return ; 
327347        } 
348+         // don't insert the imports before the item list/block expr's opening curly brace 
349+         ImportScope :: Module ( item_list)  => item_list. l_curly_token ( ) , 
328350        // don't insert the imports before the item list's opening curly brace 
329-         ImportScope :: Module ( item_list)  => match  item_list. l_curly_token ( )  { 
330-             Some ( b)  => { 
331-                 cov_mark:: hit!( insert_group_empty_module) ; 
332-                 ted:: insert ( ted:: Position :: after ( & b) ,  make:: tokens:: single_newline ( ) ) ; 
333-                 ted:: insert ( ted:: Position :: after ( & b) ,  use_item. syntax ( ) ) ; 
334-             } 
335-             None  => { 
336-                 // This should never happens, broken module syntax node 
337-                 ted:: insert ( 
338-                     ted:: Position :: first_child_of ( scope_syntax) , 
339-                     make:: tokens:: blank_line ( ) , 
340-                 ) ; 
341-                 ted:: insert ( ted:: Position :: first_child_of ( scope_syntax) ,  use_item. syntax ( ) ) ; 
342-             } 
343-         } , 
351+         ImportScope :: Block ( block)  => block. l_curly_token ( ) , 
352+     } ; 
353+     match  l_curly { 
354+         Some ( b)  => { 
355+             cov_mark:: hit!( insert_group_empty_module) ; 
356+             ted:: insert ( ted:: Position :: after ( & b) ,  make:: tokens:: single_newline ( ) ) ; 
357+             ted:: insert ( ted:: Position :: after ( & b) ,  use_item. syntax ( ) ) ; 
358+         } 
359+         None  => { 
360+             // This should never happens, broken module syntax node 
361+             ted:: insert ( ted:: Position :: first_child_of ( scope_syntax) ,  make:: tokens:: blank_line ( ) ) ; 
362+             ted:: insert ( ted:: Position :: first_child_of ( scope_syntax) ,  use_item. syntax ( ) ) ; 
363+         } 
344364    } 
345365} 
346366
0 commit comments