From 91c176f142757d19d33e5aa546083626e87e9c19 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Fri, 19 Jul 2024 19:26:04 -0700 Subject: [PATCH] refactor jotai-scope to use the new store.unstable_derive api --- __tests__/ScopeProvider/01_basic_spec.tsx | 4 +- __tests__/ScopeProvider/03_nested.tsx | 5 + __tests__/ScopeProvider/05_derived_self.tsx | 4 + package.json | 4 +- src/ScopeProvider.tsx | 78 +++++++ src/ScopeProvider/ScopeProvider.tsx | 68 ------ src/ScopeProvider/patchedStore.ts | 39 ---- src/ScopeProvider/scope.ts | 246 -------------------- src/ScopeProvider/types.ts | 5 - src/index.ts | 2 +- yarn.lock | 4 +- 11 files changed, 94 insertions(+), 365 deletions(-) create mode 100644 src/ScopeProvider.tsx delete mode 100644 src/ScopeProvider/ScopeProvider.tsx delete mode 100644 src/ScopeProvider/patchedStore.ts delete mode 100644 src/ScopeProvider/scope.ts delete mode 100644 src/ScopeProvider/types.ts diff --git a/__tests__/ScopeProvider/01_basic_spec.tsx b/__tests__/ScopeProvider/01_basic_spec.tsx index ec0ed43..2fbe2ce 100644 --- a/__tests__/ScopeProvider/01_basic_spec.tsx +++ b/__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) diff --git a/__tests__/ScopeProvider/03_nested.tsx b/__tests__/ScopeProvider/03_nested.tsx index 7632089..76d7299 100644 --- a/__tests__/ScopeProvider/03_nested.tsx +++ b/__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(); const increaseUnscopedBase1 = '.unscoped.setBase1'; diff --git a/__tests__/ScopeProvider/05_derived_self.tsx b/__tests__/ScopeProvider/05_derived_self.tsx index e63953b..1983c83 100644 --- a/__tests__/ScopeProvider/05_derived_self.tsx +++ b/__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(); expect( diff --git a/package.json b/package.json index c61bea7..a00f438 100644 --- a/package.json +++ b/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" } } diff --git a/src/ScopeProvider.tsx b/src/ScopeProvider.tsx new file mode 100644 index 0000000..111fa8d --- /dev/null +++ b/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; +type NamedStore = Store & { name?: string }; + +type ScopeProviderProps = { + atoms: Iterable>; + 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>; + }) { + return ( + !isEqualSet(scopedAtoms, current.scopedAtoms) || + current.baseStore !== baseStore + ); + }, + }; + } + + const [{ hasChanged, scopedStore }, setState] = useState(initialize); + if (hasChanged({ scopedAtoms, baseStore })) { + setState(initialize); + } + return {children}; +} + +function isEqualSet(a: Set, b: Set) { + 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>, + 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; +} diff --git a/src/ScopeProvider/ScopeProvider.tsx b/src/ScopeProvider/ScopeProvider.tsx deleted file mode 100644 index 44479f3..0000000 --- a/src/ScopeProvider/ScopeProvider.tsx +++ /dev/null @@ -1,68 +0,0 @@ -import { Provider, useStore } from 'jotai/react'; -import { - createContext, - useContext, - useState, - type PropsWithChildren, -} from 'react'; -import { createScope, type Scope } from './scope'; -import type { AnyAtom, Store } from './types'; -import { createPatchedStore, isTopLevelScope } from './patchedStore'; - -const ScopeContext = createContext<{ - scope: Scope | undefined; - baseStore: Store | undefined; -}>({ scope: undefined, baseStore: undefined }); - -export function ScopeProvider({ - atoms, - children, - debugName, -}: PropsWithChildren<{ atoms: Iterable; debugName?: string }>) { - const parentStore: Store = useStore(); - let { scope: parentScope, baseStore = parentStore } = - useContext(ScopeContext); - // if this ScopeProvider is the first descendant scope under Provider then it is the top level scope - // https://github.com/jotaijs/jotai-scope/pull/33#discussion_r1604268003 - if (isTopLevelScope(parentStore)) { - parentScope = undefined; - baseStore = parentStore; - } - - // atomSet is used to detect if the atoms prop has changed. - const atomSet = new Set(atoms); - - function initialize() { - const scope = createScope(atoms, parentScope, debugName); - return { - patchedStore: createPatchedStore(baseStore, scope), - scopeContext: { scope, baseStore }, - hasChanged(current: { - baseStore: Store; - parentScope: Scope | undefined; - atomSet: Set; - }) { - return ( - parentScope !== current.parentScope || - !isEqualSet(atomSet, current.atomSet) || - current.baseStore !== baseStore - ); - }, - }; - } - - const [state, setState] = useState(initialize); - const { hasChanged, scopeContext, patchedStore } = state; - if (hasChanged({ parentScope, atomSet, baseStore })) { - setState(initialize); - } - return ( - - {children} - - ); -} - -function isEqualSet(a: Set, b: Set) { - return a === b || (a.size === b.size && Array.from(a).every((v) => b.has(v))); -} diff --git a/src/ScopeProvider/patchedStore.ts b/src/ScopeProvider/patchedStore.ts deleted file mode 100644 index fd00f0c..0000000 --- a/src/ScopeProvider/patchedStore.ts +++ /dev/null @@ -1,39 +0,0 @@ -import type { Scope } from './scope'; -import type { Store } from './types'; - -function PatchedStore() {} - -/** - * @returns a patched store that intercepts get and set calls to apply the scope - */ -export function createPatchedStore(baseStore: Store, scope: Scope): Store { - const store: Store = { - ...baseStore, - get(anAtom, ...args) { - const [scopedAtom] = scope.getAtom(anAtom); - return baseStore.get(scopedAtom, ...args); - }, - set(anAtom, ...args) { - const [scopedAtom, implicitScope] = scope.getAtom(anAtom); - const restore = scope.prepareWriteAtom(scopedAtom, anAtom, implicitScope); - try { - return baseStore.set(scopedAtom, ...args); - } finally { - restore?.(); - } - }, - sub(anAtom, ...args) { - const [scopedAtom] = scope.getAtom(anAtom); - return baseStore.sub(scopedAtom, ...args); - }, - // TODO: update this patch to support devtools - }; - return Object.assign(Object.create(PatchedStore.prototype), store); -} - -/** - * @returns true if the current scope is the first descendant scope under Provider - */ -export function isTopLevelScope(parentStore: Store) { - return !(parentStore instanceof PatchedStore); -} diff --git a/src/ScopeProvider/scope.ts b/src/ScopeProvider/scope.ts deleted file mode 100644 index fff10a1..0000000 --- a/src/ScopeProvider/scope.ts +++ /dev/null @@ -1,246 +0,0 @@ -import { atom, type Atom } from 'jotai'; -import { type AnyAtom, type AnyWritableAtom } from './types'; - -export type Scope = { - /** - * Returns a scoped atom from the original atom. - * @param anAtom - * @param implicitScope the atom is implicitly scoped in the provided scope - * @returns the scoped atom and the scope of the atom - */ - getAtom: (anAtom: T, implicitScope?: Scope) => [T, Scope?]; - /** - * @modifies the atom's write function for atoms that can hold a value - * @returns a function to restore the original write function - */ - prepareWriteAtom: ( - anAtom: T, - originalAtom: T, - implicitScope?: Scope, - ) => (() => void) | undefined; - - /** - * @debug - */ - name?: string; - - /** - * @debug - */ - toString?: () => string; -}; - -const globalScopeKey: { name?: string } = {}; -if (process.env.NODE_ENV !== 'production') { - globalScopeKey.name = 'unscoped'; - globalScopeKey.toString = toString; -} - -type GlobalScopeKey = typeof globalScopeKey; - -export function createScope( - atoms: Iterable, - parentScope: Scope | undefined, - scopeName?: string | undefined, -): Scope { - const explicit = new WeakMap(); - const implicit = new WeakMap(); - type ScopeMap = WeakMap; - const inherited = new WeakMap(); - - const currentScope: Scope = { - getAtom, - prepareWriteAtom(anAtom, originalAtom, implicitScope) { - if ( - originalAtom.read === defaultRead && - isWritableAtom(originalAtom) && - isWritableAtom(anAtom) && - originalAtom.write !== defaultWrite && - currentScope !== implicitScope - ) { - // atom is writable with init and holds a value - // we need to preserve the value, so we don't want to copy the atom - // instead, we need to override write until the write is finished - const { write } = originalAtom; - anAtom.write = createScopedWrite( - originalAtom.write.bind( - originalAtom, - ) as (typeof originalAtom)['write'], - implicitScope, - ); - return () => { - anAtom.write = write; - }; - } - return undefined; - }, - }; - - if (scopeName && process.env.NODE_ENV !== 'production') { - currentScope.name = scopeName; - currentScope.toString = toString; - } - // populate explicitly scoped atoms - for (const anAtom of atoms) { - explicit.set(anAtom, [cloneAtom(anAtom, currentScope), currentScope]); - } - - /** - * Returns a scoped atom from the original atom. - * @param anAtom - * @param implicitScope the atom is implicitly scoped in the provided scope - * @returns the scoped atom and the scope of the atom - */ - function getAtom( - anAtom: T, - implicitScope?: Scope, - ): [T, Scope?] { - if (explicit.has(anAtom)) { - return explicit.get(anAtom) as [T, Scope]; - } - if (implicitScope === currentScope) { - // dependencies of explicitly scoped atoms are implicitly scoped - // implicitly scoped atoms are only accessed by implicit and explicit scoped atoms - if (!implicit.has(anAtom)) { - implicit.set(anAtom, [cloneAtom(anAtom, implicitScope), implicitScope]); - } - return implicit.get(anAtom) as [T, Scope]; - } - const scopeKey = implicitScope ?? globalScopeKey; - if (parentScope) { - // inherited atoms are copied so they can access scoped atoms - // but they are not explicitly scoped - // dependencies of inherited atoms first check if they are explicitly scoped - // otherwise they use their original scope's atom - if (!inherited.get(scopeKey)?.has(anAtom)) { - const [ancestorAtom, explicitScope] = parentScope.getAtom( - anAtom, - implicitScope, - ); - setInheritedAtom( - inheritAtom(ancestorAtom, anAtom, explicitScope), - anAtom, - implicitScope, - explicitScope, - ); - } - return inherited.get(scopeKey)!.get(anAtom) as [T, Scope]; - } - if (!inherited.get(scopeKey)?.has(anAtom)) { - // non-primitive atoms may need to access scoped atoms - // so we need to create a copy of the atom - setInheritedAtom(inheritAtom(anAtom, anAtom), anAtom); - } - return inherited.get(scopeKey)!.get(anAtom) as [T, Scope?]; - } - - function setInheritedAtom( - scopedAtom: T, - originalAtom: T, - implicitScope?: Scope, - explicitScope?: Scope, - ) { - const scopeKey = implicitScope ?? globalScopeKey; - if (!inherited.has(scopeKey)) { - inherited.set(scopeKey, new WeakMap()); - } - inherited.get(scopeKey)!.set( - originalAtom, - [ - scopedAtom, // - explicitScope, - ].filter(Boolean) as [T, Scope?], - ); - } - - /** - * @returns a copy of the atom for derived atoms or the original atom for primitive and writable atoms - */ - function inheritAtom( - anAtom: Atom, - originalAtom: Atom, - implicitScope?: Scope, - ) { - if (originalAtom.read !== defaultRead) { - return cloneAtom(originalAtom, implicitScope); - } - return anAtom; - } - - /** - * @returns a scoped copy of the atom - */ - function cloneAtom(originalAtom: Atom, implicitScope?: Scope) { - // avoid reading `init` to preserve lazy initialization - const scopedAtom: Atom = Object.create( - Object.getPrototypeOf(originalAtom), - Object.getOwnPropertyDescriptors(originalAtom), - ); - - if (scopedAtom.read !== defaultRead) { - scopedAtom.read = createScopedRead( - originalAtom.read.bind(originalAtom), - implicitScope, - ); - } - - if ( - isWritableAtom(scopedAtom) && - isWritableAtom(originalAtom) && - scopedAtom.write !== defaultWrite - ) { - scopedAtom.write = createScopedWrite( - originalAtom.write.bind(originalAtom), - implicitScope, - ); - } - - return scopedAtom; - } - - function createScopedRead>( - read: T['read'], - implicitScope?: Scope, - ): T['read'] { - return function scopedRead(get, opts) { - return read( - function scopedGet(a) { - const [scopedAtom] = getAtom(a, implicitScope); - return get(scopedAtom); - }, // - opts, - ); - }; - } - - function createScopedWrite( - write: T['write'], - implicitScope?: Scope, - ): T['write'] { - return function scopedWrite(get, set, ...args) { - return write( - function scopedGet(a) { - const [scopedAtom] = getAtom(a, implicitScope); - return get(scopedAtom); - }, - function scopedSet(a, ...v) { - const [scopedAtom] = getAtom(a, implicitScope); - return set(scopedAtom, ...v); - }, - ...args, - ); - }; - } - - return currentScope; -} - -function isWritableAtom(anAtom: AnyAtom): anAtom is AnyWritableAtom { - return 'write' in anAtom; -} - -const { read: defaultRead, write: defaultWrite } = atom(null); - -function toString(this: { name: string }) { - return this.name; -} diff --git a/src/ScopeProvider/types.ts b/src/ScopeProvider/types.ts deleted file mode 100644 index 27952a0..0000000 --- a/src/ScopeProvider/types.ts +++ /dev/null @@ -1,5 +0,0 @@ -import type { Atom, WritableAtom, getDefaultStore } from 'jotai'; - -export type AnyAtom = Atom | WritableAtom; -export type AnyWritableAtom = WritableAtom; -export type Store = ReturnType; diff --git a/src/index.ts b/src/index.ts index 0c90ebb..0da2644 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,2 +1,2 @@ export { createIsolation } from './createIsolation'; -export { ScopeProvider } from './ScopeProvider/ScopeProvider'; +export { ScopeProvider } from './ScopeProvider'; diff --git a/yarn.lock b/yarn.lock index 64fe2fc..99071b6 100644 --- a/yarn.lock +++ b/yarn.lock @@ -5340,9 +5340,9 @@ jest@^29.7.0: import-local "^3.0.2" jest-cli "^29.7.0" -"jotai@github:pmndrs/jotai#unstable_derive": +"jotai@https://pkg.csb.dev/pmndrs/jotai/commit/b30da262/jotai": version "2.9.0" - resolved "https://codeload.github.com/pmndrs/jotai/tar.gz/b30da262aa0668b707f86d34f55f05d8c968e364" + resolved "https://pkg.csb.dev/pmndrs/jotai/commit/b30da262/jotai#81493df4d9bd51048cdd6d77b5051265f2c04463" "js-tokens@^3.0.0 || ^4.0.0", js-tokens@^4.0.0: version "4.0.0"