@@ -19,40 +19,74 @@ pub fn jest() -> impl Pass {
1919#[ derive( Default ) ]
2020struct Jest {
2121 imported : Vec < Id > ,
22+ mock_vars : Vec < Id > ,
2223}
2324
2425impl Jest {
2526 fn visit_mut_stmt_like < T > ( & mut self , orig : & mut Vec < T > )
2627 where
2728 T : StmtLike + VisitMutWith < Self > ,
2829 {
30+ // First pass to collect mock variable declarations
2931 for item in & mut * orig {
32+ // Check for declarations with identifiers starting with "mock"
33+ if let Some ( Stmt :: Decl ( Decl :: Var ( var_decl) ) ) = item. as_stmt ( ) {
34+ for decl in & var_decl. decls {
35+ if let Pat :: Ident ( ident) = & decl. name {
36+ let name = & * ident. id . sym ;
37+ if name. starts_with ( "mock" ) {
38+ self . mock_vars . push ( ident. id . to_id ( ) ) ;
39+ }
40+ }
41+ }
42+ }
43+
44+ // Standard visitation
3045 item. visit_mut_with ( self ) ;
3146 }
3247
3348 let items = orig. take ( ) ;
3449
3550 let mut new = Vec :: with_capacity ( items. len ( ) ) ;
3651 let mut hoisted = Vec :: with_capacity ( 8 ) ;
52+
3753 items. into_iter ( ) . for_each ( |item| {
3854 match item. try_into_stmt ( ) {
39- Ok ( stmt) => match & stmt {
40- Stmt :: Expr ( ExprStmt { expr, .. } ) => match & * * expr {
41- Expr :: Call ( CallExpr {
55+ Ok ( stmt) => {
56+ // Check for mock variable declarations to hoist
57+ if let Stmt :: Decl ( Decl :: Var ( var_decl) ) = & stmt {
58+ let mut should_hoist = false ;
59+ for decl in & var_decl. decls {
60+ if let Pat :: Ident ( ident) = & decl. name {
61+ if self . mock_vars . iter ( ) . any ( |id| * id == ident. id . to_id ( ) ) {
62+ should_hoist = true ;
63+ break ;
64+ }
65+ }
66+ }
67+
68+ if should_hoist {
69+ hoisted. push ( T :: from ( stmt) ) ;
70+ return ;
71+ }
72+ }
73+
74+ // Check for jest calls to hoist
75+ if let Stmt :: Expr ( ExprStmt { expr, .. } ) = & stmt {
76+ if let Expr :: Call ( CallExpr {
4277 callee : Callee :: Expr ( callee) ,
4378 ..
44- } ) => {
79+ } ) = & * * expr
80+ {
4581 if self . should_hoist ( callee) {
46- hoisted. push ( T :: from ( stmt) )
47- } else {
48- new. push ( T :: from ( stmt) )
82+ hoisted. push ( T :: from ( stmt) ) ;
83+ return ;
4984 }
5085 }
51- _ => new. push ( T :: from ( stmt) ) ,
52- } ,
86+ }
5387
54- _ => new. push ( T :: from ( stmt) ) ,
55- } ,
88+ new. push ( T :: from ( stmt) ) ;
89+ }
5690 Err ( node) => new. push ( node) ,
5791 } ;
5892 } ) ;
@@ -82,6 +116,7 @@ impl VisitMut for Jest {
82116 noop_visit_mut_type ! ( ) ;
83117
84118 fn visit_mut_module_items ( & mut self , items : & mut Vec < ModuleItem > ) {
119+ // First collect imported jest methods
85120 for item in items. iter ( ) {
86121 if let ModuleItem :: ModuleDecl ( ModuleDecl :: Import ( ImportDecl {
87122 specifiers, src, ..
@@ -118,6 +153,7 @@ impl VisitMut for Jest {
118153 }
119154 }
120155
156+ // Now process and hoist as needed
121157 self . visit_mut_stmt_like ( items)
122158 }
123159
0 commit comments