refactor: drop support for dev_* methods (#180)

* refactor: drop support for dev_* methods

* chore: update jotai dev version for testing purposes

* docs: update docs

* refactor: remove tree-shaking

* chore: update jotai package

Author: arjunvegda

README.md

@@ -8,14 +8,13 @@
 
 - Debug 🐞 atom values with ease
 - ⏳ Time-travel through your atoms and find bugs faster than before
-- Out-of-the-box 🔌 support for async/suspendible atoms
+- Out-of-the-box 🔌 support for async/suspendable atoms
 - Built-in Dark mode 🌗
 - ✅ Supports custom `store`
 - ✅ Works with provider-less mode
 - ✅ Works with Next.js
 - ✅ Supports custom `nonce` for CSP
 - ✅ Hides private atoms with ability to configure
-- ✅ Tree-shakable and built for non-production environments
 - ✅ Parses all the JavaScript values with JSON Tree view
 - ✅ Diff checking with additions and deletion highlights
 
@@ -29,7 +28,7 @@
 
 ## ☝️ Prerequisites
 
-- Jotai version `>=2.9.0`
+- Jotai version `>=v2.12.3`
 - React version `>=17.0.0`
 
 ## 📦 Setup
@@ -124,7 +123,6 @@ type DevToolsProps = {
 
 ```tsx
 import { DevTools } from 'jotai-devtools';
-// Note that this may get included in your production builds. Please import it conditionally if you want to avoid that
 import 'jotai-devtools/styles.css';
 
 const App = () => {
@@ -142,7 +140,6 @@ const App = () => {
 ```tsx
 import { createStore } from 'jotai';
 import { DevTools } from 'jotai-devtools';
-// Note that this may get included in your production builds. Please import it conditionally if you want to avoid that
 import 'jotai-devtools/styles.css';
 
 const customStore = createStore();

package.json

@@ -140,7 +140,6 @@
     "jest": "^29.7.0",
     "jest-environment-jsdom": "^29.7.0",
     "jest-watch-typeahead": "^2.2.2",
-    "jotai": "^2.11.3",
     "jotai-tanstack-query": "^0.7.2",
     "lint-staged": "^15.3.0",
     "postcss": "^8.4.49",
@@ -172,10 +171,12 @@
     "@redux-devtools/extension": "^3.3.0",
     "clsx": "^2.1.1",
     "javascript-stringify": "^2.1.0",
+    "jotai": "^2.12.3",
     "jsondiffpatch": "^0.5.0",
     "react-base16-styling": "^0.9.1",
     "react-error-boundary": "^5.0.0",
     "react-json-tree": "^0.18.0",
     "react-resizable-panels": "2.0.10"
-  }
+  },
+  "packageManager": "pnpm@8.5.1+sha512.5ba7a52fb370c346bf302202933a646f15ba7fdf5342e5f80ee2680e35a9dfd5f4f0f283904a60e1c823ddde84b606000dc7505a5fde33c35d94b255137460a4"
 }

pnpm-lock.yaml

@@ -126,14 +126,10 @@ const DevToolsProvider = ({ children }: React.PropsWithChildren) => {
   );
 };
 
-export const InternalDevTools = (props: DevToolsProps): JSX.Element | null => {
-  if (__DEV__) {
-    return (
-      <DevToolsProvider>
-        <DevToolsMain {...props} />
-      </DevToolsProvider>
-    );
-  }
-
-  return <></>;
+export const DevTools = (props: DevToolsProps): JSX.Element | null => {
+  return (
+    <DevToolsProvider>
+      <DevToolsMain {...props} />
+    </DevToolsProvider>
+  );
 };

src/DevTools/DevTools.tsx

@@ -1,11 +1,4 @@
-import { InternalDevTools } from './DevTools';
 export type { DevToolsProps } from './DevTools';
-
-// This is a workaround to make DevTools tree-shakable in production builds
-// This is due to a limitation in tsup where it does not support preserving signatures
-// of exports or generating seperate chunks for exports
-export const DevTools: typeof InternalDevTools = __DEV__
-  ? InternalDevTools
-  : () => null;
+export { DevTools } from './DevTools';
 
 export * from './utils';

src/index.ts

@@ -1,12 +1,20 @@
 import { useStore } from 'jotai/react';
 import type { Atom, WritableAtom, createStore } from 'jotai/vanilla';
 import type {
-  INTERNAL_DevStoreRev4,
-  INTERNAL_PrdStore,
-} from 'jotai/vanilla/store';
+  INTERNAL_AtomState,
+  INTERNAL_Store,
+} from 'jotai/vanilla/internals';
+
+export type DevStore = {
+  get_internal_weak_map: () => {
+    get: (atom: Atom<unknown>) => INTERNAL_AtomState | undefined;
+  };
+  get_mounted_atoms: () => Set<Atom<unknown>>;
+  restore_atoms: (values: Iterable<readonly [Atom<unknown>, unknown]>) => void;
+};
 
 export type StoreWithoutDevMethods = ReturnType<typeof createStore>;
-export type StoreWithDevMethods = INTERNAL_DevStoreRev4 & INTERNAL_PrdStore;
+export type StoreWithDevMethods = INTERNAL_Store & DevStore;
 
 export type Store = StoreWithoutDevMethods | StoreWithDevMethods;
 

src/types.ts

@@ -1,8 +1,14 @@
-import { Atom } from 'jotai';
+import { Atom, WritableAtom } from 'jotai';
+import { INTERNAL_overrideCreateStore } from 'jotai/vanilla';
+import {
+  INTERNAL_buildStoreRev1 as INTERNAL_buildStore,
+  INTERNAL_initializeStoreHooks,
+} from 'jotai/vanilla/internals';
 import {
   AnyAtom,
   AnyAtomError,
   AnyAtomValue,
+  DevStore,
   Store,
   StoreWithDevMethods,
 } from '../../types';
@@ -33,17 +39,13 @@ type DevToolsStoreMethods = {
 
 type WithDevToolsStore<S extends Store> = S & DevToolsStoreMethods;
 
-const isDevStore = (store: Store | undefined): store is StoreWithDevMethods => {
-  return store ? 'dev4_get_internal_weak_map' in store : false;
-};
-
 export const isDevToolsStore = (
   store: Store | WithDevToolsStore<Store>,
 ): store is WithDevToolsStore<Store> => {
   return 'subscribeStore' in store;
 };
 
-const __composeDevTools = (
+const __composeWithDevTools = (
   store: StoreWithDevMethods,
 ): WithDevToolsStore<StoreWithDevMethods> => {
   const { sub, set, get } = store;
@@ -120,11 +122,11 @@ const __composeDevTools = (
   };
 
   (store as WithDevToolsStore<typeof store>).getMountedAtoms = () => {
-    return store.dev4_get_mounted_atoms();
+    return store.get_mounted_atoms();
   };
 
   (store as WithDevToolsStore<typeof store>).getAtomState = (atom) => {
-    const aState = store.dev4_get_internal_weak_map().get(atom);
+    const aState = store.get_internal_weak_map().get(atom);
 
     if (aState) {
       return { v: aState.v, e: aState.e, d: new Set(aState.d.keys()) };
@@ -134,26 +136,98 @@ const __composeDevTools = (
   };
 
   (store as WithDevToolsStore<typeof store>).getMountedAtomState = (atom) => {
-    const aState = store.dev4_get_internal_weak_map().get(atom);
+    const aState = store.get_internal_weak_map().get(atom);
 
-    if (aState && aState.m) {
+    if (aState && 'm' in aState) {
       return {
-        l: aState.m.l,
-        t: aState.m.t,
+        l: (aState.m as any).l,
+        t: (aState.m as any).t,
       };
     }
 
     return undefined;
   };
 
   (store as WithDevToolsStore<typeof store>).restoreAtoms = (values) => {
-    store.dev4_restore_atoms(values);
+    store.restore_atoms(values);
     storeListeners.forEach((l) => l({ type: 'restore' }));
   };
 
   return store as typeof store & DevToolsStoreMethods;
 };
 
+const createDevStore = (): StoreWithDevMethods => {
+  let inRestoreAtom = 0;
+  const storeHooks = INTERNAL_initializeStoreHooks({});
+  const atomStateMap = new WeakMap();
+  const mountedAtoms = new WeakMap();
+  const store = INTERNAL_buildStore(
+    atomStateMap,
+    mountedAtoms,
+    undefined,
+    undefined,
+    undefined,
+    undefined,
+    storeHooks,
+    undefined,
+    (atom, get, set, ...args) => {
+      if (inRestoreAtom) {
+        return set(atom, ...(args as any));
+      }
+      return atom.write(get, set, ...(args as any));
+    },
+  );
+  const debugMountedAtoms = new Set<Atom<unknown>>();
+  storeHooks.m.add(undefined, (atom) => {
+    debugMountedAtoms.add(atom);
+    const atomState = atomStateMap.get(atom);
+    // For DevStoreRev4 compatibility
+    (atomState as any).m = mountedAtoms.get(atom);
+  });
+  storeHooks.u.add(undefined, (atom) => {
+    debugMountedAtoms.delete(atom);
+    const atomState = atomStateMap.get(atom);
+    // For DevStoreRev4 compatibility
+    delete (atomState as any).m;
+  });
+  const devStore: DevStore = {
+    // store dev methods (these are tentative and subject to change without notice)
+    get_internal_weak_map: () => atomStateMap,
+    get_mounted_atoms: () => debugMountedAtoms,
+    restore_atoms: (values) => {
+      const restoreAtom: WritableAtom<null, [], void> = {
+        read: () => null,
+        write: (_get, set) => {
+          ++inRestoreAtom;
+          try {
+            for (const [atom, value] of values) {
+              if ('init' in atom) {
+                set(atom as never, value);
+              }
+            }
+          } finally {
+            --inRestoreAtom;
+          }
+        },
+      };
+      store.set(restoreAtom);
+    },
+  };
+
+  return {
+    ...store,
+    ...devStore,
+  };
+};
+
+const isDevStore = (store: Store): store is StoreWithDevMethods => {
+  return 'get_internal_weak_map' in store;
+};
+
+INTERNAL_overrideCreateStore((prev) => {
+  return createDevStore;
+});
+
 export const composeWithDevTools = (
   store: Store,
 ): typeof store | WithDevToolsStore<typeof store> => {
@@ -163,7 +237,7 @@ export const composeWithDevTools = (
   }
 
   if (isDevStore(store)) {
-    return __composeDevTools(store);
+    return __composeWithDevTools(store);
   }
 
   return store;

src/utils/internals/compose-with-devtools.ts

@@ -19,8 +19,6 @@ const baseConfig: Options = {
   // Outputs `dist/index.js` and `dist/utils.js`
   entry: {
     index: 'src/index.ts',
-    // Workaround to generate seperate chunks for DevTools so we could export a null component for production builds
-    internal__devtools: 'src/DevTools/index.ts',
     utils: 'src/utils/index.ts',
   },
   loader: {