Improve docs for ReactRefreshTransformVisitor fields

crates/react_refresh/src/lib.rs

@@ -19,11 +19,17 @@ pub struct ReactRefreshTransformVisitor {
     atom_import_map: AtomImportMap,
     #[allow(dead_code)]
     file_name: FileName,
+    /// We're currently at the top level
     top_level: bool,
-
-    object_path: Vec<String>,
-    /// [true] if any atom was replaced.
+    /// 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<String>,
 }
 
 fn create_react_refresh_call_expr_(key: String, atom_expr: &CallExpr) -> CallExpr {
@@ -82,7 +88,7 @@ impl ReactRefreshTransformVisitor {
             file_name,
             top_level: false,
             used_atom: false,
-            object_path: Vec::new(),
+            access_path: Vec::new(),
         }
     }
 
@@ -91,9 +97,9 @@ impl ReactRefreshTransformVisitor {
             FileName::Real(ref real_file_name) => format!(
                 "{}/{}",
                 real_file_name.display(),
-                self.object_path.join(".")
+                self.access_path.join(".")
             ),
-            _ => self.object_path.join("."),
+            _ => self.access_path.join("."),
         }
     }
 }
@@ -105,6 +111,34 @@ impl VisitMut for ReactRefreshTransformVisitor {
         self.atom_import_map.visit_import_decl(import);
     }
 
+    fn visit_mut_module_items(&mut self, items: &mut Vec<ModuleItem>) {
+        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(),
+                get(name, inst) { 
+                  if (this.cache.has(name)) {
+                    return this.cache.get(name)
+                  }
+                  this.cache.set(name, inst)
+                  return inst
+                },
+              }" as Stmt
+            );
+            let mi: ModuleItem = jotai_cache_stmt.into();
+            items.insert(0, mi);
+        }
+    }
+
+    fn visit_mut_stmts(&mut self, stmts: &mut Vec<Stmt>) {
+        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) {
         if !self.top_level {
             return;
@@ -120,23 +154,26 @@ impl VisitMut for ReactRefreshTransformVisitor {
             "wat".to_string()
         };
 
-        self.object_path.push(key);
+        self.access_path.push(key);
         var_declarator.visit_mut_children_with(self);
-        self.object_path.pop();
+        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_array_lit(&mut self, array: &mut ArrayLit) {
         if !self.top_level {
             return;
         }
         for (i, child) in array.elems.iter_mut().enumerate() {
-            self.object_path.push(i.to_string());
+            self.access_path.push(i.to_string());
             child.visit_mut_with(self);
-            self.object_path.pop();
+            self.access_path.pop();
         }
     }
 
@@ -150,26 +187,24 @@ impl VisitMut for ReactRefreshTransformVisitor {
             match prop {
                 PropOrSpread::Prop(ref mut prop) => match prop.as_mut() {
                     Prop::Shorthand(ref mut s) => {
-                        self.object_path.push(s.sym.to_string());
+                        self.access_path.push(s.sym.to_string());
                         prop.visit_mut_with(self);
-                        self.object_path.pop();
+                        self.access_path.pop();
                     }
                     Prop::KeyValue(ref mut kv) => {
-                        self.object_path.push(show_prop_name(&kv.key));
+                        self.access_path.push(show_prop_name(&kv.key));
                         prop.visit_mut_with(self);
-                        self.object_path.pop();
+                        self.access_path.pop();
                     }
-                    // TODO: need to add in something here to avoid collisions
                     _ => prop.visit_mut_with(self),
                 },
-                // TODO: need to add in something here to avoid collisions
                 _ => prop.visit_mut_with(self),
             }
         }
     }
 
     fn visit_mut_call_expr(&mut self, call_expr: &mut CallExpr) {
-        // If this is an atom usage, replace it with the cached version.
+        // 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) {
@@ -182,34 +217,6 @@ impl VisitMut for ReactRefreshTransformVisitor {
         }
         call_expr.visit_mut_children_with(self);
     }
-
-    fn visit_mut_module_items(&mut self, items: &mut Vec<ModuleItem>) {
-        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(),
-                get(name, inst) { 
-                  if (this.cache.has(name)) {
-                    return this.cache.get(name)
-                  }
-                  this.cache.set(name, inst)
-                  return inst
-                },
-              }" as Stmt
-            );
-            let mi: ModuleItem = jotai_cache_stmt.into();
-            items.insert(0, mi);
-        }
-    }
-
-    fn visit_mut_stmts(&mut self, stmts: &mut Vec<Stmt>) {
-        let top_level = self.top_level;
-        self.top_level = false;
-        stmts.visit_mut_children_with(self);
-        self.top_level = top_level;
-    }
 }
 
 pub fn react_refresh(config: Config, file_name: FileName) -> impl Fold {
@@ -783,6 +790,41 @@ function keepThese() {
 "#
     );
 
+    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)),
+};
+"#
+    );
+
     test_inline!(
         Syntax::default(),
         |_| transform(None, Some(FileName::Anon)),