React Refresh support (#782)

* Setup code for React refresh plugin

* Finish initial impl and first test

* Add plugin to preset
Setup build for plugin

* Add validation
Cleanup

* Add more tests
Small fix for default exports

* Add preset tests

* Add preset tests

* Update Parcel CSB

.codesandbox/ci.json

@@ -7,7 +7,7 @@
     "simple-snowpack-react-o1gmx",
     "next-js-uo1h0",
     "next-js-with-custom-babel-config-komw9",
-    "react-with-custom-babel-config-2dpyh"
+    "react-with-custom-babel-config-z1ebx"
   ],
   "node": "12"
 }

package.json

@@ -89,6 +89,12 @@
       "import": "./esm/babel/plugin-debug-label.mjs",
       "default": "./babel/plugin-debug-label.js"
     },
+    "./babel/plugin-react-refresh": {
+      "types": "./babel/plugin-react-refresh.d.ts",
+      "module": "./esm/babel/plugin-react-refresh.js",
+      "import": "./esm/babel/plugin-react-refresh.mjs",
+      "default": "./babel/plugin-react-refresh.js"
+    },
     "./babel/preset": {
       "types": "./babel/preset.d.ts",
       "module": "./esm/babel/preset.js",
@@ -115,6 +121,7 @@
     "build:redux": "rollup -c --config-redux",
     "build:urql": "rollup -c --config-urql",
     "build:babel:plugin-debug-label": "rollup -c --config-babel_plugin-debug-label",
+    "build:babel:plugin-react-refresh": "rollup -c --config-babel_plugin-react-refresh",
     "build:babel:preset": "rollup -c --config-babel_preset",
     "postbuild": "yarn copy",
     "eslint": "eslint --fix '*.{js,json}' '{src,tests,benchmarks}/**/*.{ts,tsx}'",

src/babel/plugin-react-refresh.ts

@@ -0,0 +1,72 @@
+import babel, { PluginObj } from '@babel/core'
+import templateBuilder from '@babel/template'
+import { isAtom } from './utils'
+
+export default function reactRefreshPlugin({
+  types: t,
+}: typeof babel): PluginObj {
+  return {
+    pre({ opts }) {
+      if (!opts.filename) {
+        throw new Error('Filename must be available')
+      }
+    },
+    visitor: {
+      Program: {
+        exit(path) {
+          const jotaiAtomCache = templateBuilder(`
+          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
+            },
+          }`)()
+          path.unshiftContainer('body', jotaiAtomCache)
+        },
+      },
+      ExportDefaultDeclaration(nodePath, state) {
+        const { node } = nodePath
+        if (
+          t.isCallExpression(node.declaration) &&
+          isAtom(t, node.declaration.callee)
+        ) {
+          const filename = state.filename || 'unknown'
+          const atomKey = `${filename}/defaultExport`
+
+          const buildExport = templateBuilder(
+            `export default globalThis.jotaiAtomCache.get(%%atomKey%%, %%atom%%)`
+          )
+          const ast = buildExport({
+            atomKey: t.stringLiteral(atomKey),
+            atom: node.declaration,
+          })
+          nodePath.replaceWith(ast as babel.Node)
+        }
+      },
+      VariableDeclarator(nodePath, state) {
+        if (
+          t.isIdentifier(nodePath.node.id) &&
+          t.isCallExpression(nodePath.node.init) &&
+          isAtom(t, nodePath.node.init.callee)
+        ) {
+          const filename = state.filename || 'unknown'
+          const atomKey = `${filename}/${nodePath.node.id.name}`
+
+          const buildAtomDeclaration = templateBuilder(
+            `const %%atomIdentifier%% = globalThis.jotaiAtomCache.get(%%atomKey%%, %%atom%%)`
+          )
+          const ast = buildAtomDeclaration({
+            atomIdentifier: t.identifier(nodePath.node.id.name),
+            atomKey: t.stringLiteral(atomKey),
+            atom: nodePath.node.init,
+          })
+          nodePath.parentPath.replaceWith(ast as babel.Node)
+        }
+      },
+    },
+  }
+}

src/babel/preset.ts

@@ -1,8 +1,9 @@
 import babel from '@babel/core'
 import pluginDebugLabel from './plugin-debug-label'
+import pluginReactRefresh from './plugin-react-refresh'
 
 export default function jotaiPreset(): { plugins: babel.PluginItem[] } {
   return {
-    plugins: [pluginDebugLabel],
+    plugins: [pluginDebugLabel, pluginReactRefresh],
   }
 }

src/babel/utils.ts

@@ -0,0 +1,18 @@
+import { types } from '@babel/core'
+
+export function isAtom(
+  t: typeof types,
+  callee: babel.types.Expression | babel.types.V8IntrinsicIdentifier
+) {
+  if (t.isIdentifier(callee) && callee.name === 'atom') {
+    return true
+  }
+
+  if (t.isMemberExpression(callee)) {
+    const { property } = callee
+    if (t.isIdentifier(property) && property.name === 'atom') {
+      return true
+    }
+  }
+  return false
+}

tests/babel/plugin-react-refresh.test.ts

@@ -0,0 +1,143 @@
+import path from 'path'
+import { transformSync } from '@babel/core'
+
+const plugin = path.join(__dirname, '../../src/babel/plugin-react-refresh')
+
+const transform = (code: string, filename?: string) =>
+  transformSync(code, {
+    babelrc: false,
+    configFile: false,
+    filename,
+    root: '.',
+    plugins: [[plugin]],
+  })?.code
+
+it('Should add a cache for a single atom', () => {
+  expect(transform(`const countAtom = atom(0);`, '/src/atoms/index.ts'))
+    .toMatchInlineSnapshot(`
+    "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;
+      }
+
+    };
+    const countAtom = globalThis.jotaiAtomCache.get(\\"/src/atoms/index.ts/countAtom\\", atom(0));"
+  `)
+})
+
+it('Should add a cache for multiple atoms', () => {
+  expect(
+    transform(
+      `
+  const countAtom = atom(0);
+  const doubleAtom = atom((get) => get(countAtom) * 2);
+  `,
+      '/src/atoms/index.ts'
+    )
+  ).toMatchInlineSnapshot(`
+    "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;
+      }
+
+    };
+    const countAtom = globalThis.jotaiAtomCache.get(\\"/src/atoms/index.ts/countAtom\\", atom(0));
+    const doubleAtom = globalThis.jotaiAtomCache.get(\\"/src/atoms/index.ts/doubleAtom\\", atom(get => get(countAtom) * 2));"
+  `)
+})
+
+it('Should add a cache for multiple exported atoms', () => {
+  expect(
+    transform(
+      `
+  export const countAtom = atom(0);
+  export const doubleAtom = atom((get) => get(countAtom) * 2);
+  `,
+      '/src/atoms/index.ts'
+    )
+  ).toMatchInlineSnapshot(`
+    "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;
+      }
+
+    };
+    export const countAtom = globalThis.jotaiAtomCache.get(\\"/src/atoms/index.ts/countAtom\\", atom(0));
+    export const doubleAtom = globalThis.jotaiAtomCache.get(\\"/src/atoms/index.ts/doubleAtom\\", atom(get => get(countAtom) * 2));"
+  `)
+})
+
+it('Should add a cache for a default exported atom', () => {
+  expect(transform(`export default atom(0);`, '/src/atoms/index.ts'))
+    .toMatchInlineSnapshot(`
+    "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;
+      }
+
+    };
+    export default globalThis.jotaiAtomCache.get(\\"/src/atoms/index.ts/defaultExport\\", atom(0));"
+  `)
+})
+
+it('Should add a cache for mixed exports of atoms', () => {
+  expect(
+    transform(
+      `
+  export const countAtom = atom(0);
+  export default atom((get) => get(countAtom) * 2);
+  `,
+      '/src/atoms/index.ts'
+    )
+  ).toMatchInlineSnapshot(`
+    "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;
+      }
+
+    };
+    export const countAtom = globalThis.jotaiAtomCache.get(\\"/src/atoms/index.ts/countAtom\\", atom(0));
+    export default globalThis.jotaiAtomCache.get(\\"/src/atoms/index.ts/defaultExport\\", atom(get => get(countAtom) * 2));"
+  `)
+})
+
+it('Should fail if no filename is available', () => {
+  expect(() => transform(`const countAtom = atom(0);`)).toThrowError(
+    'Filename must be available'
+  )
+})