From b608508ca31b2b20a77d5822735a86829ac8cdba Mon Sep 17 00:00:00 2001 From: martinhath Date: Tue, 3 Sep 2024 13:45:41 +0200 Subject: [PATCH] react_refresh: Rewrite main loop (#38) * react_refresh: Don't visit arrow function nodes Since atoms created in a function context are new atoms, there's nothing in the arrow function sub-tree that we should transform. The default behavior seems to be to visit everything, whereas an empty impl stops the walk. Another driveby fix: store previous `top_level` in `visit_mut_stmts` instead of always resetting it to `true`. I don't have a failing test for this (so it might be fine), but it's a pattern that's explicitly called out in the docs [0]. [0]: https://swc.rs/docs/plugin/ecmascript/cheatsheet#make-your-handlers-stateless Fixes #36 * react_refresh: Rewrite main loop This turned into an all-in-one commit: * Move the logic to the `visit_mut_*` functions from VisitMut. * Make atom key names more flexible to support e.g. #26 The largest change is the fact that we perform most of the work directly in the visitor functions. For applicable calls to `atom` (and friends) we replace the `CallExpr` right then and there. This makes it so that we don't need to deal with `export` handling, or accidentally replace entire expressions with erroneous cached atoms (See e.g. the failing test in #36). This also solves other open issues, like #26, #34, and #35. Fixes #26 Fixes #34 Fixes #35 * Add test case for #35 * Improve docs for ReactRefreshTransformVisitor fields * Replace "wat" placeholder --- crates/react_refresh/src/lib.rs | 447 ++++++++++++++++++-------------- 1 file changed, 259 insertions(+), 188 deletions(-) diff --git a/crates/react_refresh/src/lib.rs b/crates/react_refresh/src/lib.rs index 38174da..dd09c7d 100644 --- a/crates/react_refresh/src/lib.rs +++ b/crates/react_refresh/src/lib.rs @@ -2,12 +2,10 @@ use common::{parse_plugin_config, AtomImportMap, Config}; use swc_core::{ + common::FileName, common::DUMMY_SP, - common::{util::take::Take, FileName}, ecma::{ ast::*, - atoms::JsWord, - utils::{ModuleItemLike, StmtLike}, visit::{as_folder, noop_visit_mut_type, Fold, FoldWith, VisitMut, VisitMutWith}, }, plugin::{ @@ -19,16 +17,23 @@ use swc_core::{ pub struct ReactRefreshTransformVisitor { atom_import_map: AtomImportMap, - current_var_declarator: Option, - refresh_atom_var_decl: Option, #[allow(dead_code)] file_name: FileName, - exporting: bool, + /// We're currently at the top level top_level: bool, + /// Any atom was used. + used_atom: bool, + /// Path to the current expression when walking object and array literals. + /// For instance, when walking this expression: + /// ```js + /// const foo = [{}, { bar: [ 123 ]}] + /// ``` + /// the path will be `["foo", "1", "bar", "0"]` when visiting `123`. + access_path: Vec, } -fn create_react_refresh_call_expr(key: String, atom_expr: &CallExpr) -> Box { - Box::new(Expr::Call(CallExpr { +fn create_react_refresh_call_expr_(key: String, atom_expr: &CallExpr) -> CallExpr { + CallExpr { span: DUMMY_SP, callee: Callee::Expr(Box::new(Expr::Member(MemberExpr { span: DUMMY_SP, @@ -54,27 +59,25 @@ fn create_react_refresh_call_expr(key: String, atom_expr: &CallExpr) -> Box VarDeclarator { - let atom_name = atom_name_id.0.clone(); - VarDeclarator { - name: Pat::Ident(Ident::new(atom_name, DUMMY_SP.with_ctxt(atom_name_id.1)).into()), - span: DUMMY_SP.with_ctxt(atom_name_id.1), - init: Some(create_react_refresh_call_expr(key, atom_expr)), - definite: false, } } -fn create_cache_key(atom_name: &JsWord, file_name: &FileName) -> String { - match file_name { - FileName::Real(real_file_name) => format!("{}/{}", real_file_name.display(), atom_name), - _ => atom_name.to_string(), +fn show_prop_name(pn: &PropName) -> String { + use PropName::*; + match pn { + Ident(ref i) => i.sym.to_string(), + Str(ref s) => s.value.to_string(), + Num(ref n) => n + .raw + .as_ref() + .expect("Num(c).raw should be Some") + .to_string(), + Computed(ref c) => format!("computed:{:?}", c.span), + BigInt(ref b) => b + .raw + .as_ref() + .expect("BigInt(b).raw should be Some") + .to_string(), } } @@ -82,130 +85,36 @@ impl ReactRefreshTransformVisitor { pub fn new(config: Config, file_name: FileName) -> Self { Self { atom_import_map: AtomImportMap::new(config.atom_names), - current_var_declarator: None, - refresh_atom_var_decl: None, file_name, - exporting: false, top_level: false, + used_atom: false, + access_path: Vec::new(), } } - fn visit_mut_stmt_like(&mut self, stmts: &mut Vec) - where - Vec: VisitMutWith, - T: VisitMutWith + StmtLike + ModuleItemLike, - { - let mut stmts_updated: Vec = Vec::with_capacity(stmts.len()); - let mut is_atom_present: bool = false; - - for stmt in stmts.take() { - let exporting_old = self.exporting; - let stmt = match stmt.try_into_stmt() { - Ok(mut stmt) => { - stmt.visit_mut_with(self); - ::from_stmt(stmt) - } - Err(node) => match node.try_into_module_decl() { - Ok(mut module_decl) => { - match module_decl { - ModuleDecl::ExportDefaultExpr(mut default_export) => { - if !self.atom_import_map.is_atom_import(&default_export.expr) { - default_export.visit_mut_with(self); - stmts_updated.push( - ::try_from_module_decl( - default_export.into(), - ) - .unwrap(), - ); - continue; - } - is_atom_present = true; - - let atom_name: JsWord = match &self.file_name { - FileName::Real(real_file_name) => { - real_file_name.file_stem().map_or_else( - || "default_atom".into(), - |real_file_name| { - real_file_name.to_string_lossy().into() - }, - ) - } - _ => "default_atom".into(), - }; - - // export default expression - stmts_updated.push( - ::try_from_module_decl( - ModuleDecl::ExportDefaultExpr(ExportDefaultExpr { - expr: create_react_refresh_call_expr( - create_cache_key(&atom_name, &self.file_name), - default_export.expr.as_call().unwrap(), - ), - span: DUMMY_SP, - }), - ) - .unwrap(), - ); - continue; - } - ModuleDecl::ExportDecl(mut export_decl) => { - if let Decl::Var(mut var_decl) = export_decl.decl.clone() { - if let [VarDeclarator { - init: Some(init_expr), - // TODO: handle remaining expressions too. See #35. - .. - }] = var_decl.decls.as_mut_slice() - { - if self.atom_import_map.is_atom_import(&*init_expr) { - self.exporting = true; - } - } - } - export_decl.visit_mut_with(self); - ::try_from_module_decl(export_decl.into()) - .unwrap() - } - _ => { - module_decl.visit_mut_with(self); - ::try_from_module_decl(module_decl).unwrap() - } - } - } - Err(..) => unreachable!(), - }, - }; - - if self.refresh_atom_var_decl.is_none() { - stmts_updated.push(stmt); - continue; - } + fn create_cache_key(&self) -> String { + match self.file_name { + FileName::Real(ref real_file_name) => format!( + "{}/{}", + real_file_name.display(), + self.access_path.join(".") + ), + _ => self.access_path.join("."), + } + } +} - is_atom_present = true; +impl VisitMut for ReactRefreshTransformVisitor { + noop_visit_mut_type!(); - let updated_decl = Decl::Var(Box::new(VarDecl { - span: DUMMY_SP, - kind: VarDeclKind::Const, - declare: false, - decls: vec![self.refresh_atom_var_decl.take().unwrap()], - })); - - if self.exporting { - stmts_updated.push( - ::try_from_module_decl(ModuleDecl::ExportDecl( - ExportDecl { - span: DUMMY_SP, - decl: updated_decl, - }, - )) - .unwrap(), - ) - } else { - stmts_updated.push(::from_stmt(Stmt::Decl(updated_decl))) - } - self.exporting = exporting_old; - } + fn visit_mut_import_decl(&mut self, import: &mut ImportDecl) { + self.atom_import_map.visit_import_decl(import); + } - if is_atom_present { + fn visit_mut_module_items(&mut self, items: &mut Vec) { + self.top_level = true; + items.visit_mut_children_with(self); + if self.used_atom { let jotai_cache_stmt = quote!( "globalThis.jotaiAtomCache = globalThis.jotaiAtomCache || { cache: new Map(), @@ -218,21 +127,16 @@ impl ReactRefreshTransformVisitor { }, }" as Stmt ); - let mut stmts_with_cache: Vec = Vec::with_capacity(stmts_updated.len() + 1); - stmts_with_cache.push(::from_stmt(jotai_cache_stmt)); - stmts_with_cache.append(&mut stmts_updated); - stmts_updated = stmts_with_cache + let mi: ModuleItem = jotai_cache_stmt.into(); + items.insert(0, mi); } - - *stmts = stmts_updated; } -} - -impl VisitMut for ReactRefreshTransformVisitor { - noop_visit_mut_type!(); - fn visit_mut_import_decl(&mut self, import: &mut ImportDecl) { - self.atom_import_map.visit_import_decl(import); + fn visit_mut_stmts(&mut self, stmts: &mut Vec) { + let top_level = self.top_level; + self.top_level = false; + stmts.visit_mut_children_with(self); + self.top_level = top_level; } fn visit_mut_var_declarator(&mut self, var_declarator: &mut VarDeclarator) { @@ -240,56 +144,78 @@ impl VisitMut for ReactRefreshTransformVisitor { return; } - let old_var_declarator = self.current_var_declarator.take(); - - self.current_var_declarator = if let Pat::Ident(BindingIdent { - id: Ident { span, sym, .. }, + let key = if let Pat::Ident(BindingIdent { + id: Ident { sym, .. }, .. }) = &var_declarator.name { - Some((sym.clone(), span.ctxt)) + sym.to_string() } else { - None + "[missing-declarator]".to_string() }; + self.access_path.push(key); var_declarator.visit_mut_children_with(self); - - self.current_var_declarator = old_var_declarator; + self.access_path.pop(); } fn visit_mut_arrow_expr(&mut self, _: &mut ArrowExpr) { - // Don't touch this sub-tree + // Arrow expressions is (maybe) the only way for expressions to not be at the top level and + // not have us visit `stmts` before. Since we record whether we're on the top level in + // `visit_mut_stmts`, we need to make sure we don't visit the body here, so that any atoms + // aren't erroneously cached. } - fn visit_mut_call_expr(&mut self, call_expr: &mut CallExpr) { - if self.current_var_declarator.is_none() { + fn visit_mut_array_lit(&mut self, array: &mut ArrayLit) { + if !self.top_level { return; } - - call_expr.visit_mut_children_with(self); - - let atom_name = self.current_var_declarator.as_ref().unwrap(); - if let Callee::Expr(expr) = &call_expr.callee { - if self.atom_import_map.is_atom_import(expr) { - self.refresh_atom_var_decl = Some(create_react_refresh_var_decl( - atom_name.clone(), - create_cache_key(&atom_name.0, &self.file_name), - call_expr, - )) - } + for (i, child) in array.elems.iter_mut().enumerate() { + self.access_path.push(i.to_string()); + child.visit_mut_with(self); + self.access_path.pop(); } } - fn visit_mut_module_items(&mut self, items: &mut Vec) { - self.top_level = true; - self.visit_mut_stmt_like(items); + fn visit_mut_object_lit(&mut self, object: &mut ObjectLit) { + if !self.top_level { + return; + } + // For each prop in the object we need to record the path down to build up the ind-path + // down to any atoms in the literal. + for prop in object.props.iter_mut() { + match prop { + PropOrSpread::Prop(ref mut prop) => match prop.as_mut() { + Prop::Shorthand(ref mut s) => { + self.access_path.push(s.sym.to_string()); + prop.visit_mut_with(self); + self.access_path.pop(); + } + Prop::KeyValue(ref mut kv) => { + self.access_path.push(show_prop_name(&kv.key)); + prop.visit_mut_with(self); + self.access_path.pop(); + } + _ => prop.visit_mut_with(self), + }, + _ => prop.visit_mut_with(self), + } + } } - fn visit_mut_stmts(&mut self, stmts: &mut Vec) { - let top_level = self.top_level; - self.top_level = false; - self.visit_mut_stmt_like(stmts); - self.top_level = top_level; + fn visit_mut_call_expr(&mut self, call_expr: &mut CallExpr) { + // If this is an atom, replace it with the cached `get` expression. + if self.top_level { + if let Callee::Expr(expr) = &call_expr.callee { + if self.atom_import_map.is_atom_import(expr) { + *call_expr = + create_react_refresh_call_expr_(self.create_cache_key(), call_expr); + self.used_atom = true; + return; + } + } + } + call_expr.visit_mut_children_with(self); } } @@ -526,7 +452,7 @@ globalThis.jotaiAtomCache = globalThis.jotaiAtomCache || { }, } import { atom } from "jotai"; -export default globalThis.jotaiAtomCache.get("atoms.ts/atoms", atom(0));"# +export default globalThis.jotaiAtomCache.get("atoms.ts/", atom(0));"# ); test_inline!( @@ -548,7 +474,7 @@ globalThis.jotaiAtomCache = globalThis.jotaiAtomCache || { }, } import { atom } from "jotai"; -export default globalThis.jotaiAtomCache.get("countAtom.ts/countAtom", atom(0));"# +export default globalThis.jotaiAtomCache.get("countAtom.ts/", atom(0));"# ); test_inline!( @@ -573,7 +499,7 @@ globalThis.jotaiAtomCache = globalThis.jotaiAtomCache || { }, } import { atom } from "jotai"; -export default globalThis.jotaiAtomCache.get("src/atoms/countAtom.ts/countAtom", atom(0));"# +export default globalThis.jotaiAtomCache.get("src/atoms/countAtom.ts/", atom(0));"# ); test_inline!( @@ -781,6 +707,151 @@ function getAtom() { } const getAtom2 = () => atom(2); const getAtom3 = () => { return atom(3) }; +"# + ); + + test_inline!( + Syntax::default(), + |_| transform(None, Some(FileName::Anon)), + atom_in_atom_reader_stmt, + r#" +import { atom } from "jotai"; + +export const state = atom(() => { + return atom(0); +});"#, + r#" +globalThis.jotaiAtomCache = globalThis.jotaiAtomCache || { + cache: new Map(), + get(name, inst) { + if (this.cache.has(name)) { + return this.cache.get(name) + } + this.cache.set(name, inst) + return inst + }, +} +import { atom } from "jotai"; + +export const state = globalThis.jotaiAtomCache.get("state", atom(() => { + return atom(0); +}));"# + ); + + test_inline!( + Syntax::default(), + |_| transform(None, Some(FileName::Anon)), + array_and_object_top_level, + r#" +import { atom } from "jotai"; + +const arr = [ + atom(3), + atom(4), +]; + +const obj = { + five: atom(5), + six: atom(6), +}; + +function keepThese() { + const a = [atom(7)]; + const b = { eight: atom(8) }; +} +"#, + r#" +globalThis.jotaiAtomCache = globalThis.jotaiAtomCache || { + cache: new Map(), + get(name, inst) { + if (this.cache.has(name)) { + return this.cache.get(name) + } + this.cache.set(name, inst) + return inst + }, +} +import { atom } from "jotai"; + +const arr = [ + globalThis.jotaiAtomCache.get("arr.0", atom(3)), + globalThis.jotaiAtomCache.get("arr.1", atom(4)), +]; + +const obj = { + five: globalThis.jotaiAtomCache.get("obj.five", atom(5)), + six: globalThis.jotaiAtomCache.get("obj.six", atom(6)), +}; + +function keepThese() { + const a = [atom(7)]; + const b = { eight: atom(8) }; +} +"# + ); + + test_inline!( + Syntax::default(), + |_| transform(None, Some(FileName::Anon)), + object_edge_cases, + r#" +import { atom } from "jotai"; + +const obj = { + five: atom(5), + six: atom(6), + ...({ + six: atom(66), + }) +}; +"#, + r#" +globalThis.jotaiAtomCache = globalThis.jotaiAtomCache || { + cache: new Map(), + get(name, inst) { + if (this.cache.has(name)) { + return this.cache.get(name) + } + this.cache.set(name, inst) + return inst + }, +} +import { atom } from "jotai"; + +const obj = { + five: globalThis.jotaiAtomCache.get("obj.five", atom(5)), + six: globalThis.jotaiAtomCache.get("obj.six", atom(6)), + ...{ + six: globalThis.jotaiAtomCache.get("obj.six", atom(66)), + } +}; +"# + ); + + test_inline!( + Syntax::default(), + |_| transform(None, Some(FileName::Anon)), + compound_export, + r#" +import { atom } from "jotai"; + +export const one = atom(1), + two = atom(2); +"#, + r#" +globalThis.jotaiAtomCache = globalThis.jotaiAtomCache || { + cache: new Map(), + get(name, inst) { + if (this.cache.has(name)) { + return this.cache.get(name) + } + this.cache.set(name, inst) + return inst + }, +} +import { atom } from "jotai"; + +export const one = globalThis.jotaiAtomCache.get("one", atom(1)), two = globalThis.jotaiAtomCache.get("two", atom(2)); "# ); }