[WIP] Refactor to use the new jotai store derive method

__tests__/ScopeProvider/01_basic_spec.tsx

@@ -554,7 +554,7 @@ describe('Counter', () => {
   });
 
   /*
-    base, derivedA(base), derivedB(base),
+    base, derivedA(base), derivedB(base)
     S0[base]: base0
     S1[base]: base1
     S2[base]: base2
@@ -629,7 +629,7 @@ describe('Counter', () => {
   });
 
   /*
-    baseA, baseB, baseC, derived(baseA + baseB + baseC),
+    baseA, baseB, baseC, derived(baseA + baseB + baseC)
     S0[     ]: derived(baseA0 + baseB0 + baseC0)
     S1[baseB]: derived(baseA0 + baseB1 + baseC0)
     S2[baseC]: derived(baseA0 + baseB1 + baseC2)

__tests__/ScopeProvider/03_nested.tsx

@@ -79,6 +79,11 @@ function App() {
 }
 
 describe('Counter', () => {
+  /*
+  baseA, baseB, baseC
+  S1[baseA]: baseA1 baseB0 baseC0
+  S2[baseB]: baseA1 baseB2 baseC0
+  */
   test('nested primitive atoms are correctly scoped', () => {
     const { container } = render(<App />);
     const increaseUnscopedBase1 = '.unscoped.setBase1';

__tests__/ScopeProvider/05_derived_self.tsx

@@ -46,6 +46,10 @@ function App() {
 }
 
 describe('Self', () => {
+  /*
+  baseA, derivedB(baseA, derivedB)
+  S1[baseA]: baseA1, derivedB0(baseA1, derivedB0)
+  */
   test('derived dep scope is preserved in self reference', () => {
     const { container } = render(<App />);
     expect(

package.json

@@ -75,7 +75,7 @@
     "html-webpack-plugin": "^5.5.3",
     "jest": "^29.7.0",
     "jest-environment-jsdom": "^29.7.0",
-    "jotai": "2.9.0",
+    "jotai": "https://pkg.csb.dev/pmndrs/jotai/commit/b30da262/jotai",
     "microbundle": "^0.15.1",
     "npm-run-all": "^4.1.5",
     "prettier": "^3.0.3",
@@ -90,7 +90,7 @@
     "webpack-dev-server": "^4.15.1"
   },
   "peerDependencies": {
-    "jotai": ">=2.9.0",
+    "jotai": ">=2.9.1",
     "react": ">=17.0.0"
   }
 }

src/ScopeProvider.tsx

@@ -0,0 +1,78 @@
+import { type ReactNode, useState } from 'react';
+import { Provider, useStore } from 'jotai/react';
+import type { Atom, getDefaultStore } from 'jotai/vanilla';
+
+type Store = ReturnType<typeof getDefaultStore>;
+type NamedStore = Store & { name?: string };
+
+type ScopeProviderProps = {
+  atoms: Iterable<Atom<unknown>>;
+  debugName?: string;
+  store?: Store;
+  children: ReactNode;
+};
+export function ScopeProvider(props: ScopeProviderProps) {
+  const { atoms, children, debugName, ...options } = props;
+  const baseStore = useStore(options);
+  const scopedAtoms = new Set(atoms);
+
+  function initialize() {
+    return {
+      scopedStore: createScopedStore(baseStore, scopedAtoms, debugName),
+      hasChanged(current: {
+        baseStore: Store;
+        scopedAtoms: Set<Atom<unknown>>;
+      }) {
+        return (
+          !isEqualSet(scopedAtoms, current.scopedAtoms) ||
+          current.baseStore !== baseStore
+        );
+      },
+    };
+  }
+
+  const [{ hasChanged, scopedStore }, setState] = useState(initialize);
+  if (hasChanged({ scopedAtoms, baseStore })) {
+    setState(initialize);
+  }
+  return <Provider store={scopedStore}>{children}</Provider>;
+}
+
+function isEqualSet(a: Set<unknown>, b: Set<unknown>) {
+  return a === b || (a.size === b.size && Array.from(a).every((v) => b.has(v)));
+}
+
+/**
+ * @returns a derived store that intercepts get and set calls to apply the scope
+ */
+export function createScopedStore(
+  baseStore: Store,
+  scopedAtoms: Set<Atom<unknown>>,
+  debugName?: string,
+) {
+  const derivedStore: NamedStore = baseStore.unstable_derive((getAtomState) => {
+    const scopedAtomStateMap = new WeakMap();
+    const scopedAtomStateSet = new WeakSet();
+    return [
+      (atom, originAtomState) => {
+        if (
+          scopedAtomStateSet.has(originAtomState as never) ||
+          scopedAtoms.has(atom)
+        ) {
+          let atomState = scopedAtomStateMap.get(atom);
+          if (!atomState) {
+            atomState = { d: new Map(), p: new Set(), n: 0 };
+            scopedAtomStateMap.set(atom, atomState);
+            scopedAtomStateSet.add(atomState);
+          }
+          return atomState;
+        }
+        return getAtomState(atom, originAtomState);
+      },
+    ];
+  });
+  if (debugName) {
+    derivedStore.name = debugName;
+  }
+  return derivedStore;
+}