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

crates/react_refresh/src/lib.rs

@@ -93,8 +93,9 @@ impl ReactRefreshTransformVisitor {
     fn visit_mut_stmt_like<T>(&mut self, stmts: &mut Vec<T>)
     where
         Vec<T>: VisitMutWith<Self>,
-        T: VisitMutWith<Self> + StmtLike + ModuleItemLike,
+        T: VisitMutWith<Self> + StmtLike + ModuleItemLike + std::fmt::Debug,
     {
+        dbg!(&stmts, self.top_level);
         let mut stmts_updated: Vec<T> = Vec::with_capacity(stmts.len());
         let mut is_atom_present: bool = false;
 
@@ -152,6 +153,7 @@ impl ReactRefreshTransformVisitor {
                                 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()
                                     {
@@ -256,6 +258,10 @@ impl VisitMut for ReactRefreshTransformVisitor {
         self.current_var_declarator = old_var_declarator;
     }
 
+    fn visit_mut_arrow_expr(&mut self, _: &mut ArrowExpr) {
+        // Don't touch this sub-tree
+    }
+
     fn visit_mut_call_expr(&mut self, call_expr: &mut CallExpr) {
         if self.current_var_declarator.is_none() {
             return;
@@ -281,9 +287,10 @@ impl VisitMut for ReactRefreshTransformVisitor {
     }
 
     fn visit_mut_stmts(&mut self, stmts: &mut Vec<Stmt>) {
+        let top_level = self.top_level;
         self.top_level = false;
         self.visit_mut_stmt_like(stmts);
-        self.top_level = true;
+        self.top_level = top_level;
     }
 }
 
@@ -728,4 +735,53 @@ function createAtom(ov) {
 const value1Atom = createAtom('Hello String!');
 const countAtom = globalThis.jotaiAtomCache.get("atoms.ts/countAtom", atom(0));"#
     );
+
+    test_inline!(
+        Syntax::default(),
+        |_| transform(None, Some(FileName::Anon)),
+        nested_top_level_atoms,
+        r#"
+import { atom } from "jotai";
+
+const three = atom(atom(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";
+const three = globalThis.jotaiAtomCache.get("three", atom(atom(atom(0))));
+"#
+    );
+
+    test_inline!(
+        Syntax::default(),
+        |_| transform(None, Some(FileName::Anon)),
+        higher_order_fn_to_atom,
+        r#"
+import { atom } from "jotai";
+
+function getAtom() {
+    return atom(1);
+}
+const getAtom2 = () => atom(2);
+const getAtom3 = () => { return atom(3) };
+"#,
+        r#"
+import { atom } from "jotai";
+
+function getAtom() {
+    return atom(1);
+}
+const getAtom2 = () => atom(2);
+const getAtom3 = () => { return atom(3) };
+"#
+    );
 }