From 2bde4e9efbd2a62ee51c64c9f991f9e48cd9aeda Mon Sep 17 00:00:00 2001
From: martinhath <martin@vind.ai>
Date: Thu, 29 Aug 2024 09:45:38 +0200
Subject: [PATCH] react_refresh: Don't visit arrow function nodes (#37)

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 | 57 ++++++++++++++++++++++++++++++++-
 1 file changed, 56 insertions(+), 1 deletion(-)

diff --git a/crates/react_refresh/src/lib.rs b/crates/react_refresh/src/lib.rs
index 0761320..38174da 100644
--- a/crates/react_refresh/src/lib.rs
+++ b/crates/react_refresh/src/lib.rs
@@ -152,6 +152,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 +257,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 +286,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 +734,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) };
+"#
+    );
 }