From 30690e50255075a01eaf96edb445a0b4b71f7eca Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 9 Jul 2024 15:56:32 +0900 Subject: [PATCH 01/31] experimental(core): expose unstable_derive instead of unstable_is --- src/vanilla/atom.ts | 1 - src/vanilla/store.ts | 45 +++++++++++++++++++++++++++++--------------- 2 files changed, 30 insertions(+), 16 deletions(-) diff --git a/src/vanilla/atom.ts b/src/vanilla/atom.ts index b2cfcba268..6da084b612 100644 --- a/src/vanilla/atom.ts +++ b/src/vanilla/atom.ts @@ -40,7 +40,6 @@ type OnMount = < export interface Atom { toString: () => string read: Read - unstable_is?(a: Atom): boolean debugLabel?: string /** * To ONLY be used by Jotai libraries to mark atoms as private. Subject to change. diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index f8d317d07d..a41e6965bf 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -8,9 +8,6 @@ type OnUnmount = () => void type Getter = Parameters[0] type Setter = Parameters[1] -const isSelfAtom = (atom: AnyAtom, a: AnyAtom): boolean => - atom.unstable_is ? atom.unstable_is(a) : a === atom - const hasInitialValue = >( atom: T, ): atom is T & (T extends Atom ? { init: Value } : never) => @@ -251,6 +248,12 @@ type DevStoreRev4 = { dev4_restore_atoms: (values: Iterable) => void } +// internal & unstable type +type StoreArgs = [ + atomStateMap: WeakMap, + isSelfAtom: (atom: AnyAtom, a: AnyAtom) => boolean, +] + type PrdStore = { get: (atom: Atom) => Value set: ( @@ -258,14 +261,17 @@ type PrdStore = { ...args: Args ) => Result sub: (atom: AnyAtom, listener: () => void) => () => void + unstable_derive: (fn: (...args: StoreArgs) => StoreArgs) => Store } type Store = PrdStore | (PrdStore & DevStoreRev4) export type INTERNAL_DevStoreRev4 = DevStoreRev4 export type INTERNAL_PrdStore = PrdStore -export const createStore = (): Store => { - const atomStateMap = new WeakMap() +const buildStore = ( + atomStateMap: WeakMap, + isSelfAtom: (atom: AnyAtom, a: AnyAtom) => boolean, +): Store => { // for debugging purpose only let debugMountedAtoms: Set @@ -685,11 +691,21 @@ export const createStore = (): Store => { } } + const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) => { + const derivedArgs = fn(atomStateMap, isSelfAtom) + const derivedStore = buildStore(...derivedArgs) + return derivedStore + } + + const store: Store = { + get: readAtom, + set: writeAtom, + sub: subscribeAtom, + unstable_derive, + } + if (import.meta.env?.MODE !== 'production') { - const store: Store = { - get: readAtom, - set: writeAtom, - sub: subscribeAtom, + const devStore: DevStoreRev4 = { // store dev methods (these are tentative and subject to change without notice) dev4_get_internal_weak_map: () => atomStateMap, dev4_get_mounted_atoms: () => debugMountedAtoms, @@ -711,15 +727,14 @@ export const createStore = (): Store => { flushPending(pending) }, } - return store - } - return { - get: readAtom, - set: writeAtom, - sub: subscribeAtom, + Object.assign(store, devStore) } + return store } +export const createStore = (): Store => + buildStore(new WeakMap(), (atom, a) => atom === a) + let defaultStore: Store | undefined export const getDefaultStore = (): Store => { From 634136712f3b8eadb7470249e02293ccd348cd6c Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 9 Jul 2024 15:58:55 +0900 Subject: [PATCH 02/31] chore: remove empty line --- src/vanilla/store.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index a41e6965bf..0071279caa 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -703,7 +703,6 @@ const buildStore = ( sub: subscribeAtom, unstable_derive, } - if (import.meta.env?.MODE !== 'production') { const devStore: DevStoreRev4 = { // store dev methods (these are tentative and subject to change without notice) From a885f65a624118f8ce24b7229a24e24dbd2aaa30 Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 9 Jul 2024 16:07:30 +0900 Subject: [PATCH 03/31] chore types --- src/vanilla/store.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 0071279caa..67b2e8df57 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -249,7 +249,7 @@ type DevStoreRev4 = { } // internal & unstable type -type StoreArgs = [ +type StoreArgs = readonly [ atomStateMap: WeakMap, isSelfAtom: (atom: AnyAtom, a: AnyAtom) => boolean, ] @@ -269,8 +269,8 @@ export type INTERNAL_DevStoreRev4 = DevStoreRev4 export type INTERNAL_PrdStore = PrdStore const buildStore = ( - atomStateMap: WeakMap, - isSelfAtom: (atom: AnyAtom, a: AnyAtom) => boolean, + atomStateMap: StoreArgs[0], + isSelfAtom: StoreArgs[1], ): Store => { // for debugging purpose only let debugMountedAtoms: Set From 1c352298fdc1adfab7e0e784530cc794488aa310 Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 10 Jul 2024 09:27:56 +0900 Subject: [PATCH 04/31] follow resolveAtom in #2609 --- src/vanilla/store.ts | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 67b2e8df57..f051790a3c 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -251,7 +251,7 @@ type DevStoreRev4 = { // internal & unstable type type StoreArgs = readonly [ atomStateMap: WeakMap, - isSelfAtom: (atom: AnyAtom, a: AnyAtom) => boolean, + resolveAtom: (atom: T) => T, ] type PrdStore = { @@ -270,7 +270,7 @@ export type INTERNAL_PrdStore = PrdStore const buildStore = ( atomStateMap: StoreArgs[0], - isSelfAtom: StoreArgs[1], + resolveAtom: StoreArgs[1], ): Store => { // for debugging purpose only let debugMountedAtoms: Set @@ -388,7 +388,8 @@ const buildStore = ( atomState.d.clear() let isSync = true const getter: Getter = (a: Atom) => { - if (isSelfAtom(atom, a)) { + a = resolveAtom(a) + if (a === (atom as AnyAtom)) { const aState = getAtomState(a) if (!isAtomStateInitialized(aState)) { if (hasInitialValue(a)) { @@ -468,7 +469,7 @@ const buildStore = ( } const readAtom = (atom: Atom): Value => - returnAtomValue(readAtomState(undefined, atom)) + returnAtomValue(readAtomState(undefined, resolveAtom(atom))) const recomputeDependents = (pending: Pending, atom: AnyAtom) => { const getDependents = (a: AnyAtom): Set => { @@ -497,7 +498,6 @@ const buildStore = ( } markedAtoms.add(n) for (const m of getDependents(n)) { - // we shouldn't use isSelfAtom here. if (n !== m) { visit(m) } @@ -543,13 +543,14 @@ const buildStore = ( ...args: Args ): Result => { const getter: Getter = (a: Atom) => - returnAtomValue(readAtomState(pending, a)) + returnAtomValue(readAtomState(pending, resolveAtom(a))) const setter: Setter = ( a: WritableAtom, ...args: As ) => { let r: R | undefined - if (isSelfAtom(atom, a)) { + a = resolveAtom(a) + if (a === (atom as AnyAtom)) { if (!hasInitialValue(a)) { // NOTE technically possible but restricted as it may cause bugs throw new Error('atom not writable') @@ -565,7 +566,7 @@ const buildStore = ( recomputeDependents(pending, a) } } else { - r = writeAtomState(pending, a as AnyWritableAtom, ...args) as R + r = writeAtomState(pending, a, ...args) as R } flushPending(pending) return r as R @@ -579,7 +580,7 @@ const buildStore = ( ...args: Args ): Result => { const pending = createPending() - const result = writeAtomState(pending, atom, ...args) + const result = writeAtomState(pending, resolveAtom(atom), ...args) flushPending(pending) return result } @@ -678,6 +679,7 @@ const buildStore = ( } const subscribeAtom = (atom: AnyAtom, listener: () => void) => { + atom = resolveAtom(atom) const pending = createPending() const mounted = mountAtom(pending, atom) flushPending(pending) @@ -692,7 +694,7 @@ const buildStore = ( } const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) => { - const derivedArgs = fn(atomStateMap, isSelfAtom) + const derivedArgs = fn(atomStateMap, resolveAtom) const derivedStore = buildStore(...derivedArgs) return derivedStore } @@ -732,7 +734,7 @@ const buildStore = ( } export const createStore = (): Store => - buildStore(new WeakMap(), (atom, a) => atom === a) + buildStore(new WeakMap(), (atom) => atom) let defaultStore: Store | undefined From e92f1c10c8df24989e10e30d25cbb1e721e6ccd8 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Wed, 10 Jul 2024 10:42:26 +0900 Subject: [PATCH 05/31] store test from #2609 --- tests/vanilla/store.test.tsx | 156 +++++++++++++++++++++++++++++++++++ 1 file changed, 156 insertions(+) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 652a551ca4..4147c0bdd9 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -555,3 +555,159 @@ describe('aborting atoms', () => { expect(callAfterAbort).toHaveBeenCalledTimes(1) }) }) + +describe('unstable_resolve resolves the correct value for', () => { + function nextTask() { + return new Promise((resolve) => setTimeout(resolve)) + } + + it('primitive atom', async () => { + const store = createStore() + store.unstable_resolve = (atom) => { + if (atom === (pseudo as Atom)) { + return a as unknown as typeof atom + } + return atom + } + + const pseudo = atom('pseudo') as typeof a + pseudo.debugLabel = 'pseudo' + pseudo.onMount = (setSelf) => setSelf((v) => v + ':pseudo-mounted') + const a = atom('a') + a.debugLabel = 'a' + a.onMount = (setSelf) => setSelf((v) => v + ':a-mounted') + + expect(store.get(pseudo)).toBe('a') + const callback = vi.fn() + store.sub(pseudo, callback) + expect(store.get(pseudo)).toBe('a:a-mounted') + store.set(pseudo, (v) => v + ':a-updated') + expect(store.get(pseudo)).toBe('a:a-mounted:a-updated') + await nextTask() + expect(store.get(pseudo)).toBe('a:a-mounted:a-updated') + }) + + it('derived atom', async () => { + const store = createStore() + store.unstable_resolve = (atom) => { + if (atom === (pseudo as Atom)) { + return a as unknown as typeof atom + } + return atom + } + + const pseudo = atom('pseudo') as typeof a + pseudo.debugLabel = 'pseudo' + pseudo.onMount = (setSelf) => setSelf((v) => v + ':pseudo-mounted') + const a = atom('a') + a.debugLabel = 'a' + a.onMount = (setSelf) => setSelf((v) => v + ':a-mounted') + const b = atom('b') + b.debugLabel = 'b' + const c = atom((get) => get(pseudo) + get(b)) + c.debugLabel = 'c' + expect(store.get(c)).toBe('ab') + const d = atom( + (_get, { setSelf }) => setTimeout(setSelf, 0, 'd'), + (get, set, v) => set(pseudo, get(pseudo) + v), + ) + store.get(d) + await nextTask() + expect(store.get(a)).toBe('ad') + expect(store.get(c)).toBe('adb') + expect(store.get(pseudo)).toEqual('ad') + const callback = vi.fn() + store.sub(c, callback) + expect(store.get(pseudo)).toBe('ad:a-mounted') + delete store.unstable_resolve + await nextTask() + expect(store.get(pseudo)).toEqual('pseudo') + store.sub(pseudo, callback) + expect(store.get(pseudo)).toEqual('pseudo:pseudo-mounted') + }) + + it('writable atom', async () => { + const store = createStore() + store.unstable_resolve = (atom) => { + if (atom === (pseudo as Atom)) { + return a as unknown as typeof atom + } + return atom + } + + const pseudoWriteFn = vi.fn() + const pseudo = atom('pseudo', pseudoWriteFn) as unknown as typeof a + pseudo.debugLabel = 'pseudo' + pseudo.onMount = (setSelf) => setSelf('pseudo-mounted') + const a = atom('a', (get, set, value: string) => { + set(pseudo, get(pseudo) + ':' + value) + return () => set(pseudo, get(pseudo) + ':a-unmounted') + }) + a.debugLabel = 'a' + a.onMount = (setSelf) => setSelf('a-mounted') + expect(store.get(pseudo)).toBe('a') + const callback = vi.fn() + const unsub = store.sub(pseudo, callback) + await nextTask() + expect(store.get(pseudo)).toBe('a:a-mounted') + const value = store.set(pseudo, 'a-updated') + expect(pseudoWriteFn).not.toHaveBeenCalled() + expect(typeof value).toBe('function') + expect(store.get(pseudo)).toBe('a:a-mounted:a-updated') + unsub() + await nextTask() + expect(store.get(pseudo)).toBe('a:a-mounted:a-updated:a-unmounted') + }) + + it('this in read and write', async () => { + const store = createStore() + store.unstable_resolve = (atom) => { + if (atom === (pseudo as Atom)) { + return this_read as unknown as typeof atom + } + return atom + } + + const pseudo = atom('pseudo') as typeof this_read + pseudo.debugLabel = 'pseudo' + let i = 0 + const this_read = atom(function read(this: any, get) { + return i++ % 2 ? get(this) : 'this_read' + }) + this_read.debugLabel = 'this_read' + expect(store.get(pseudo)).toBe('this_read') + + store.unstable_resolve = (atom) => { + if (atom === (pseudo as Atom)) { + return this_write as unknown as typeof atom + } + return atom + } + + const this_write = atom( + 'this', + function write(this: any, get, set, value: string) { + set(this, get(this) + ':' + value) + return () => set(this, get(this) + ':this_write-unmounted') + }, + ) + this_write.debugLabel = 'this_write' + this_write.onMount = (setSelf) => setSelf('this_write-mounted') + expect(store.get(pseudo)).toBe('this') + const callback = vi.fn() + const unsub = store.sub(pseudo, callback) + await nextTask() + expect(store.get(pseudo)).toBe('this:this_write-mounted') + expect(callback).not.toHaveBeenCalledOnce() + store.set(this_write, 'this_write-updated') + expect(store.get(pseudo)).toBe('this:this_write-mounted:this_write-updated') + await nextTask() + unsub() + await nextTask() + expect(store.get(pseudo)).toBe( + 'this:this_write-mounted:this_write-updated:this_write-unmounted', + ) + delete store.unstable_resolve + expect(store.get(pseudo)).toBe('pseudo') + }) +}) From d72e26baa352ab7b84b98a98404194d3de81a8ac Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 10 Jul 2024 11:12:27 +0900 Subject: [PATCH 06/31] adjust store.test.tsx for unstable_derive --- tests/vanilla/store.test.tsx | 52 +++++++++++++++++++++++------------- 1 file changed, 34 insertions(+), 18 deletions(-) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 4147c0bdd9..ba4356f28f 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -1,7 +1,7 @@ import { waitFor } from '@testing-library/dom' import { assert, describe, expect, it, vi } from 'vitest' import { atom, createStore } from 'jotai/vanilla' -import type { Getter } from 'jotai/vanilla' +import type { Atom, Getter } from 'jotai/vanilla' it('should not fire on subscribe', async () => { const store = createStore() @@ -556,19 +556,22 @@ describe('aborting atoms', () => { }) }) -describe('unstable_resolve resolves the correct value for', () => { +describe('unstable_derive resolves the correct value for', () => { function nextTask() { return new Promise((resolve) => setTimeout(resolve)) } it('primitive atom', async () => { - const store = createStore() - store.unstable_resolve = (atom) => { - if (atom === (pseudo as Atom)) { + const unstable_resolve = >(atom: T): T => { + if (atom === (pseudo as Atom)) { return a as unknown as typeof atom } return atom } + const store = createStore().unstable_derive((atomStateMap) => [ + atomStateMap, + unstable_resolve, + ]) const pseudo = atom('pseudo') as typeof a pseudo.debugLabel = 'pseudo' @@ -588,13 +591,18 @@ describe('unstable_resolve resolves the correct value for', () => { }) it('derived atom', async () => { - const store = createStore() - store.unstable_resolve = (atom) => { - if (atom === (pseudo as Atom)) { + let unstable_resolve: + | (>(atom: T) => T) + | undefined = (atom) => { + if (atom === (pseudo as Atom)) { return a as unknown as typeof atom } return atom } + const store = createStore().unstable_derive((atomStateMap) => [ + atomStateMap, + (atom) => unstable_resolve?.(atom) ?? atom, + ]) const pseudo = atom('pseudo') as typeof a pseudo.debugLabel = 'pseudo' @@ -619,7 +627,7 @@ describe('unstable_resolve resolves the correct value for', () => { const callback = vi.fn() store.sub(c, callback) expect(store.get(pseudo)).toBe('ad:a-mounted') - delete store.unstable_resolve + unstable_resolve = undefined await nextTask() expect(store.get(pseudo)).toEqual('pseudo') store.sub(pseudo, callback) @@ -627,13 +635,16 @@ describe('unstable_resolve resolves the correct value for', () => { }) it('writable atom', async () => { - const store = createStore() - store.unstable_resolve = (atom) => { - if (atom === (pseudo as Atom)) { + const unstable_resolve = >(atom: T): T => { + if (atom === (pseudo as Atom)) { return a as unknown as typeof atom } return atom } + const store = createStore().unstable_derive((atomStateMap) => [ + atomStateMap, + unstable_resolve, + ]) const pseudoWriteFn = vi.fn() const pseudo = atom('pseudo', pseudoWriteFn) as unknown as typeof a @@ -660,13 +671,18 @@ describe('unstable_resolve resolves the correct value for', () => { }) it('this in read and write', async () => { - const store = createStore() - store.unstable_resolve = (atom) => { - if (atom === (pseudo as Atom)) { + let unstable_resolve: + | (>(atom: T) => T) + | undefined = (atom) => { + if (atom === (pseudo as Atom)) { return this_read as unknown as typeof atom } return atom } + const store = createStore().unstable_derive((atomStateMap) => [ + atomStateMap, + (atom) => unstable_resolve?.(atom) ?? atom, + ]) const pseudo = atom('pseudo') as typeof this_read pseudo.debugLabel = 'pseudo' @@ -677,8 +693,8 @@ describe('unstable_resolve resolves the correct value for', () => { this_read.debugLabel = 'this_read' expect(store.get(pseudo)).toBe('this_read') - store.unstable_resolve = (atom) => { - if (atom === (pseudo as Atom)) { + unstable_resolve = (atom) => { + if (atom === (pseudo as Atom)) { return this_write as unknown as typeof atom } return atom @@ -707,7 +723,7 @@ describe('unstable_resolve resolves the correct value for', () => { expect(store.get(pseudo)).toBe( 'this:this_write-mounted:this_write-updated:this_write-unmounted', ) - delete store.unstable_resolve + unstable_resolve = undefined expect(store.get(pseudo)).toBe('pseudo') }) }) From 6f4f41d6624e8b828707e7b8b139baf3ca9558fc Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 10 Jul 2024 11:33:49 +0900 Subject: [PATCH 07/31] limit atom state map type --- src/vanilla/store.ts | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index f051790a3c..94631048cf 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -243,14 +243,19 @@ const flushPending = (pending: Pending) => { // for debugging purpose only type DevStoreRev4 = { - dev4_get_internal_weak_map: () => WeakMap + dev4_get_internal_weak_map: () => AtomStateMap dev4_get_mounted_atoms: () => Set dev4_restore_atoms: (values: Iterable) => void } +type AtomStateMap = { + get(key: Atom): AtomState | undefined + set(key: Atom, value: AtomState): void +} + // internal & unstable type type StoreArgs = readonly [ - atomStateMap: WeakMap, + atomStateMap: AtomStateMap, resolveAtom: (atom: T) => T, ] @@ -280,7 +285,7 @@ const buildStore = ( } const getAtomState = (atom: Atom) => { - let atomState = atomStateMap.get(atom) as AtomState | undefined + let atomState = atomStateMap.get(atom) if (!atomState) { atomState = { d: new Map(), p: new Set(), n: 0 } atomStateMap.set(atom, atomState) @@ -734,7 +739,7 @@ const buildStore = ( } export const createStore = (): Store => - buildStore(new WeakMap(), (atom) => atom) + buildStore(new WeakMap() as AtomStateMap, (atom) => atom) let defaultStore: Store | undefined From 18d55e187ccc42458b1cefbeaf11dd9782d292ba Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 10 Jul 2024 11:48:57 +0900 Subject: [PATCH 08/31] add test with atom state map --- tests/vanilla/store.test.tsx | 43 +++++++++++++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index ba4356f28f..4326f0634d 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -556,7 +556,48 @@ describe('aborting atoms', () => { }) }) -describe('unstable_derive resolves the correct value for', () => { +describe('unstable_derive with atomStateMap resolves the correct value for', () => { + it('primitive atom', async () => { + const a = atom('a') + a.onMount = (setSelf) => setSelf((v) => v + ':mounted') + const pseudo = atom('a') as typeof a + pseudo.onMount = (setSelf) => setSelf((v) => v + ':mounted') + + const unstable_resolve = >(atom: T): T => { + if (atom === (pseudo as Atom)) { + return a as unknown as typeof atom + } + return atom + } + const store = createStore() + const derivedStore = store.unstable_derive((atomStateMap, resolveAtom) => [ + { + get(key) { + return atomStateMap.get(unstable_resolve(key)) + }, + set(key, value) { + atomStateMap.set(unstable_resolve(key), value) + }, + }, + resolveAtom, // unchanged + ]) + + expect(store.get(pseudo)).toBe('a') + expect(derivedStore.get(pseudo)).toBe('a') + + derivedStore.sub(pseudo, vi.fn()) + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(pseudo)).toBe('a') + expect(derivedStore.get(pseudo)).toBe('a:mounted') + + derivedStore.set(pseudo, (v) => v + ':updated') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(pseudo)).toBe('a') + expect(derivedStore.get(pseudo)).toBe('a:mounted:updated') + }) +}) + +describe('unstable_derive with resolveAtom resolves the correct value for', () => { function nextTask() { return new Promise((resolve) => setTimeout(resolve)) } From b53f905720a8018ef0b876b986ba6ab325e826f2 Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 10 Jul 2024 12:41:16 +0900 Subject: [PATCH 09/31] update test using unstable_derive --- tests/vanilla/store.test.tsx | 57 +++++++++++++++++++----------------- 1 file changed, 30 insertions(+), 27 deletions(-) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 4326f0634d..ebccc7bcef 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -556,44 +556,47 @@ describe('aborting atoms', () => { }) }) -describe('unstable_derive with atomStateMap resolves the correct value for', () => { +describe('unstable_derive for scoping atoms', () => { it('primitive atom', async () => { const a = atom('a') a.onMount = (setSelf) => setSelf((v) => v + ':mounted') - const pseudo = atom('a') as typeof a - pseudo.onMount = (setSelf) => setSelf((v) => v + ':mounted') + const scopedAtoms = new Set>([a]) - const unstable_resolve = >(atom: T): T => { - if (atom === (pseudo as Atom)) { - return a as unknown as typeof atom - } - return atom - } const store = createStore() - const derivedStore = store.unstable_derive((atomStateMap, resolveAtom) => [ - { - get(key) { - return atomStateMap.get(unstable_resolve(key)) - }, - set(key, value) { - atomStateMap.set(unstable_resolve(key), value) + const derivedStore = store.unstable_derive((atomStateMap, resolveAtom) => { + const scopedAtomStateMap = new WeakMap() as typeof atomStateMap + return [ + { + get(key) { + if (scopedAtoms.has(key)) { + return scopedAtomStateMap.get(key) + } + return atomStateMap.get(key) + }, + set(key, value) { + if (scopedAtoms.has(key)) { + scopedAtomStateMap.set(key, value) + } else { + atomStateMap.set(key, value) + } + }, }, - }, - resolveAtom, // unchanged - ]) + resolveAtom, // unchanged + ] + }) - expect(store.get(pseudo)).toBe('a') - expect(derivedStore.get(pseudo)).toBe('a') + expect(store.get(a)).toBe('a') + expect(derivedStore.get(a)).toBe('a') - derivedStore.sub(pseudo, vi.fn()) + derivedStore.sub(a, vi.fn()) await new Promise((resolve) => setTimeout(resolve)) - expect(store.get(pseudo)).toBe('a') - expect(derivedStore.get(pseudo)).toBe('a:mounted') + expect(store.get(a)).toBe('a') + expect(derivedStore.get(a)).toBe('a:mounted') - derivedStore.set(pseudo, (v) => v + ':updated') + derivedStore.set(a, (v) => v + ':updated') await new Promise((resolve) => setTimeout(resolve)) - expect(store.get(pseudo)).toBe('a') - expect(derivedStore.get(pseudo)).toBe('a:mounted:updated') + expect(store.get(a)).toBe('a') + expect(derivedStore.get(a)).toBe('a:mounted:updated') }) }) From 22d2647bfc4d79be58311f53f22394372b2c7505 Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 10 Jul 2024 21:15:30 +0900 Subject: [PATCH 10/31] drop resolveAtom, add derive atom test --- src/vanilla/store.ts | 23 ++-- tests/vanilla/store.test.tsx | 196 ++++++----------------------------- 2 files changed, 36 insertions(+), 183 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 94631048cf..3fc8b87aa3 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -254,10 +254,7 @@ type AtomStateMap = { } // internal & unstable type -type StoreArgs = readonly [ - atomStateMap: AtomStateMap, - resolveAtom: (atom: T) => T, -] +type StoreArgs = readonly [atomStateMap: AtomStateMap] type PrdStore = { get: (atom: Atom) => Value @@ -273,10 +270,7 @@ type Store = PrdStore | (PrdStore & DevStoreRev4) export type INTERNAL_DevStoreRev4 = DevStoreRev4 export type INTERNAL_PrdStore = PrdStore -const buildStore = ( - atomStateMap: StoreArgs[0], - resolveAtom: StoreArgs[1], -): Store => { +const buildStore = (atomStateMap: StoreArgs[0]): Store => { // for debugging purpose only let debugMountedAtoms: Set @@ -393,7 +387,6 @@ const buildStore = ( atomState.d.clear() let isSync = true const getter: Getter = (a: Atom) => { - a = resolveAtom(a) if (a === (atom as AnyAtom)) { const aState = getAtomState(a) if (!isAtomStateInitialized(aState)) { @@ -474,7 +467,7 @@ const buildStore = ( } const readAtom = (atom: Atom): Value => - returnAtomValue(readAtomState(undefined, resolveAtom(atom))) + returnAtomValue(readAtomState(undefined, atom)) const recomputeDependents = (pending: Pending, atom: AnyAtom) => { const getDependents = (a: AnyAtom): Set => { @@ -548,13 +541,12 @@ const buildStore = ( ...args: Args ): Result => { const getter: Getter = (a: Atom) => - returnAtomValue(readAtomState(pending, resolveAtom(a))) + returnAtomValue(readAtomState(pending, a)) const setter: Setter = ( a: WritableAtom, ...args: As ) => { let r: R | undefined - a = resolveAtom(a) if (a === (atom as AnyAtom)) { if (!hasInitialValue(a)) { // NOTE technically possible but restricted as it may cause bugs @@ -585,7 +577,7 @@ const buildStore = ( ...args: Args ): Result => { const pending = createPending() - const result = writeAtomState(pending, resolveAtom(atom), ...args) + const result = writeAtomState(pending, atom, ...args) flushPending(pending) return result } @@ -684,7 +676,6 @@ const buildStore = ( } const subscribeAtom = (atom: AnyAtom, listener: () => void) => { - atom = resolveAtom(atom) const pending = createPending() const mounted = mountAtom(pending, atom) flushPending(pending) @@ -699,7 +690,7 @@ const buildStore = ( } const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) => { - const derivedArgs = fn(atomStateMap, resolveAtom) + const derivedArgs = fn(atomStateMap) const derivedStore = buildStore(...derivedArgs) return derivedStore } @@ -739,7 +730,7 @@ const buildStore = ( } export const createStore = (): Store => - buildStore(new WeakMap() as AtomStateMap, (atom) => atom) + buildStore(new WeakMap() as AtomStateMap) let defaultStore: Store | undefined diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index ebccc7bcef..db13748e26 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -563,7 +563,7 @@ describe('unstable_derive for scoping atoms', () => { const scopedAtoms = new Set>([a]) const store = createStore() - const derivedStore = store.unstable_derive((atomStateMap, resolveAtom) => { + const derivedStore = store.unstable_derive((atomStateMap) => { const scopedAtomStateMap = new WeakMap() as typeof atomStateMap return [ { @@ -581,7 +581,6 @@ describe('unstable_derive for scoping atoms', () => { } }, }, - resolveAtom, // unchanged ] }) @@ -598,176 +597,39 @@ describe('unstable_derive for scoping atoms', () => { expect(store.get(a)).toBe('a') expect(derivedStore.get(a)).toBe('a:mounted:updated') }) -}) - -describe('unstable_derive with resolveAtom resolves the correct value for', () => { - function nextTask() { - return new Promise((resolve) => setTimeout(resolve)) - } - - it('primitive atom', async () => { - const unstable_resolve = >(atom: T): T => { - if (atom === (pseudo as Atom)) { - return a as unknown as typeof atom - } - return atom - } - const store = createStore().unstable_derive((atomStateMap) => [ - atomStateMap, - unstable_resolve, - ]) - - const pseudo = atom('pseudo') as typeof a - pseudo.debugLabel = 'pseudo' - pseudo.onMount = (setSelf) => setSelf((v) => v + ':pseudo-mounted') - const a = atom('a') - a.debugLabel = 'a' - a.onMount = (setSelf) => setSelf((v) => v + ':a-mounted') - - expect(store.get(pseudo)).toBe('a') - const callback = vi.fn() - store.sub(pseudo, callback) - expect(store.get(pseudo)).toBe('a:a-mounted') - store.set(pseudo, (v) => v + ':a-updated') - expect(store.get(pseudo)).toBe('a:a-mounted:a-updated') - await nextTask() - expect(store.get(pseudo)).toBe('a:a-mounted:a-updated') - }) it('derived atom', async () => { - let unstable_resolve: - | (>(atom: T) => T) - | undefined = (atom) => { - if (atom === (pseudo as Atom)) { - return a as unknown as typeof atom - } - return atom - } - const store = createStore().unstable_derive((atomStateMap) => [ - atomStateMap, - (atom) => unstable_resolve?.(atom) ?? atom, - ]) - - const pseudo = atom('pseudo') as typeof a - pseudo.debugLabel = 'pseudo' - pseudo.onMount = (setSelf) => setSelf((v) => v + ':pseudo-mounted') const a = atom('a') - a.debugLabel = 'a' - a.onMount = (setSelf) => setSelf((v) => v + ':a-mounted') - const b = atom('b') - b.debugLabel = 'b' - const c = atom((get) => get(pseudo) + get(b)) - c.debugLabel = 'c' - expect(store.get(c)).toBe('ab') - const d = atom( - (_get, { setSelf }) => setTimeout(setSelf, 0, 'd'), - (get, set, v) => set(pseudo, get(pseudo) + v), - ) - store.get(d) - await nextTask() - expect(store.get(a)).toBe('ad') - expect(store.get(c)).toBe('adb') - expect(store.get(pseudo)).toEqual('ad') - const callback = vi.fn() - store.sub(c, callback) - expect(store.get(pseudo)).toBe('ad:a-mounted') - unstable_resolve = undefined - await nextTask() - expect(store.get(pseudo)).toEqual('pseudo') - store.sub(pseudo, callback) - expect(store.get(pseudo)).toEqual('pseudo:pseudo-mounted') - }) - - it('writable atom', async () => { - const unstable_resolve = >(atom: T): T => { - if (atom === (pseudo as Atom)) { - return a as unknown as typeof atom - } - return atom - } - const store = createStore().unstable_derive((atomStateMap) => [ - atomStateMap, - unstable_resolve, - ]) - - const pseudoWriteFn = vi.fn() - const pseudo = atom('pseudo', pseudoWriteFn) as unknown as typeof a - pseudo.debugLabel = 'pseudo' - pseudo.onMount = (setSelf) => setSelf('pseudo-mounted') - const a = atom('a', (get, set, value: string) => { - set(pseudo, get(pseudo) + ':' + value) - return () => set(pseudo, get(pseudo) + ':a-unmounted') - }) - a.debugLabel = 'a' - a.onMount = (setSelf) => setSelf('a-mounted') - expect(store.get(pseudo)).toBe('a') - const callback = vi.fn() - const unsub = store.sub(pseudo, callback) - await nextTask() - expect(store.get(pseudo)).toBe('a:a-mounted') - const value = store.set(pseudo, 'a-updated') - expect(pseudoWriteFn).not.toHaveBeenCalled() - expect(typeof value).toBe('function') - expect(store.get(pseudo)).toBe('a:a-mounted:a-updated') - unsub() - await nextTask() - expect(store.get(pseudo)).toBe('a:a-mounted:a-updated:a-unmounted') - }) + const b = atom((get) => get(a)) + const scopedAtoms = new Set>([a]) - it('this in read and write', async () => { - let unstable_resolve: - | (>(atom: T) => T) - | undefined = (atom) => { - if (atom === (pseudo as Atom)) { - return this_read as unknown as typeof atom - } - return atom - } - const store = createStore().unstable_derive((atomStateMap) => [ - atomStateMap, - (atom) => unstable_resolve?.(atom) ?? atom, - ]) - - const pseudo = atom('pseudo') as typeof this_read - pseudo.debugLabel = 'pseudo' - let i = 0 - const this_read = atom(function read(this: any, get) { - return i++ % 2 ? get(this) : 'this_read' + const store = createStore() + const derivedStore = store.unstable_derive((atomStateMap) => { + const scopedAtomStateMap = new WeakMap() as typeof atomStateMap + return [ + { + get(key) { + if (scopedAtoms.has(key)) { + return scopedAtomStateMap.get(key) + } + return atomStateMap.get(key) + }, + set(key, value) { + if (scopedAtoms.has(key)) { + scopedAtomStateMap.set(key, value) + } else { + atomStateMap.set(key, value) + } + }, + }, + ] }) - this_read.debugLabel = 'this_read' - expect(store.get(pseudo)).toBe('this_read') - unstable_resolve = (atom) => { - if (atom === (pseudo as Atom)) { - return this_write as unknown as typeof atom - } - return atom - } - - const this_write = atom( - 'this', - function write(this: any, get, set, value: string) { - set(this, get(this) + ':' + value) - return () => set(this, get(this) + ':this_write-unmounted') - }, - ) - this_write.debugLabel = 'this_write' - this_write.onMount = (setSelf) => setSelf('this_write-mounted') - expect(store.get(pseudo)).toBe('this') - const callback = vi.fn() - const unsub = store.sub(pseudo, callback) - await nextTask() - expect(store.get(pseudo)).toBe('this:this_write-mounted') - expect(callback).not.toHaveBeenCalledOnce() - store.set(this_write, 'this_write-updated') - expect(store.get(pseudo)).toBe('this:this_write-mounted:this_write-updated') - await nextTask() - unsub() - await nextTask() - expect(store.get(pseudo)).toBe( - 'this:this_write-mounted:this_write-updated:this_write-unmounted', - ) - unstable_resolve = undefined - expect(store.get(pseudo)).toBe('pseudo') + expect(store.get(b)).toBe('a') + expect(derivedStore.get(b)).toBe('a') + derivedStore.set(a, 'b') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(b)).toBe('a') + expect(derivedStore.get(b)).toBe('b') }) }) From bd9c6c0b6733854c8f1ad261d5b3a548f00157e3 Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 10 Jul 2024 21:21:25 +0900 Subject: [PATCH 11/31] simplify unstable_derive --- src/vanilla/store.ts | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 3fc8b87aa3..932aac0067 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -254,7 +254,10 @@ type AtomStateMap = { } // internal & unstable type -type StoreArgs = readonly [atomStateMap: AtomStateMap] +type StoreArgs = readonly [ + atomStateMap: AtomStateMap, + // possible other arguments in the future +] type PrdStore = { get: (atom: Atom) => Value @@ -689,11 +692,8 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { } } - const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) => { - const derivedArgs = fn(atomStateMap) - const derivedStore = buildStore(...derivedArgs) - return derivedStore - } + const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) => + buildStore(...fn(atomStateMap)) const store: Store = { get: readAtom, From 803b515011542ac802fcec9a737e76e78e8a7220 Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 16 Jul 2024 21:40:37 +0900 Subject: [PATCH 12/31] wip: this works but i do not like it --- src/vanilla/store.ts | 172 +++++++++++++++++++++-------------- tests/vanilla/store.test.tsx | 80 +++++++++++++++- 2 files changed, 182 insertions(+), 70 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 932aac0067..84430b0f52 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -241,24 +241,30 @@ const flushPending = (pending: Pending) => { } } -// for debugging purpose only -type DevStoreRev4 = { - dev4_get_internal_weak_map: () => AtomStateMap - dev4_get_mounted_atoms: () => Set - dev4_restore_atoms: (values: Iterable) => void -} - type AtomStateMap = { get(key: Atom): AtomState | undefined set(key: Atom, value: AtomState): void } +type GetAtomState = ( + atomStateMap: AtomStateMap, + atom: Atom, + context: unknown, +) => [atomState: AtomState, context: unknown] + // internal & unstable type type StoreArgs = readonly [ atomStateMap: AtomStateMap, - // possible other arguments in the future + getAtomState: GetAtomState, ] +// for debugging purpose only +type DevStoreRev4 = { + dev4_get_internal_weak_map: () => AtomStateMap + dev4_get_mounted_atoms: () => Set + dev4_restore_atoms: (values: Iterable) => void +} + type PrdStore = { get: (atom: Atom) => Value set: ( @@ -268,12 +274,16 @@ type PrdStore = { sub: (atom: AnyAtom, listener: () => void) => () => void unstable_derive: (fn: (...args: StoreArgs) => StoreArgs) => Store } + type Store = PrdStore | (PrdStore & DevStoreRev4) export type INTERNAL_DevStoreRev4 = DevStoreRev4 export type INTERNAL_PrdStore = PrdStore -const buildStore = (atomStateMap: StoreArgs[0]): Store => { +const buildStore = ( + atomStateMap: StoreArgs[0], + getAtomState: StoreArgs[1], +): Store => { // for debugging purpose only let debugMountedAtoms: Set @@ -281,17 +291,9 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { debugMountedAtoms = new Set() } - const getAtomState = (atom: Atom) => { - let atomState = atomStateMap.get(atom) - if (!atomState) { - atomState = { d: new Map(), p: new Set(), n: 0 } - atomStateMap.set(atom, atomState) - } - return atomState - } - const setAtomStateValueOrPromise = ( atom: AnyAtom, + context: unknown, atomState: AtomState, valueOrPromise: unknown, abortPromise = () => {}, @@ -314,7 +316,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { ) if (continuablePromise.status === PENDING) { for (const a of atomState.d.keys()) { - const aState = getAtomState(a) + const [aState] = getAtomState(atomStateMap, a, context) addPendingContinuablePromiseToDependency( atom, continuablePromise, @@ -342,13 +344,14 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const addDependency = ( pending: Pending | undefined, atom: Atom, + context: unknown, a: AnyAtom, aState: AtomState, ) => { if (import.meta.env?.MODE !== 'production' && a === atom) { throw new Error('[Bug] atom cannot depend on itself') } - const atomState = getAtomState(atom) + const [atomState] = getAtomState(atomStateMap, atom, context) atomState.d.set(a, aState.n) const continuablePromise = getPendingContinuablePromise(atomState) if (continuablePromise) { @@ -363,10 +366,11 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const readAtomState = ( pending: Pending | undefined, atom: Atom, + context: unknown, force?: (a: AnyAtom) => boolean, ): AtomState => { + const [atomState, nextContext] = getAtomState(atomStateMap, atom, context) // See if we can skip recomputing this atom. - const atomState = getAtomState(atom) if (!force?.(atom) && isAtomStateInitialized(atomState)) { // If the atom is mounted, we can use the cache. // because it should have been updated by dependencies. @@ -380,7 +384,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { ([a, n]) => // Recursively, read the atom state of the dependency, and // check if the atom epoch number is unchanged - readAtomState(pending, a, force).n === n, + readAtomState(pending, a, nextContext, force).n === n, ) ) { return atomState @@ -391,10 +395,10 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { let isSync = true const getter: Getter = (a: Atom) => { if (a === (atom as AnyAtom)) { - const aState = getAtomState(a) + const [aState] = getAtomState(atomStateMap, a, nextContext) if (!isAtomStateInitialized(aState)) { if (hasInitialValue(a)) { - setAtomStateValueOrPromise(a, aState, a.init) + setAtomStateValueOrPromise(a, nextContext, aState, a.init) } else { // NOTE invalid derived atoms can reach here throw new Error('no atom init') @@ -403,13 +407,13 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { return returnAtomValue(aState) } // a !== atom - const aState = readAtomState(pending, a, force) + const aState = readAtomState(pending, a, nextContext, force) if (isSync) { - addDependency(pending, atom, a, aState) + addDependency(pending, atom, nextContext, a, aState) } else { const pending = createPending() - addDependency(pending, atom, a, aState) - mountDependencies(pending, atom, atomState) + addDependency(pending, atom, nextContext, a, aState) + mountDependencies(pending, atom, nextContext, atomState) flushPending(pending) } return returnAtomValue(aState) @@ -447,13 +451,14 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const valueOrPromise = atom.read(getter, options as never) setAtomStateValueOrPromise( atom, + nextContext, atomState, valueOrPromise, () => controller?.abort(), () => { if (atomState.m) { const pending = createPending() - mountDependencies(pending, atom, atomState) + mountDependencies(pending, atom, nextContext, atomState) flushPending(pending) } }, @@ -470,11 +475,15 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { } const readAtom = (atom: Atom): Value => - returnAtomValue(readAtomState(undefined, atom)) + returnAtomValue(readAtomState(undefined, atom, undefined)) - const recomputeDependents = (pending: Pending, atom: AnyAtom) => { - const getDependents = (a: AnyAtom): Set => { - const aState = getAtomState(a) + const recomputeDependents = ( + pending: Pending, + atom: AnyAtom, + context: unknown, + ) => { + const getDependents = (a: AnyAtom) => { + const [aState, nextContext] = getAtomState(atomStateMap, a, context) const dependents = new Set(aState.m?.t) for (const atomWithPendingContinuablePromise of aState.p) { dependents.add(atomWithPendingContinuablePromise) @@ -482,7 +491,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { getPendingDependents(pending, a)?.forEach((dependent) => { dependents.add(dependent) }) - return dependents + return [dependents, nextContext] as const } // This is a topological sort via depth-first search, slightly modified from @@ -491,14 +500,15 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { // Step 1: traverse the dependency graph to build the topsorted atom list // We don't bother to check for cycles, which simplifies the algorithm. - const topsortedAtoms: AnyAtom[] = [] + const topsortedAtoms: (readonly [atom: AnyAtom, context: unknown])[] = [] const markedAtoms = new Set() const visit = (n: AnyAtom) => { if (markedAtoms.has(n)) { return } markedAtoms.add(n) - for (const m of getDependents(n)) { + const [dependents, nextContext] = getDependents(n) + for (const m of dependents) { if (n !== m) { visit(m) } @@ -506,7 +516,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { // The algorithm calls for pushing onto the front of the list. For // performance, we will simply push onto the end, and then will iterate in // reverse order later. - topsortedAtoms.push(n) + topsortedAtoms.push([n, nextContext]) } // Visit the root atom. This is the only atom in the dependency graph // without incoming edges, which is one reason we can simplify the algorithm @@ -516,8 +526,8 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const changedAtoms = new Set([atom]) const isMarked = (a: AnyAtom) => markedAtoms.has(a) for (let i = topsortedAtoms.length - 1; i >= 0; --i) { - const a = topsortedAtoms[i]! - const aState = getAtomState(a) + const [a, nextContext] = topsortedAtoms[i]! + const [aState] = getAtomState(atomStateMap, a, nextContext) const prevEpochNumber = aState.n let hasChangedDeps = false for (const dep of aState.d.keys()) { @@ -527,8 +537,8 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { } } if (hasChangedDeps) { - readAtomState(pending, a, isMarked) - mountDependencies(pending, a, aState) + readAtomState(pending, a, nextContext, isMarked) + mountDependencies(pending, a, nextContext, aState) if (prevEpochNumber !== aState.n) { addPendingAtom(pending, a, aState) changedAtoms.add(a) @@ -541,10 +551,11 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const writeAtomState = ( pending: Pending, atom: WritableAtom, + context: unknown, ...args: Args ): Result => { const getter: Getter = (a: Atom) => - returnAtomValue(readAtomState(pending, a)) + returnAtomValue(readAtomState(pending, a, context)) const setter: Setter = ( a: WritableAtom, ...args: As @@ -555,18 +566,18 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { // NOTE technically possible but restricted as it may cause bugs throw new Error('atom not writable') } - const aState = getAtomState(a) + const [aState, nextContext] = getAtomState(atomStateMap, a, context) const hasPrevValue = 'v' in aState const prevValue = aState.v const v = args[0] as V - setAtomStateValueOrPromise(a, aState, v) - mountDependencies(pending, a, aState) + setAtomStateValueOrPromise(a, nextContext, aState, v) + mountDependencies(pending, a, nextContext, aState) if (!hasPrevValue || !Object.is(prevValue, aState.v)) { addPendingAtom(pending, a, aState) - recomputeDependents(pending, a) + recomputeDependents(pending, a, nextContext) } } else { - r = writeAtomState(pending, a, ...args) as R + r = writeAtomState(pending, a, context, ...args) as R } flushPending(pending) return r as R @@ -580,7 +591,12 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { ...args: Args ): Result => { const pending = createPending() - const result = writeAtomState(pending, atom, ...args) + const result = writeAtomState( + pending, + atom, + getAtomState(atomStateMap, atom, undefined)[1], + ...args, + ) flushPending(pending) return result } @@ -588,19 +604,20 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const mountDependencies = ( pending: Pending, atom: AnyAtom, + context: unknown, atomState: AtomState, ) => { if (atomState.m && !getPendingContinuablePromise(atomState)) { for (const a of atomState.d.keys()) { if (!atomState.m.d.has(a)) { - const aMounted = mountAtom(pending, a) + const aMounted = mountAtom(pending, a, context) aMounted.t.add(atom) atomState.m.d.add(a) } } for (const a of atomState.m.d || []) { if (!atomState.d.has(a)) { - const aMounted = unmountAtom(pending, a) + const aMounted = unmountAtom(pending, a, context) aMounted?.t.delete(atom) atomState.m.d.delete(a) } @@ -608,14 +625,18 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { } } - const mountAtom = (pending: Pending, atom: AnyAtom): Mounted => { - const atomState = getAtomState(atom) + const mountAtom = ( + pending: Pending, + atom: AnyAtom, + context: unknown, + ): Mounted => { + const [atomState, nextContext] = getAtomState(atomStateMap, atom, context) if (!atomState.m) { // recompute atom state - readAtomState(pending, atom) + readAtomState(pending, atom, nextContext) // mount dependencies first for (const a of atomState.d.keys()) { - const aMounted = mountAtom(pending, a) + const aMounted = mountAtom(pending, a, nextContext) aMounted.t.add(atom) } // mount self @@ -632,7 +653,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const { onMount } = atom addPendingFunction(pending, () => { const onUnmount = onMount((...args) => - writeAtomState(pending, atom, ...args), + writeAtomState(pending, atom, nextContext, ...args), ) if (onUnmount) { mounted.u = onUnmount @@ -646,12 +667,15 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const unmountAtom = ( pending: Pending, atom: AnyAtom, + context: unknown, ): Mounted | undefined => { - const atomState = getAtomState(atom) + const [atomState, nextContext] = getAtomState(atomStateMap, atom, context) if ( atomState.m && !atomState.m.l.size && - !Array.from(atomState.m.t).some((a) => getAtomState(a).m) + !Array.from(atomState.m.t).some( + (a) => getAtomState(atomStateMap, a, nextContext)[0].m, + ) ) { // unmount self const onUnmount = atomState.m.u @@ -664,7 +688,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { } // unmount dependencies for (const a of atomState.d.keys()) { - const aMounted = unmountAtom(pending, a) + const aMounted = unmountAtom(pending, a, nextContext) aMounted?.t.delete(atom) } // abort pending promise @@ -680,20 +704,20 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const subscribeAtom = (atom: AnyAtom, listener: () => void) => { const pending = createPending() - const mounted = mountAtom(pending, atom) + const mounted = mountAtom(pending, atom, undefined) flushPending(pending) const listeners = mounted.l listeners.add(listener) return () => { listeners.delete(listener) const pending = createPending() - unmountAtom(pending, atom) + unmountAtom(pending, atom, undefined) flushPending(pending) } } const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) => - buildStore(...fn(atomStateMap)) + buildStore(...fn(atomStateMap, getAtomState)) const store: Store = { get: readAtom, @@ -710,14 +734,18 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const pending = createPending() for (const [atom, value] of values) { if (hasInitialValue(atom)) { - const aState = getAtomState(atom) + const [aState, context] = getAtomState( + atomStateMap, + atom, + undefined, + ) const hasPrevValue = 'v' in aState const prevValue = aState.v - setAtomStateValueOrPromise(atom, aState, value) - mountDependencies(pending, atom, aState) + setAtomStateValueOrPromise(atom, context, aState, value) + mountDependencies(pending, atom, context, aState) if (!hasPrevValue || !Object.is(prevValue, aState.v)) { addPendingAtom(pending, atom, aState) - recomputeDependents(pending, atom) + recomputeDependents(pending, atom, context) } } } @@ -729,8 +757,18 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { return store } -export const createStore = (): Store => - buildStore(new WeakMap() as AtomStateMap) +export const createStore = (): Store => { + const atomStateMap = new WeakMap() as AtomStateMap + const getAtomState: GetAtomState = (atomStateMap, atom, context) => { + let atomState = atomStateMap.get(atom) + if (!atomState) { + atomState = { d: new Map(), p: new Set(), n: 0 } + atomStateMap.set(atom, atomState) + } + return [atomState, context] + } + return buildStore(atomStateMap, getAtomState) +} let defaultStore: Store | undefined diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index db13748e26..99e4a29966 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -563,7 +563,7 @@ describe('unstable_derive for scoping atoms', () => { const scopedAtoms = new Set>([a]) const store = createStore() - const derivedStore = store.unstable_derive((atomStateMap) => { + const derivedStore = store.unstable_derive((atomStateMap, getAtomState) => { const scopedAtomStateMap = new WeakMap() as typeof atomStateMap return [ { @@ -581,6 +581,7 @@ describe('unstable_derive for scoping atoms', () => { } }, }, + getAtomState, ] }) @@ -598,13 +599,13 @@ describe('unstable_derive for scoping atoms', () => { expect(derivedStore.get(a)).toBe('a:mounted:updated') }) - it('derived atom', async () => { + it('derived atom (scoping primitive)', async () => { const a = atom('a') const b = atom((get) => get(a)) const scopedAtoms = new Set>([a]) const store = createStore() - const derivedStore = store.unstable_derive((atomStateMap) => { + const derivedStore = store.unstable_derive((atomStateMap, getAtomState) => { const scopedAtomStateMap = new WeakMap() as typeof atomStateMap return [ { @@ -622,6 +623,7 @@ describe('unstable_derive for scoping atoms', () => { } }, }, + getAtomState, ] }) @@ -632,4 +634,76 @@ describe('unstable_derive for scoping atoms', () => { expect(store.get(b)).toBe('a') expect(derivedStore.get(b)).toBe('b') }) + + it('derived atom (scoping derived)', async () => { + const a = atom('a') + const b = atom( + (get) => get(a), + (_get, set, v: string) => { + set(a, v) + }, + ) + const scopedAtoms = new Set>([b]) + + const store = createStore() + const derivedStore = store.unstable_derive((atomStateMap, getAtomState) => { + const scopedAtomStateMap = new WeakMap() as typeof atomStateMap + return [ + atomStateMap, + (atomStateMap, atom, unknownContext: unknown) => { + let context = unknownContext as + | { scope?: typeof scopedAtomStateMap } + | undefined + if (scopedAtoms.has(atom)) { + context = { + ...context, + scope: scopedAtomStateMap, + } + } + if (context?.scope !== scopedAtomStateMap) { + return getAtomState(atomStateMap, atom, context) + } + let atomState = context.scope.get(atom) + if (!atomState) { + atomState = { d: new Map(), p: new Set(), n: 0 } + context.scope.set(atom, atomState) + } + return [atomState, context] + }, + ] + }) + + expect(store.get(a)).toBe('a') + expect(store.get(b)).toBe('a') + expect(derivedStore.get(a)).toBe('a') + expect(derivedStore.get(b)).toBe('a') + + store.set(a, 'a2') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(a)).toBe('a2') + expect(store.get(b)).toBe('a2') + expect(derivedStore.get(a)).toBe('a2') + expect(derivedStore.get(b)).toBe('a') + + store.set(b, 'a3') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(a)).toBe('a3') + expect(store.get(b)).toBe('a3') + expect(derivedStore.get(a)).toBe('a3') + expect(derivedStore.get(b)).toBe('a') + + derivedStore.set(a, 'a4') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(a)).toBe('a4') + expect(store.get(b)).toBe('a4') + expect(derivedStore.get(a)).toBe('a4') + expect(derivedStore.get(b)).toBe('a') + + derivedStore.set(b, 'a5') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(a)).toBe('a4') + expect(store.get(b)).toBe('a4') + expect(derivedStore.get(a)).toBe('a4') + expect(derivedStore.get(b)).toBe('a5') + }) }) From 849741094c1cddd012b1bd0d71c09d381d85561b Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 16 Jul 2024 23:22:53 +0900 Subject: [PATCH 13/31] separate atomstate and atomcontext --- src/vanilla/store.ts | 154 +++++++++++++++++++---------------- tests/vanilla/store.test.tsx | 137 ++++++++++++++++--------------- 2 files changed, 156 insertions(+), 135 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 84430b0f52..26ecc78dee 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -149,6 +149,12 @@ type AtomState = { e?: AnyError } +const createDefaultAtomState = (): AtomState => ({ + d: new Map(), + p: new Set(), + n: 0, +}) + const isAtomStateInitialized = (atomState: AtomState) => 'v' in atomState || 'e' in atomState @@ -241,26 +247,28 @@ const flushPending = (pending: Pending) => { } } -type AtomStateMap = { - get(key: Atom): AtomState | undefined - set(key: Atom, value: AtomState): void +type AtomContext = unknown + +type GetAtomState = { + (atom: Atom): AtomState | undefined + ( + atom: Atom, + context: AtomContext, + createDefault: () => AtomState, + ): AtomState } -type GetAtomState = ( - atomStateMap: AtomStateMap, - atom: Atom, - context: unknown, -) => [atomState: AtomState, context: unknown] +type GetAtomContext = (atom: AnyAtom, prevContext?: AtomContext) => AtomContext // internal & unstable type type StoreArgs = readonly [ - atomStateMap: AtomStateMap, getAtomState: GetAtomState, + getAtomContext: GetAtomContext, ] // for debugging purpose only type DevStoreRev4 = { - dev4_get_internal_weak_map: () => AtomStateMap + dev4_get_internal_weak_map: () => { get: GetAtomState } dev4_get_mounted_atoms: () => Set dev4_restore_atoms: (values: Iterable) => void } @@ -281,8 +289,8 @@ export type INTERNAL_DevStoreRev4 = DevStoreRev4 export type INTERNAL_PrdStore = PrdStore const buildStore = ( - atomStateMap: StoreArgs[0], - getAtomState: StoreArgs[1], + getAtomState: StoreArgs[0], + getAtomContext: StoreArgs[1], ): Store => { // for debugging purpose only let debugMountedAtoms: Set @@ -293,7 +301,7 @@ const buildStore = ( const setAtomStateValueOrPromise = ( atom: AnyAtom, - context: unknown, + context: AtomContext, atomState: AtomState, valueOrPromise: unknown, abortPromise = () => {}, @@ -316,7 +324,7 @@ const buildStore = ( ) if (continuablePromise.status === PENDING) { for (const a of atomState.d.keys()) { - const [aState] = getAtomState(atomStateMap, a, context) + const aState = getAtomState(a, context, createDefaultAtomState) addPendingContinuablePromiseToDependency( atom, continuablePromise, @@ -344,14 +352,14 @@ const buildStore = ( const addDependency = ( pending: Pending | undefined, atom: Atom, - context: unknown, + context: AtomContext, a: AnyAtom, aState: AtomState, ) => { if (import.meta.env?.MODE !== 'production' && a === atom) { throw new Error('[Bug] atom cannot depend on itself') } - const [atomState] = getAtomState(atomStateMap, atom, context) + const atomState = getAtomState(atom, context, createDefaultAtomState) atomState.d.set(a, aState.n) const continuablePromise = getPendingContinuablePromise(atomState) if (continuablePromise) { @@ -366,10 +374,10 @@ const buildStore = ( const readAtomState = ( pending: Pending | undefined, atom: Atom, - context: unknown, + context: AtomContext, force?: (a: AnyAtom) => boolean, ): AtomState => { - const [atomState, nextContext] = getAtomState(atomStateMap, atom, context) + const atomState = getAtomState(atom, context, createDefaultAtomState) // See if we can skip recomputing this atom. if (!force?.(atom) && isAtomStateInitialized(atomState)) { // If the atom is mounted, we can use the cache. @@ -384,7 +392,8 @@ const buildStore = ( ([a, n]) => // Recursively, read the atom state of the dependency, and // check if the atom epoch number is unchanged - readAtomState(pending, a, nextContext, force).n === n, + readAtomState(pending, a, getAtomContext(a, context), force).n === + n, ) ) { return atomState @@ -394,8 +403,9 @@ const buildStore = ( atomState.d.clear() let isSync = true const getter: Getter = (a: Atom) => { + const nextContext = getAtomContext(a, context) if (a === (atom as AnyAtom)) { - const [aState] = getAtomState(atomStateMap, a, nextContext) + const aState = getAtomState(a, nextContext, createDefaultAtomState) if (!isAtomStateInitialized(aState)) { if (hasInitialValue(a)) { setAtomStateValueOrPromise(a, nextContext, aState, a.init) @@ -451,14 +461,14 @@ const buildStore = ( const valueOrPromise = atom.read(getter, options as never) setAtomStateValueOrPromise( atom, - nextContext, + context, atomState, valueOrPromise, () => controller?.abort(), () => { if (atomState.m) { const pending = createPending() - mountDependencies(pending, atom, nextContext, atomState) + mountDependencies(pending, atom, context, atomState) flushPending(pending) } }, @@ -475,15 +485,16 @@ const buildStore = ( } const readAtom = (atom: Atom): Value => - returnAtomValue(readAtomState(undefined, atom, undefined)) + returnAtomValue(readAtomState(undefined, atom, getAtomContext(atom))) const recomputeDependents = ( pending: Pending, atom: AnyAtom, - context: unknown, + context: AtomContext, ) => { - const getDependents = (a: AnyAtom) => { - const [aState, nextContext] = getAtomState(atomStateMap, a, context) + const getDependents = (a: AnyAtom, context: AtomContext) => { + const aState = getAtomState(a, context, createDefaultAtomState) + const nextContext = getAtomContext(a, context) const dependents = new Set(aState.m?.t) for (const atomWithPendingContinuablePromise of aState.p) { dependents.add(atomWithPendingContinuablePromise) @@ -500,34 +511,35 @@ const buildStore = ( // Step 1: traverse the dependency graph to build the topsorted atom list // We don't bother to check for cycles, which simplifies the algorithm. - const topsortedAtoms: (readonly [atom: AnyAtom, context: unknown])[] = [] + const topsortedAtoms: (readonly [atom: AnyAtom, context: AtomContext])[] = + [] const markedAtoms = new Set() - const visit = (n: AnyAtom) => { + const visit = (n: AnyAtom, context: AtomContext) => { if (markedAtoms.has(n)) { return } markedAtoms.add(n) - const [dependents, nextContext] = getDependents(n) + const [dependents, nextContext] = getDependents(n, context) for (const m of dependents) { if (n !== m) { - visit(m) + visit(m, nextContext) } } // The algorithm calls for pushing onto the front of the list. For // performance, we will simply push onto the end, and then will iterate in // reverse order later. - topsortedAtoms.push([n, nextContext]) + topsortedAtoms.push([n, context]) } // Visit the root atom. This is the only atom in the dependency graph // without incoming edges, which is one reason we can simplify the algorithm - visit(atom) + visit(atom, context) // Step 2: use the topsorted atom list to recompute all affected atoms // Track what's changed, so that we can short circuit when possible const changedAtoms = new Set([atom]) const isMarked = (a: AnyAtom) => markedAtoms.has(a) for (let i = topsortedAtoms.length - 1; i >= 0; --i) { const [a, nextContext] = topsortedAtoms[i]! - const [aState] = getAtomState(atomStateMap, a, nextContext) + const aState = getAtomState(a, nextContext, createDefaultAtomState) const prevEpochNumber = aState.n let hasChangedDeps = false for (const dep of aState.d.keys()) { @@ -551,22 +563,23 @@ const buildStore = ( const writeAtomState = ( pending: Pending, atom: WritableAtom, - context: unknown, + context: AtomContext, ...args: Args ): Result => { const getter: Getter = (a: Atom) => - returnAtomValue(readAtomState(pending, a, context)) + returnAtomValue(readAtomState(pending, a, getAtomContext(a, context))) const setter: Setter = ( a: WritableAtom, ...args: As ) => { + const nextContext = getAtomContext(a, context) let r: R | undefined if (a === (atom as AnyAtom)) { if (!hasInitialValue(a)) { // NOTE technically possible but restricted as it may cause bugs throw new Error('atom not writable') } - const [aState, nextContext] = getAtomState(atomStateMap, a, context) + const aState = getAtomState(a, nextContext, createDefaultAtomState) const hasPrevValue = 'v' in aState const prevValue = aState.v const v = args[0] as V @@ -577,7 +590,7 @@ const buildStore = ( recomputeDependents(pending, a, nextContext) } } else { - r = writeAtomState(pending, a, context, ...args) as R + r = writeAtomState(pending, a, nextContext, ...args) as R } flushPending(pending) return r as R @@ -591,12 +604,7 @@ const buildStore = ( ...args: Args ): Result => { const pending = createPending() - const result = writeAtomState( - pending, - atom, - getAtomState(atomStateMap, atom, undefined)[1], - ...args, - ) + const result = writeAtomState(pending, atom, getAtomContext(atom), ...args) flushPending(pending) return result } @@ -604,7 +612,7 @@ const buildStore = ( const mountDependencies = ( pending: Pending, atom: AnyAtom, - context: unknown, + context: AtomContext, atomState: AtomState, ) => { if (atomState.m && !getPendingContinuablePromise(atomState)) { @@ -628,15 +636,15 @@ const buildStore = ( const mountAtom = ( pending: Pending, atom: AnyAtom, - context: unknown, + context: AtomContext, ): Mounted => { - const [atomState, nextContext] = getAtomState(atomStateMap, atom, context) + const atomState = getAtomState(atom, context, createDefaultAtomState) if (!atomState.m) { // recompute atom state - readAtomState(pending, atom, nextContext) + readAtomState(pending, atom, context) // mount dependencies first for (const a of atomState.d.keys()) { - const aMounted = mountAtom(pending, a, nextContext) + const aMounted = mountAtom(pending, a, getAtomContext(a, context)) aMounted.t.add(atom) } // mount self @@ -653,7 +661,7 @@ const buildStore = ( const { onMount } = atom addPendingFunction(pending, () => { const onUnmount = onMount((...args) => - writeAtomState(pending, atom, nextContext, ...args), + writeAtomState(pending, atom, context, ...args), ) if (onUnmount) { mounted.u = onUnmount @@ -667,14 +675,15 @@ const buildStore = ( const unmountAtom = ( pending: Pending, atom: AnyAtom, - context: unknown, + context: AtomContext, ): Mounted | undefined => { - const [atomState, nextContext] = getAtomState(atomStateMap, atom, context) + const atomState = getAtomState(atom, context, createDefaultAtomState) if ( atomState.m && !atomState.m.l.size && !Array.from(atomState.m.t).some( - (a) => getAtomState(atomStateMap, a, nextContext)[0].m, + (a) => + getAtomState(a, getAtomContext(a, context), createDefaultAtomState).m, ) ) { // unmount self @@ -688,7 +697,7 @@ const buildStore = ( } // unmount dependencies for (const a of atomState.d.keys()) { - const aMounted = unmountAtom(pending, a, nextContext) + const aMounted = unmountAtom(pending, a, getAtomContext(a, context)) aMounted?.t.delete(atom) } // abort pending promise @@ -717,7 +726,7 @@ const buildStore = ( } const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) => - buildStore(...fn(atomStateMap, getAtomState)) + buildStore(...fn(getAtomState, getAtomContext)) const store: Store = { get: readAtom, @@ -728,23 +737,24 @@ const buildStore = ( if (import.meta.env?.MODE !== 'production') { const devStore: DevStoreRev4 = { // store dev methods (these are tentative and subject to change without notice) - dev4_get_internal_weak_map: () => atomStateMap, + dev4_get_internal_weak_map: () => ({ get: getAtomState }), dev4_get_mounted_atoms: () => debugMountedAtoms, dev4_restore_atoms: (values) => { const pending = createPending() for (const [atom, value] of values) { if (hasInitialValue(atom)) { - const [aState, context] = getAtomState( - atomStateMap, + const context = getAtomContext(atom) + const atomState = getAtomState( atom, - undefined, + context, + createDefaultAtomState, ) - const hasPrevValue = 'v' in aState - const prevValue = aState.v - setAtomStateValueOrPromise(atom, context, aState, value) - mountDependencies(pending, atom, context, aState) - if (!hasPrevValue || !Object.is(prevValue, aState.v)) { - addPendingAtom(pending, atom, aState) + const hasPrevValue = 'v' in atomState + const prevValue = atomState.v + setAtomStateValueOrPromise(atom, context, atomState, value) + mountDependencies(pending, atom, context, atomState) + if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { + addPendingAtom(pending, atom, atomState) recomputeDependents(pending, atom, context) } } @@ -758,16 +768,20 @@ const buildStore = ( } export const createStore = (): Store => { - const atomStateMap = new WeakMap() as AtomStateMap - const getAtomState: GetAtomState = (atomStateMap, atom, context) => { + const atomStateMap = new WeakMap() + const getAtomState = ((atom, _context, createDefault) => { let atomState = atomStateMap.get(atom) if (!atomState) { - atomState = { d: new Map(), p: new Set(), n: 0 } + if (!createDefault) { + return undefined + } + atomState = createDefault() atomStateMap.set(atom, atomState) } - return [atomState, context] - } - return buildStore(atomStateMap, getAtomState) + return atomState + }) as GetAtomState + const getAtomContext: GetAtomContext = (_atom, context) => context + return buildStore(getAtomState, getAtomContext) } let defaultStore: Store | undefined diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 99e4a29966..7a46dc7979 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -563,27 +563,28 @@ describe('unstable_derive for scoping atoms', () => { const scopedAtoms = new Set>([a]) const store = createStore() - const derivedStore = store.unstable_derive((atomStateMap, getAtomState) => { - const scopedAtomStateMap = new WeakMap() as typeof atomStateMap - return [ - { - get(key) { - if (scopedAtoms.has(key)) { - return scopedAtomStateMap.get(key) + const derivedStore = store.unstable_derive( + (getAtomState, getAtomContext) => { + const scopedAtomStateMap = new WeakMap() + return [ + ((atom, context, createDefault) => { + if (scopedAtoms.has(atom)) { + let atomState = scopedAtomStateMap.get(atom) + if (!atomState) { + if (!createDefault) { + return undefined + } + atomState = createDefault() + scopedAtomStateMap.set(atom, atomState) + } + return atomState } - return atomStateMap.get(key) - }, - set(key, value) { - if (scopedAtoms.has(key)) { - scopedAtomStateMap.set(key, value) - } else { - atomStateMap.set(key, value) - } - }, - }, - getAtomState, - ] - }) + return getAtomState(atom, context, createDefault) + }) as typeof getAtomState, + getAtomContext, + ] + }, + ) expect(store.get(a)).toBe('a') expect(derivedStore.get(a)).toBe('a') @@ -605,27 +606,28 @@ describe('unstable_derive for scoping atoms', () => { const scopedAtoms = new Set>([a]) const store = createStore() - const derivedStore = store.unstable_derive((atomStateMap, getAtomState) => { - const scopedAtomStateMap = new WeakMap() as typeof atomStateMap - return [ - { - get(key) { - if (scopedAtoms.has(key)) { - return scopedAtomStateMap.get(key) - } - return atomStateMap.get(key) - }, - set(key, value) { - if (scopedAtoms.has(key)) { - scopedAtomStateMap.set(key, value) - } else { - atomStateMap.set(key, value) + const derivedStore = store.unstable_derive( + (getAtomState, getAtomContext) => { + const scopedAtomStateMap = new WeakMap() + return [ + ((atom, context, createDefault) => { + if (scopedAtoms.has(atom)) { + let atomState = scopedAtomStateMap.get(atom) + if (!atomState) { + if (!createDefault) { + return undefined + } + atomState = createDefault() + scopedAtomStateMap.set(atom, atomState) + } + return atomState } - }, - }, - getAtomState, - ] - }) + return getAtomState(atom, context, createDefault) + }) as typeof getAtomState, + getAtomContext, + ] + }, + ) expect(store.get(b)).toBe('a') expect(derivedStore.get(b)).toBe('a') @@ -646,32 +648,37 @@ describe('unstable_derive for scoping atoms', () => { const scopedAtoms = new Set>([b]) const store = createStore() - const derivedStore = store.unstable_derive((atomStateMap, getAtomState) => { - const scopedAtomStateMap = new WeakMap() as typeof atomStateMap - return [ - atomStateMap, - (atomStateMap, atom, unknownContext: unknown) => { - let context = unknownContext as - | { scope?: typeof scopedAtomStateMap } - | undefined - if (scopedAtoms.has(atom)) { - context = { - ...context, - scope: scopedAtomStateMap, + const derivedStore = store.unstable_derive( + (getAtomState, getAtomContext) => { + const scopedAtomStateMap = new WeakMap() + return [ + ((atom, context, createDefault) => { + if ( + (context as { scoped: unknown } | undefined)?.scoped || + scopedAtoms.has(atom) + ) { + let atomState = scopedAtomStateMap.get(atom) + if (!atomState) { + if (!createDefault) { + return undefined + } + atomState = createDefault() + scopedAtomStateMap.set(atom, atomState) + } + return atomState } - } - if (context?.scope !== scopedAtomStateMap) { - return getAtomState(atomStateMap, atom, context) - } - let atomState = context.scope.get(atom) - if (!atomState) { - atomState = { d: new Map(), p: new Set(), n: 0 } - context.scope.set(atom, atomState) - } - return [atomState, context] - }, - ] - }) + return getAtomState(atom, context, createDefault) + }) as typeof getAtomState, + (atom, context) => { + context = getAtomContext(atom, context) + if (scopedAtoms.has(atom)) { + return { ...(context as object), scoped: true } + } + return context + }, + ] + }, + ) expect(store.get(a)).toBe('a') expect(store.get(b)).toBe('a') From 0b05c61179f51213840c07c606423729c7ea52d7 Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 16 Jul 2024 23:30:44 +0900 Subject: [PATCH 14/31] for old ts --- tests/vanilla/store.test.tsx | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 7a46dc7979..2d8280bb81 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -567,7 +567,11 @@ describe('unstable_derive for scoping atoms', () => { (getAtomState, getAtomContext) => { const scopedAtomStateMap = new WeakMap() return [ - ((atom, context, createDefault) => { + (( + atom: Atom, + context: unknown, + createDefault: () => NonNullable>, + ) => { if (scopedAtoms.has(atom)) { let atomState = scopedAtomStateMap.get(atom) if (!atomState) { @@ -610,7 +614,11 @@ describe('unstable_derive for scoping atoms', () => { (getAtomState, getAtomContext) => { const scopedAtomStateMap = new WeakMap() return [ - ((atom, context, createDefault) => { + (( + atom: Atom, + context: unknown, + createDefault: () => NonNullable>, + ) => { if (scopedAtoms.has(atom)) { let atomState = scopedAtomStateMap.get(atom) if (!atomState) { @@ -652,7 +660,11 @@ describe('unstable_derive for scoping atoms', () => { (getAtomState, getAtomContext) => { const scopedAtomStateMap = new WeakMap() return [ - ((atom, context, createDefault) => { + (( + atom: Atom, + context: unknown, + createDefault: () => NonNullable>, + ) => { if ( (context as { scoped: unknown } | undefined)?.scoped || scopedAtoms.has(atom) From b6610e4282621d22de195bc2f3c1debbdd7ad177 Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 16 Jul 2024 23:38:26 +0900 Subject: [PATCH 15/31] fix --- src/vanilla/store.ts | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 26ecc78dee..d6ceb054f6 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -494,15 +494,20 @@ const buildStore = ( ) => { const getDependents = (a: AnyAtom, context: AtomContext) => { const aState = getAtomState(a, context, createDefaultAtomState) - const nextContext = getAtomContext(a, context) - const dependents = new Set(aState.m?.t) + const dependents = new Map() + for (const dependent of aState.m?.t || []) { + dependents.set(dependent, getAtomContext(dependent, context)) + } for (const atomWithPendingContinuablePromise of aState.p) { - dependents.add(atomWithPendingContinuablePromise) + dependents.set( + atomWithPendingContinuablePromise, + getAtomContext(atomWithPendingContinuablePromise, context), + ) } getPendingDependents(pending, a)?.forEach((dependent) => { - dependents.add(dependent) + dependents.set(dependent, getAtomContext(dependent, context)) }) - return [dependents, nextContext] as const + return dependents } // This is a topological sort via depth-first search, slightly modified from @@ -519,10 +524,10 @@ const buildStore = ( return } markedAtoms.add(n) - const [dependents, nextContext] = getDependents(n, context) - for (const m of dependents) { + const dependents = getDependents(n, context) + for (const [m, c] of dependents) { if (n !== m) { - visit(m, nextContext) + visit(m, c) } } // The algorithm calls for pushing onto the front of the list. For From ed6f32d08891fe6bf5a4a69b5164f9568adc4c67 Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 16 Jul 2024 23:51:42 +0900 Subject: [PATCH 16/31] fix 2 --- src/vanilla/store.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index d6ceb054f6..7f3ff8384e 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -419,11 +419,11 @@ const buildStore = ( // a !== atom const aState = readAtomState(pending, a, nextContext, force) if (isSync) { - addDependency(pending, atom, nextContext, a, aState) + addDependency(pending, atom, context, a, aState) } else { const pending = createPending() - addDependency(pending, atom, nextContext, a, aState) - mountDependencies(pending, atom, nextContext, atomState) + addDependency(pending, atom, context, a, aState) + mountDependencies(pending, atom, context, atomState) flushPending(pending) } return returnAtomValue(aState) From 442f1189c237d597b13a8774e0701f364d7f4290 Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 17 Jul 2024 00:14:20 +0900 Subject: [PATCH 17/31] simplify a bit with dev breaking changes --- src/vanilla/store.ts | 74 +++++++++++++-------------------- tests/vanilla/store.test.tsx | 49 +++++++--------------- tests/vanilla/storedev.test.tsx | 28 ++++++------- 3 files changed, 56 insertions(+), 95 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 7f3ff8384e..c693be6faf 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -149,12 +149,6 @@ type AtomState = { e?: AnyError } -const createDefaultAtomState = (): AtomState => ({ - d: new Map(), - p: new Set(), - n: 0, -}) - const isAtomStateInitialized = (atomState: AtomState) => 'v' in atomState || 'e' in atomState @@ -249,14 +243,10 @@ const flushPending = (pending: Pending) => { type AtomContext = unknown -type GetAtomState = { - (atom: Atom): AtomState | undefined - ( - atom: Atom, - context: AtomContext, - createDefault: () => AtomState, - ): AtomState -} +type GetAtomState = ( + atom: Atom, + context: AtomContext, +) => AtomState type GetAtomContext = (atom: AnyAtom, prevContext?: AtomContext) => AtomContext @@ -267,10 +257,10 @@ type StoreArgs = readonly [ ] // for debugging purpose only -type DevStoreRev4 = { - dev4_get_internal_weak_map: () => { get: GetAtomState } - dev4_get_mounted_atoms: () => Set - dev4_restore_atoms: (values: Iterable) => void +type DevStoreRev5 = { + dev5_get_atom_state: GetAtomState + dev5_get_mounted_atoms: () => Set + dev5_restore_atoms: (values: Iterable) => void } type PrdStore = { @@ -283,9 +273,9 @@ type PrdStore = { unstable_derive: (fn: (...args: StoreArgs) => StoreArgs) => Store } -type Store = PrdStore | (PrdStore & DevStoreRev4) +type Store = PrdStore | (PrdStore & DevStoreRev5) -export type INTERNAL_DevStoreRev4 = DevStoreRev4 +export type INTERNAL_DevStoreRev5 = DevStoreRev5 export type INTERNAL_PrdStore = PrdStore const buildStore = ( @@ -324,7 +314,7 @@ const buildStore = ( ) if (continuablePromise.status === PENDING) { for (const a of atomState.d.keys()) { - const aState = getAtomState(a, context, createDefaultAtomState) + const aState = getAtomState(a, context) addPendingContinuablePromiseToDependency( atom, continuablePromise, @@ -359,7 +349,7 @@ const buildStore = ( if (import.meta.env?.MODE !== 'production' && a === atom) { throw new Error('[Bug] atom cannot depend on itself') } - const atomState = getAtomState(atom, context, createDefaultAtomState) + const atomState = getAtomState(atom, context) atomState.d.set(a, aState.n) const continuablePromise = getPendingContinuablePromise(atomState) if (continuablePromise) { @@ -377,7 +367,7 @@ const buildStore = ( context: AtomContext, force?: (a: AnyAtom) => boolean, ): AtomState => { - const atomState = getAtomState(atom, context, createDefaultAtomState) + const atomState = getAtomState(atom, context) // See if we can skip recomputing this atom. if (!force?.(atom) && isAtomStateInitialized(atomState)) { // If the atom is mounted, we can use the cache. @@ -405,7 +395,7 @@ const buildStore = ( const getter: Getter = (a: Atom) => { const nextContext = getAtomContext(a, context) if (a === (atom as AnyAtom)) { - const aState = getAtomState(a, nextContext, createDefaultAtomState) + const aState = getAtomState(a, nextContext) if (!isAtomStateInitialized(aState)) { if (hasInitialValue(a)) { setAtomStateValueOrPromise(a, nextContext, aState, a.init) @@ -493,7 +483,7 @@ const buildStore = ( context: AtomContext, ) => { const getDependents = (a: AnyAtom, context: AtomContext) => { - const aState = getAtomState(a, context, createDefaultAtomState) + const aState = getAtomState(a, context) const dependents = new Map() for (const dependent of aState.m?.t || []) { dependents.set(dependent, getAtomContext(dependent, context)) @@ -544,7 +534,7 @@ const buildStore = ( const isMarked = (a: AnyAtom) => markedAtoms.has(a) for (let i = topsortedAtoms.length - 1; i >= 0; --i) { const [a, nextContext] = topsortedAtoms[i]! - const aState = getAtomState(a, nextContext, createDefaultAtomState) + const aState = getAtomState(a, nextContext) const prevEpochNumber = aState.n let hasChangedDeps = false for (const dep of aState.d.keys()) { @@ -584,7 +574,7 @@ const buildStore = ( // NOTE technically possible but restricted as it may cause bugs throw new Error('atom not writable') } - const aState = getAtomState(a, nextContext, createDefaultAtomState) + const aState = getAtomState(a, nextContext) const hasPrevValue = 'v' in aState const prevValue = aState.v const v = args[0] as V @@ -643,7 +633,7 @@ const buildStore = ( atom: AnyAtom, context: AtomContext, ): Mounted => { - const atomState = getAtomState(atom, context, createDefaultAtomState) + const atomState = getAtomState(atom, context) if (!atomState.m) { // recompute atom state readAtomState(pending, atom, context) @@ -682,13 +672,12 @@ const buildStore = ( atom: AnyAtom, context: AtomContext, ): Mounted | undefined => { - const atomState = getAtomState(atom, context, createDefaultAtomState) + const atomState = getAtomState(atom, context) if ( atomState.m && !atomState.m.l.size && !Array.from(atomState.m.t).some( - (a) => - getAtomState(a, getAtomContext(a, context), createDefaultAtomState).m, + (a) => getAtomState(a, getAtomContext(a, context)).m, ) ) { // unmount self @@ -740,20 +729,16 @@ const buildStore = ( unstable_derive, } if (import.meta.env?.MODE !== 'production') { - const devStore: DevStoreRev4 = { + const devStore: DevStoreRev5 = { // store dev methods (these are tentative and subject to change without notice) - dev4_get_internal_weak_map: () => ({ get: getAtomState }), - dev4_get_mounted_atoms: () => debugMountedAtoms, - dev4_restore_atoms: (values) => { + dev5_get_atom_state: getAtomState, + dev5_get_mounted_atoms: () => debugMountedAtoms, + dev5_restore_atoms: (values) => { const pending = createPending() for (const [atom, value] of values) { if (hasInitialValue(atom)) { const context = getAtomContext(atom) - const atomState = getAtomState( - atom, - context, - createDefaultAtomState, - ) + const atomState = getAtomState(atom, context) const hasPrevValue = 'v' in atomState const prevValue = atomState.v setAtomStateValueOrPromise(atom, context, atomState, value) @@ -774,17 +759,14 @@ const buildStore = ( export const createStore = (): Store => { const atomStateMap = new WeakMap() - const getAtomState = ((atom, _context, createDefault) => { + const getAtomState: GetAtomState = (atom) => { let atomState = atomStateMap.get(atom) if (!atomState) { - if (!createDefault) { - return undefined - } - atomState = createDefault() + atomState = { d: new Map(), p: new Set(), n: 0 } atomStateMap.set(atom, atomState) } return atomState - }) as GetAtomState + } const getAtomContext: GetAtomContext = (_atom, context) => context return buildStore(getAtomState, getAtomContext) } diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 2d8280bb81..b778b660f5 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -43,7 +43,7 @@ it('should unmount with store.get', async () => { store.get(countAtom) unsub() const result = Array.from( - 'dev4_get_mounted_atoms' in store ? store.dev4_get_mounted_atoms() : [], + 'dev5_get_mounted_atoms' in store ? store.dev5_get_mounted_atoms() : [], ) expect(result).toEqual([]) }) @@ -57,7 +57,7 @@ it('should unmount dependencies with store.get', async () => { store.get(derivedAtom) unsub() const result = Array.from( - 'dev4_restore_atoms' in store ? store.dev4_get_mounted_atoms() : [], + 'dev5_restore_atoms' in store ? store.dev5_get_mounted_atoms() : [], ) expect(result).toEqual([]) }) @@ -567,24 +567,17 @@ describe('unstable_derive for scoping atoms', () => { (getAtomState, getAtomContext) => { const scopedAtomStateMap = new WeakMap() return [ - (( - atom: Atom, - context: unknown, - createDefault: () => NonNullable>, - ) => { + (atom: Atom, context: unknown) => { if (scopedAtoms.has(atom)) { let atomState = scopedAtomStateMap.get(atom) if (!atomState) { - if (!createDefault) { - return undefined - } - atomState = createDefault() + atomState = { d: new Map(), p: new Set(), n: 0 } scopedAtomStateMap.set(atom, atomState) } return atomState } - return getAtomState(atom, context, createDefault) - }) as typeof getAtomState, + return getAtomState(atom, context) + }, getAtomContext, ] }, @@ -614,24 +607,17 @@ describe('unstable_derive for scoping atoms', () => { (getAtomState, getAtomContext) => { const scopedAtomStateMap = new WeakMap() return [ - (( - atom: Atom, - context: unknown, - createDefault: () => NonNullable>, - ) => { + (atom: Atom, context: unknown) => { if (scopedAtoms.has(atom)) { let atomState = scopedAtomStateMap.get(atom) if (!atomState) { - if (!createDefault) { - return undefined - } - atomState = createDefault() + atomState = { d: new Map(), p: new Set(), n: 0 } scopedAtomStateMap.set(atom, atomState) } return atomState } - return getAtomState(atom, context, createDefault) - }) as typeof getAtomState, + return getAtomState(atom, context) + }, getAtomContext, ] }, @@ -660,27 +646,20 @@ describe('unstable_derive for scoping atoms', () => { (getAtomState, getAtomContext) => { const scopedAtomStateMap = new WeakMap() return [ - (( - atom: Atom, - context: unknown, - createDefault: () => NonNullable>, - ) => { + (atom: Atom, context: unknown) => { if ( (context as { scoped: unknown } | undefined)?.scoped || scopedAtoms.has(atom) ) { let atomState = scopedAtomStateMap.get(atom) if (!atomState) { - if (!createDefault) { - return undefined - } - atomState = createDefault() + atomState = { d: new Map(), p: new Set(), n: 0 } scopedAtomStateMap.set(atom, atomState) } return atomState } - return getAtomState(atom, context, createDefault) - }) as typeof getAtomState, + return getAtomState(atom, context) + }, (atom, context) => { context = getAtomContext(atom, context) if (scopedAtoms.has(atom)) { diff --git a/tests/vanilla/storedev.test.tsx b/tests/vanilla/storedev.test.tsx index 1771fafc28..1554769eb1 100644 --- a/tests/vanilla/storedev.test.tsx +++ b/tests/vanilla/storedev.test.tsx @@ -1,39 +1,39 @@ import { describe, expect, it, vi } from 'vitest' import { atom, createStore } from 'jotai/vanilla' import type { - INTERNAL_DevStoreRev4, + INTERNAL_DevStoreRev5, INTERNAL_PrdStore, } from 'jotai/vanilla/store' describe('[DEV-ONLY] dev-only methods rev4', () => { it('should get atom value', () => { const store = createStore() as any - if (!('dev4_get_internal_weak_map' in store)) { + if (!('dev5_get_atom_state' in store)) { throw new Error('dev methods are not available') } const countAtom = atom(0) countAtom.debugLabel = 'countAtom' store.set(countAtom, 1) - const weakMap = store.dev4_get_internal_weak_map() - expect(weakMap.get(countAtom)?.v).toEqual(1) + const getAtomState = store.dev5_get_atom_state + expect(getAtomState(countAtom).v).toEqual(1) }) it('should restore atoms and its dependencies correctly', () => { const store = createStore() as any - if (!('dev4_restore_atoms' in store)) { + if (!('dev5_restore_atoms' in store)) { throw new Error('dev methods are not available') } const countAtom = atom(0) const derivedAtom = atom((get) => get(countAtom) * 2) store.set(countAtom, 1) - store.dev4_restore_atoms([[countAtom, 2]]) + store.dev5_restore_atoms([[countAtom, 2]]) expect(store.get(countAtom)).toBe(2) expect(store.get?.(derivedAtom)).toBe(4) }) it('should restore atoms and call store listeners correctly', () => { const store = createStore() as any - if (!('dev4_restore_atoms' in store)) { + if (!('dev5_restore_atoms' in store)) { throw new Error('dev methods are not available') } const countAtom = atom(0) @@ -43,7 +43,7 @@ describe('[DEV-ONLY] dev-only methods rev4', () => { store.set(countAtom, 2) const unsubCount = store.sub(countAtom, countCb) const unsubDerived = store.sub(derivedAtom, derivedCb) - store.dev4_restore_atoms([ + store.dev5_restore_atoms([ [countAtom, 1], [derivedAtom, 2], ]) @@ -55,8 +55,8 @@ describe('[DEV-ONLY] dev-only methods rev4', () => { }) it('should return all the mounted atoms correctly', () => { - const store = createStore() as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore - if (!('dev4_get_mounted_atoms' in store)) { + const store = createStore() as INTERNAL_DevStoreRev5 & INTERNAL_PrdStore + if (!('dev5_get_mounted_atoms' in store)) { throw new Error('dev methods are not available') } const countAtom = atom(0) @@ -64,7 +64,7 @@ describe('[DEV-ONLY] dev-only methods rev4', () => { const derivedAtom = atom((get) => get(countAtom) * 2) const unsub = store.sub(derivedAtom, vi.fn()) store.set(countAtom, 1) - const result = store.dev4_get_mounted_atoms() + const result = store.dev5_get_mounted_atoms() expect( Array.from(result).sort( (a, b) => Object.keys(a).length - Object.keys(b).length, @@ -83,8 +83,8 @@ describe('[DEV-ONLY] dev-only methods rev4', () => { }) it("should return all the mounted atoms correctly after they're unsubscribed", () => { - const store = createStore() as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore - if (!('dev4_get_mounted_atoms' in store)) { + const store = createStore() as INTERNAL_DevStoreRev5 & INTERNAL_PrdStore + if (!('dev5_get_mounted_atoms' in store)) { throw new Error('dev methods are not available') } const countAtom = atom(0) @@ -93,7 +93,7 @@ describe('[DEV-ONLY] dev-only methods rev4', () => { const unsub = store.sub(derivedAtom, vi.fn()) store.set(countAtom, 1) unsub() - const result = store.dev4_get_mounted_atoms() + const result = store.dev5_get_mounted_atoms() expect(Array.from(result)).toStrictEqual([]) }) }) From b6f8914433e0bda0b2aff49104402b948d5650e4 Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 17 Jul 2024 08:21:38 +0900 Subject: [PATCH 18/31] revert to bd9c6c0, except for new test --- src/vanilla/store.ts | 211 +++++++++++++------------------- tests/vanilla/store.test.tsx | 4 +- tests/vanilla/storedev.test.tsx | 28 ++--- 3 files changed, 102 insertions(+), 141 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index c693be6faf..932aac0067 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -241,28 +241,24 @@ const flushPending = (pending: Pending) => { } } -type AtomContext = unknown - -type GetAtomState = ( - atom: Atom, - context: AtomContext, -) => AtomState +// for debugging purpose only +type DevStoreRev4 = { + dev4_get_internal_weak_map: () => AtomStateMap + dev4_get_mounted_atoms: () => Set + dev4_restore_atoms: (values: Iterable) => void +} -type GetAtomContext = (atom: AnyAtom, prevContext?: AtomContext) => AtomContext +type AtomStateMap = { + get(key: Atom): AtomState | undefined + set(key: Atom, value: AtomState): void +} // internal & unstable type type StoreArgs = readonly [ - getAtomState: GetAtomState, - getAtomContext: GetAtomContext, + atomStateMap: AtomStateMap, + // possible other arguments in the future ] -// for debugging purpose only -type DevStoreRev5 = { - dev5_get_atom_state: GetAtomState - dev5_get_mounted_atoms: () => Set - dev5_restore_atoms: (values: Iterable) => void -} - type PrdStore = { get: (atom: Atom) => Value set: ( @@ -272,16 +268,12 @@ type PrdStore = { sub: (atom: AnyAtom, listener: () => void) => () => void unstable_derive: (fn: (...args: StoreArgs) => StoreArgs) => Store } +type Store = PrdStore | (PrdStore & DevStoreRev4) -type Store = PrdStore | (PrdStore & DevStoreRev5) - -export type INTERNAL_DevStoreRev5 = DevStoreRev5 +export type INTERNAL_DevStoreRev4 = DevStoreRev4 export type INTERNAL_PrdStore = PrdStore -const buildStore = ( - getAtomState: StoreArgs[0], - getAtomContext: StoreArgs[1], -): Store => { +const buildStore = (atomStateMap: StoreArgs[0]): Store => { // for debugging purpose only let debugMountedAtoms: Set @@ -289,9 +281,17 @@ const buildStore = ( debugMountedAtoms = new Set() } + const getAtomState = (atom: Atom) => { + let atomState = atomStateMap.get(atom) + if (!atomState) { + atomState = { d: new Map(), p: new Set(), n: 0 } + atomStateMap.set(atom, atomState) + } + return atomState + } + const setAtomStateValueOrPromise = ( atom: AnyAtom, - context: AtomContext, atomState: AtomState, valueOrPromise: unknown, abortPromise = () => {}, @@ -314,7 +314,7 @@ const buildStore = ( ) if (continuablePromise.status === PENDING) { for (const a of atomState.d.keys()) { - const aState = getAtomState(a, context) + const aState = getAtomState(a) addPendingContinuablePromiseToDependency( atom, continuablePromise, @@ -342,14 +342,13 @@ const buildStore = ( const addDependency = ( pending: Pending | undefined, atom: Atom, - context: AtomContext, a: AnyAtom, aState: AtomState, ) => { if (import.meta.env?.MODE !== 'production' && a === atom) { throw new Error('[Bug] atom cannot depend on itself') } - const atomState = getAtomState(atom, context) + const atomState = getAtomState(atom) atomState.d.set(a, aState.n) const continuablePromise = getPendingContinuablePromise(atomState) if (continuablePromise) { @@ -364,11 +363,10 @@ const buildStore = ( const readAtomState = ( pending: Pending | undefined, atom: Atom, - context: AtomContext, force?: (a: AnyAtom) => boolean, ): AtomState => { - const atomState = getAtomState(atom, context) // See if we can skip recomputing this atom. + const atomState = getAtomState(atom) if (!force?.(atom) && isAtomStateInitialized(atomState)) { // If the atom is mounted, we can use the cache. // because it should have been updated by dependencies. @@ -382,8 +380,7 @@ const buildStore = ( ([a, n]) => // Recursively, read the atom state of the dependency, and // check if the atom epoch number is unchanged - readAtomState(pending, a, getAtomContext(a, context), force).n === - n, + readAtomState(pending, a, force).n === n, ) ) { return atomState @@ -393,12 +390,11 @@ const buildStore = ( atomState.d.clear() let isSync = true const getter: Getter = (a: Atom) => { - const nextContext = getAtomContext(a, context) if (a === (atom as AnyAtom)) { - const aState = getAtomState(a, nextContext) + const aState = getAtomState(a) if (!isAtomStateInitialized(aState)) { if (hasInitialValue(a)) { - setAtomStateValueOrPromise(a, nextContext, aState, a.init) + setAtomStateValueOrPromise(a, aState, a.init) } else { // NOTE invalid derived atoms can reach here throw new Error('no atom init') @@ -407,13 +403,13 @@ const buildStore = ( return returnAtomValue(aState) } // a !== atom - const aState = readAtomState(pending, a, nextContext, force) + const aState = readAtomState(pending, a, force) if (isSync) { - addDependency(pending, atom, context, a, aState) + addDependency(pending, atom, a, aState) } else { const pending = createPending() - addDependency(pending, atom, context, a, aState) - mountDependencies(pending, atom, context, atomState) + addDependency(pending, atom, a, aState) + mountDependencies(pending, atom, atomState) flushPending(pending) } return returnAtomValue(aState) @@ -451,14 +447,13 @@ const buildStore = ( const valueOrPromise = atom.read(getter, options as never) setAtomStateValueOrPromise( atom, - context, atomState, valueOrPromise, () => controller?.abort(), () => { if (atomState.m) { const pending = createPending() - mountDependencies(pending, atom, context, atomState) + mountDependencies(pending, atom, atomState) flushPending(pending) } }, @@ -475,27 +470,17 @@ const buildStore = ( } const readAtom = (atom: Atom): Value => - returnAtomValue(readAtomState(undefined, atom, getAtomContext(atom))) + returnAtomValue(readAtomState(undefined, atom)) - const recomputeDependents = ( - pending: Pending, - atom: AnyAtom, - context: AtomContext, - ) => { - const getDependents = (a: AnyAtom, context: AtomContext) => { - const aState = getAtomState(a, context) - const dependents = new Map() - for (const dependent of aState.m?.t || []) { - dependents.set(dependent, getAtomContext(dependent, context)) - } + const recomputeDependents = (pending: Pending, atom: AnyAtom) => { + const getDependents = (a: AnyAtom): Set => { + const aState = getAtomState(a) + const dependents = new Set(aState.m?.t) for (const atomWithPendingContinuablePromise of aState.p) { - dependents.set( - atomWithPendingContinuablePromise, - getAtomContext(atomWithPendingContinuablePromise, context), - ) + dependents.add(atomWithPendingContinuablePromise) } getPendingDependents(pending, a)?.forEach((dependent) => { - dependents.set(dependent, getAtomContext(dependent, context)) + dependents.add(dependent) }) return dependents } @@ -506,35 +491,33 @@ const buildStore = ( // Step 1: traverse the dependency graph to build the topsorted atom list // We don't bother to check for cycles, which simplifies the algorithm. - const topsortedAtoms: (readonly [atom: AnyAtom, context: AtomContext])[] = - [] + const topsortedAtoms: AnyAtom[] = [] const markedAtoms = new Set() - const visit = (n: AnyAtom, context: AtomContext) => { + const visit = (n: AnyAtom) => { if (markedAtoms.has(n)) { return } markedAtoms.add(n) - const dependents = getDependents(n, context) - for (const [m, c] of dependents) { + for (const m of getDependents(n)) { if (n !== m) { - visit(m, c) + visit(m) } } // The algorithm calls for pushing onto the front of the list. For // performance, we will simply push onto the end, and then will iterate in // reverse order later. - topsortedAtoms.push([n, context]) + topsortedAtoms.push(n) } // Visit the root atom. This is the only atom in the dependency graph // without incoming edges, which is one reason we can simplify the algorithm - visit(atom, context) + visit(atom) // Step 2: use the topsorted atom list to recompute all affected atoms // Track what's changed, so that we can short circuit when possible const changedAtoms = new Set([atom]) const isMarked = (a: AnyAtom) => markedAtoms.has(a) for (let i = topsortedAtoms.length - 1; i >= 0; --i) { - const [a, nextContext] = topsortedAtoms[i]! - const aState = getAtomState(a, nextContext) + const a = topsortedAtoms[i]! + const aState = getAtomState(a) const prevEpochNumber = aState.n let hasChangedDeps = false for (const dep of aState.d.keys()) { @@ -544,8 +527,8 @@ const buildStore = ( } } if (hasChangedDeps) { - readAtomState(pending, a, nextContext, isMarked) - mountDependencies(pending, a, nextContext, aState) + readAtomState(pending, a, isMarked) + mountDependencies(pending, a, aState) if (prevEpochNumber !== aState.n) { addPendingAtom(pending, a, aState) changedAtoms.add(a) @@ -558,34 +541,32 @@ const buildStore = ( const writeAtomState = ( pending: Pending, atom: WritableAtom, - context: AtomContext, ...args: Args ): Result => { const getter: Getter = (a: Atom) => - returnAtomValue(readAtomState(pending, a, getAtomContext(a, context))) + returnAtomValue(readAtomState(pending, a)) const setter: Setter = ( a: WritableAtom, ...args: As ) => { - const nextContext = getAtomContext(a, context) let r: R | undefined if (a === (atom as AnyAtom)) { if (!hasInitialValue(a)) { // NOTE technically possible but restricted as it may cause bugs throw new Error('atom not writable') } - const aState = getAtomState(a, nextContext) + const aState = getAtomState(a) const hasPrevValue = 'v' in aState const prevValue = aState.v const v = args[0] as V - setAtomStateValueOrPromise(a, nextContext, aState, v) - mountDependencies(pending, a, nextContext, aState) + setAtomStateValueOrPromise(a, aState, v) + mountDependencies(pending, a, aState) if (!hasPrevValue || !Object.is(prevValue, aState.v)) { addPendingAtom(pending, a, aState) - recomputeDependents(pending, a, nextContext) + recomputeDependents(pending, a) } } else { - r = writeAtomState(pending, a, nextContext, ...args) as R + r = writeAtomState(pending, a, ...args) as R } flushPending(pending) return r as R @@ -599,7 +580,7 @@ const buildStore = ( ...args: Args ): Result => { const pending = createPending() - const result = writeAtomState(pending, atom, getAtomContext(atom), ...args) + const result = writeAtomState(pending, atom, ...args) flushPending(pending) return result } @@ -607,20 +588,19 @@ const buildStore = ( const mountDependencies = ( pending: Pending, atom: AnyAtom, - context: AtomContext, atomState: AtomState, ) => { if (atomState.m && !getPendingContinuablePromise(atomState)) { for (const a of atomState.d.keys()) { if (!atomState.m.d.has(a)) { - const aMounted = mountAtom(pending, a, context) + const aMounted = mountAtom(pending, a) aMounted.t.add(atom) atomState.m.d.add(a) } } for (const a of atomState.m.d || []) { if (!atomState.d.has(a)) { - const aMounted = unmountAtom(pending, a, context) + const aMounted = unmountAtom(pending, a) aMounted?.t.delete(atom) atomState.m.d.delete(a) } @@ -628,18 +608,14 @@ const buildStore = ( } } - const mountAtom = ( - pending: Pending, - atom: AnyAtom, - context: AtomContext, - ): Mounted => { - const atomState = getAtomState(atom, context) + const mountAtom = (pending: Pending, atom: AnyAtom): Mounted => { + const atomState = getAtomState(atom) if (!atomState.m) { // recompute atom state - readAtomState(pending, atom, context) + readAtomState(pending, atom) // mount dependencies first for (const a of atomState.d.keys()) { - const aMounted = mountAtom(pending, a, getAtomContext(a, context)) + const aMounted = mountAtom(pending, a) aMounted.t.add(atom) } // mount self @@ -656,7 +632,7 @@ const buildStore = ( const { onMount } = atom addPendingFunction(pending, () => { const onUnmount = onMount((...args) => - writeAtomState(pending, atom, context, ...args), + writeAtomState(pending, atom, ...args), ) if (onUnmount) { mounted.u = onUnmount @@ -670,15 +646,12 @@ const buildStore = ( const unmountAtom = ( pending: Pending, atom: AnyAtom, - context: AtomContext, ): Mounted | undefined => { - const atomState = getAtomState(atom, context) + const atomState = getAtomState(atom) if ( atomState.m && !atomState.m.l.size && - !Array.from(atomState.m.t).some( - (a) => getAtomState(a, getAtomContext(a, context)).m, - ) + !Array.from(atomState.m.t).some((a) => getAtomState(a).m) ) { // unmount self const onUnmount = atomState.m.u @@ -691,7 +664,7 @@ const buildStore = ( } // unmount dependencies for (const a of atomState.d.keys()) { - const aMounted = unmountAtom(pending, a, getAtomContext(a, context)) + const aMounted = unmountAtom(pending, a) aMounted?.t.delete(atom) } // abort pending promise @@ -707,20 +680,20 @@ const buildStore = ( const subscribeAtom = (atom: AnyAtom, listener: () => void) => { const pending = createPending() - const mounted = mountAtom(pending, atom, undefined) + const mounted = mountAtom(pending, atom) flushPending(pending) const listeners = mounted.l listeners.add(listener) return () => { listeners.delete(listener) const pending = createPending() - unmountAtom(pending, atom, undefined) + unmountAtom(pending, atom) flushPending(pending) } } const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) => - buildStore(...fn(getAtomState, getAtomContext)) + buildStore(...fn(atomStateMap)) const store: Store = { get: readAtom, @@ -729,23 +702,22 @@ const buildStore = ( unstable_derive, } if (import.meta.env?.MODE !== 'production') { - const devStore: DevStoreRev5 = { + const devStore: DevStoreRev4 = { // store dev methods (these are tentative and subject to change without notice) - dev5_get_atom_state: getAtomState, - dev5_get_mounted_atoms: () => debugMountedAtoms, - dev5_restore_atoms: (values) => { + dev4_get_internal_weak_map: () => atomStateMap, + dev4_get_mounted_atoms: () => debugMountedAtoms, + dev4_restore_atoms: (values) => { const pending = createPending() for (const [atom, value] of values) { if (hasInitialValue(atom)) { - const context = getAtomContext(atom) - const atomState = getAtomState(atom, context) - const hasPrevValue = 'v' in atomState - const prevValue = atomState.v - setAtomStateValueOrPromise(atom, context, atomState, value) - mountDependencies(pending, atom, context, atomState) - if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { - addPendingAtom(pending, atom, atomState) - recomputeDependents(pending, atom, context) + const aState = getAtomState(atom) + const hasPrevValue = 'v' in aState + const prevValue = aState.v + setAtomStateValueOrPromise(atom, aState, value) + mountDependencies(pending, atom, aState) + if (!hasPrevValue || !Object.is(prevValue, aState.v)) { + addPendingAtom(pending, atom, aState) + recomputeDependents(pending, atom) } } } @@ -757,19 +729,8 @@ const buildStore = ( return store } -export const createStore = (): Store => { - const atomStateMap = new WeakMap() - const getAtomState: GetAtomState = (atom) => { - let atomState = atomStateMap.get(atom) - if (!atomState) { - atomState = { d: new Map(), p: new Set(), n: 0 } - atomStateMap.set(atom, atomState) - } - return atomState - } - const getAtomContext: GetAtomContext = (_atom, context) => context - return buildStore(getAtomState, getAtomContext) -} +export const createStore = (): Store => + buildStore(new WeakMap() as AtomStateMap) let defaultStore: Store | undefined diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index b778b660f5..07a90ae713 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -43,7 +43,7 @@ it('should unmount with store.get', async () => { store.get(countAtom) unsub() const result = Array.from( - 'dev5_get_mounted_atoms' in store ? store.dev5_get_mounted_atoms() : [], + 'dev4_get_mounted_atoms' in store ? store.dev4_get_mounted_atoms() : [], ) expect(result).toEqual([]) }) @@ -57,7 +57,7 @@ it('should unmount dependencies with store.get', async () => { store.get(derivedAtom) unsub() const result = Array.from( - 'dev5_restore_atoms' in store ? store.dev5_get_mounted_atoms() : [], + 'dev4_restore_atoms' in store ? store.dev4_get_mounted_atoms() : [], ) expect(result).toEqual([]) }) diff --git a/tests/vanilla/storedev.test.tsx b/tests/vanilla/storedev.test.tsx index 1554769eb1..1771fafc28 100644 --- a/tests/vanilla/storedev.test.tsx +++ b/tests/vanilla/storedev.test.tsx @@ -1,39 +1,39 @@ import { describe, expect, it, vi } from 'vitest' import { atom, createStore } from 'jotai/vanilla' import type { - INTERNAL_DevStoreRev5, + INTERNAL_DevStoreRev4, INTERNAL_PrdStore, } from 'jotai/vanilla/store' describe('[DEV-ONLY] dev-only methods rev4', () => { it('should get atom value', () => { const store = createStore() as any - if (!('dev5_get_atom_state' in store)) { + if (!('dev4_get_internal_weak_map' in store)) { throw new Error('dev methods are not available') } const countAtom = atom(0) countAtom.debugLabel = 'countAtom' store.set(countAtom, 1) - const getAtomState = store.dev5_get_atom_state - expect(getAtomState(countAtom).v).toEqual(1) + const weakMap = store.dev4_get_internal_weak_map() + expect(weakMap.get(countAtom)?.v).toEqual(1) }) it('should restore atoms and its dependencies correctly', () => { const store = createStore() as any - if (!('dev5_restore_atoms' in store)) { + if (!('dev4_restore_atoms' in store)) { throw new Error('dev methods are not available') } const countAtom = atom(0) const derivedAtom = atom((get) => get(countAtom) * 2) store.set(countAtom, 1) - store.dev5_restore_atoms([[countAtom, 2]]) + store.dev4_restore_atoms([[countAtom, 2]]) expect(store.get(countAtom)).toBe(2) expect(store.get?.(derivedAtom)).toBe(4) }) it('should restore atoms and call store listeners correctly', () => { const store = createStore() as any - if (!('dev5_restore_atoms' in store)) { + if (!('dev4_restore_atoms' in store)) { throw new Error('dev methods are not available') } const countAtom = atom(0) @@ -43,7 +43,7 @@ describe('[DEV-ONLY] dev-only methods rev4', () => { store.set(countAtom, 2) const unsubCount = store.sub(countAtom, countCb) const unsubDerived = store.sub(derivedAtom, derivedCb) - store.dev5_restore_atoms([ + store.dev4_restore_atoms([ [countAtom, 1], [derivedAtom, 2], ]) @@ -55,8 +55,8 @@ describe('[DEV-ONLY] dev-only methods rev4', () => { }) it('should return all the mounted atoms correctly', () => { - const store = createStore() as INTERNAL_DevStoreRev5 & INTERNAL_PrdStore - if (!('dev5_get_mounted_atoms' in store)) { + const store = createStore() as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore + if (!('dev4_get_mounted_atoms' in store)) { throw new Error('dev methods are not available') } const countAtom = atom(0) @@ -64,7 +64,7 @@ describe('[DEV-ONLY] dev-only methods rev4', () => { const derivedAtom = atom((get) => get(countAtom) * 2) const unsub = store.sub(derivedAtom, vi.fn()) store.set(countAtom, 1) - const result = store.dev5_get_mounted_atoms() + const result = store.dev4_get_mounted_atoms() expect( Array.from(result).sort( (a, b) => Object.keys(a).length - Object.keys(b).length, @@ -83,8 +83,8 @@ describe('[DEV-ONLY] dev-only methods rev4', () => { }) it("should return all the mounted atoms correctly after they're unsubscribed", () => { - const store = createStore() as INTERNAL_DevStoreRev5 & INTERNAL_PrdStore - if (!('dev5_get_mounted_atoms' in store)) { + const store = createStore() as INTERNAL_DevStoreRev4 & INTERNAL_PrdStore + if (!('dev4_get_mounted_atoms' in store)) { throw new Error('dev methods are not available') } const countAtom = atom(0) @@ -93,7 +93,7 @@ describe('[DEV-ONLY] dev-only methods rev4', () => { const unsub = store.sub(derivedAtom, vi.fn()) store.set(countAtom, 1) unsub() - const result = store.dev5_get_mounted_atoms() + const result = store.dev4_get_mounted_atoms() expect(Array.from(result)).toStrictEqual([]) }) }) From 18205d35715bf85541449401624c023792b4bdf4 Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 17 Jul 2024 08:56:39 +0900 Subject: [PATCH 19/31] refactor (skipping the new test for now) --- src/vanilla/store.ts | 167 ++++++++++++++++++----------------- tests/vanilla/store.test.tsx | 4 +- 2 files changed, 90 insertions(+), 81 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index ea79d3cb26..5c9c32ca5f 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -149,6 +149,8 @@ type AtomState = { e?: AnyError } +type GetAtomState = (atom: Atom) => AtomState + const isAtomStateInitialized = (atomState: AtomState) => 'v' in atomState || 'e' in atomState @@ -241,11 +243,77 @@ const flushPending = (pending: Pending) => { } } -// for debugging purpose only -type DevStoreRev4 = { - dev4_get_internal_weak_map: () => AtomStateMap - dev4_get_mounted_atoms: () => Set - dev4_restore_atoms: (values: Iterable) => void +const setAtomStateValueOrPromise = ( + getAtomState: GetAtomState, + atom: AnyAtom, + atomState: AtomState, + valueOrPromise: unknown, + abortPromise = () => {}, + completePromise = () => {}, +) => { + const hasPrevValue = 'v' in atomState + const prevValue = atomState.v + const pendingPromise = getPendingContinuablePromise(atomState) + if (isPromiseLike(valueOrPromise)) { + if (pendingPromise) { + if (pendingPromise !== valueOrPromise) { + pendingPromise[CONTINUE_PROMISE](valueOrPromise, abortPromise) + ++atomState.n + } + } else { + const continuablePromise = createContinuablePromise( + valueOrPromise, + abortPromise, + completePromise, + ) + if (continuablePromise.status === PENDING) { + for (const a of atomState.d.keys()) { + const aState = getAtomState(a) + addPendingContinuablePromiseToDependency( + atom, + continuablePromise, + aState, + ) + } + } + atomState.v = continuablePromise + delete atomState.e + } + } else { + if (pendingPromise) { + pendingPromise[CONTINUE_PROMISE]( + Promise.resolve(valueOrPromise), + abortPromise, + ) + } + atomState.v = valueOrPromise + delete atomState.e + } + if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { + ++atomState.n + } +} + +const addDependency = ( + pending: Pending | undefined, + getAtomState: GetAtomState, + atom: Atom, + a: AnyAtom, + aState: AtomState, +) => { + if (import.meta.env?.MODE !== 'production' && a === atom) { + throw new Error('[Bug] atom cannot depend on itself') + } + const atomState = getAtomState(atom) + atomState.d.set(a, aState.n) + const continuablePromise = getPendingContinuablePromise(atomState) + if (continuablePromise) { + addPendingContinuablePromiseToDependency(atom, continuablePromise, aState) + } + aState.m?.t.add(atom) + if (pending) { + addPendingDependent(pending, a, atom) + } } type AtomStateMap = { @@ -259,6 +327,13 @@ type StoreArgs = readonly [ // possible other arguments in the future ] +// for debugging purpose only +type DevStoreRev4 = { + dev4_get_internal_weak_map: () => AtomStateMap + dev4_get_mounted_atoms: () => Set + dev4_restore_atoms: (values: Iterable) => void +} + type PrdStore = { get: (atom: Atom) => Value set: ( @@ -268,6 +343,7 @@ type PrdStore = { sub: (atom: AnyAtom, listener: () => void) => () => void unstable_derive: (fn: (...args: StoreArgs) => StoreArgs) => Store } + type Store = PrdStore | (PrdStore & DevStoreRev4) export type INTERNAL_DevStoreRev4 = DevStoreRev4 @@ -290,76 +366,6 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { return atomState } - const setAtomStateValueOrPromise = ( - atom: AnyAtom, - atomState: AtomState, - valueOrPromise: unknown, - abortPromise = () => {}, - completePromise = () => {}, - ) => { - const hasPrevValue = 'v' in atomState - const prevValue = atomState.v - const pendingPromise = getPendingContinuablePromise(atomState) - if (isPromiseLike(valueOrPromise)) { - if (pendingPromise) { - if (pendingPromise !== valueOrPromise) { - pendingPromise[CONTINUE_PROMISE](valueOrPromise, abortPromise) - ++atomState.n - } - } else { - const continuablePromise = createContinuablePromise( - valueOrPromise, - abortPromise, - completePromise, - ) - if (continuablePromise.status === PENDING) { - for (const a of atomState.d.keys()) { - const aState = getAtomState(a) - addPendingContinuablePromiseToDependency( - atom, - continuablePromise, - aState, - ) - } - } - atomState.v = continuablePromise - delete atomState.e - } - } else { - if (pendingPromise) { - pendingPromise[CONTINUE_PROMISE]( - Promise.resolve(valueOrPromise), - abortPromise, - ) - } - atomState.v = valueOrPromise - delete atomState.e - } - if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { - ++atomState.n - } - } - const addDependency = ( - pending: Pending | undefined, - atom: Atom, - a: AnyAtom, - aState: AtomState, - ) => { - if (import.meta.env?.MODE !== 'production' && a === atom) { - throw new Error('[Bug] atom cannot depend on itself') - } - const atomState = getAtomState(atom) - atomState.d.set(a, aState.n) - const continuablePromise = getPendingContinuablePromise(atomState) - if (continuablePromise) { - addPendingContinuablePromiseToDependency(atom, continuablePromise, aState) - } - aState.m?.t.add(atom) - if (pending) { - addPendingDependent(pending, a, atom) - } - } - const readAtomState = ( pending: Pending | undefined, atom: Atom, @@ -394,7 +400,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const aState = getAtomState(a) if (!isAtomStateInitialized(aState)) { if (hasInitialValue(a)) { - setAtomStateValueOrPromise(a, aState, a.init) + setAtomStateValueOrPromise(getAtomState, a, aState, a.init) } else { // NOTE invalid derived atoms can reach here throw new Error('no atom init') @@ -405,10 +411,10 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { // a !== atom const aState = readAtomState(pending, a, force) if (isSync) { - addDependency(pending, atom, a, aState) + addDependency(pending, getAtomState, atom, a, aState) } else { const pending = createPending() - addDependency(pending, atom, a, aState) + addDependency(pending, getAtomState, atom, a, aState) mountDependencies(pending, atom, atomState) flushPending(pending) } @@ -446,6 +452,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { try { const valueOrPromise = atom.read(getter, options as never) setAtomStateValueOrPromise( + getAtomState, atom, atomState, valueOrPromise, @@ -559,7 +566,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const hasPrevValue = 'v' in aState const prevValue = aState.v const v = args[0] as V - setAtomStateValueOrPromise(a, aState, v) + setAtomStateValueOrPromise(getAtomState, a, aState, v) mountDependencies(pending, a, aState) if (!hasPrevValue || !Object.is(prevValue, aState.v)) { addPendingAtom(pending, a, aState) @@ -713,7 +720,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const atomState = getAtomState(atom) const hasPrevValue = 'v' in atomState const prevValue = atomState.v - setAtomStateValueOrPromise(atom, atomState, value) + setAtomStateValueOrPromise(getAtomState, atom, atomState, value) mountDependencies(pending, atom, atomState) if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { addPendingAtom(pending, atom, atomState) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 07a90ae713..1320aeb540 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -1,7 +1,7 @@ import { waitFor } from '@testing-library/dom' import { assert, describe, expect, it, vi } from 'vitest' import { atom, createStore } from 'jotai/vanilla' -import type { Atom, Getter } from 'jotai/vanilla' +import type { Getter } from 'jotai/vanilla' it('should not fire on subscribe', async () => { const store = createStore() @@ -556,6 +556,7 @@ describe('aborting atoms', () => { }) }) +/* skip for now describe('unstable_derive for scoping atoms', () => { it('primitive atom', async () => { const a = atom('a') @@ -705,3 +706,4 @@ describe('unstable_derive for scoping atoms', () => { expect(derivedStore.get(b)).toBe('a5') }) }) +*/ From 23b27348f8fdefc3cfa6dfde4ab490964689cf4a Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 17 Jul 2024 09:05:06 +0900 Subject: [PATCH 20/31] refactor 2 --- src/vanilla/store.ts | 55 ++++++++++++++++++++++++++------------------ 1 file changed, 32 insertions(+), 23 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 5c9c32ca5f..2f44d1f6c3 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -316,6 +316,28 @@ const addDependency = ( } } +const getDependents = ( + pending: Pending, + getAtomState: GetAtomState, + a: AnyAtom, +): Map => { + const aState = getAtomState(a) + const dependents = new Map() + for (const a of aState.m?.t || []) { + dependents.set(a, getAtomState(a)) + } + for (const atomWithPendingContinuablePromise of aState.p) { + dependents.set( + atomWithPendingContinuablePromise, + getAtomState(atomWithPendingContinuablePromise), + ) + } + getPendingDependents(pending, a)?.forEach((dependent) => { + dependents.set(dependent, getAtomState(dependent)) + }) + return dependents +} + type AtomStateMap = { get(key: Atom): AtomState | undefined set(key: Atom, value: AtomState): void @@ -480,51 +502,38 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { returnAtomValue(readAtomState(undefined, atom)) const recomputeDependents = (pending: Pending, atom: AnyAtom) => { - const getDependents = (a: AnyAtom): Set => { - const aState = getAtomState(a) - const dependents = new Set(aState.m?.t) - for (const atomWithPendingContinuablePromise of aState.p) { - dependents.add(atomWithPendingContinuablePromise) - } - getPendingDependents(pending, a)?.forEach((dependent) => { - dependents.add(dependent) - }) - return dependents - } - // This is a topological sort via depth-first search, slightly modified from // what's described here for simplicity and performance reasons: // https://en.wikipedia.org/wiki/Topological_sorting#Depth-first_search // Step 1: traverse the dependency graph to build the topsorted atom list // We don't bother to check for cycles, which simplifies the algorithm. - const topsortedAtoms: AnyAtom[] = [] + const topsortedAtoms: (readonly [AnyAtom, AtomState])[] = [] const markedAtoms = new Set() - const visit = (n: AnyAtom) => { - if (markedAtoms.has(n)) { + const visit = (a: AnyAtom, aState: AtomState) => { + if (markedAtoms.has(a)) { return } - markedAtoms.add(n) - for (const m of getDependents(n)) { - if (n !== m) { - visit(m) + markedAtoms.add(a) + for (const [m, s] of getDependents(pending, getAtomState, a)) { + if (a !== m) { + visit(m, s) } } // The algorithm calls for pushing onto the front of the list. For // performance, we will simply push onto the end, and then will iterate in // reverse order later. - topsortedAtoms.push(n) + topsortedAtoms.push([a, aState]) } // Visit the root atom. This is the only atom in the dependency graph // without incoming edges, which is one reason we can simplify the algorithm - visit(atom) + visit(atom, getAtomState(atom)) // Step 2: use the topsorted atom list to recompute all affected atoms // Track what's changed, so that we can short circuit when possible const changedAtoms = new Set([atom]) const isMarked = (a: AnyAtom) => markedAtoms.has(a) for (let i = topsortedAtoms.length - 1; i >= 0; --i) { - const a = topsortedAtoms[i]! - const aState = getAtomState(a) + const [a, aState] = topsortedAtoms[i]! const prevEpochNumber = aState.n let hasChangedDeps = false for (const dep of aState.d.keys()) { From eea418c7ea94e35ca466d1c06e4697574a94f9eb Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 17 Jul 2024 10:09:20 +0900 Subject: [PATCH 21/31] revert those refactors :-p --- src/vanilla/store.ts | 202 +++++++++++++++++++++---------------------- 1 file changed, 98 insertions(+), 104 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 2f44d1f6c3..9100984447 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -149,8 +149,6 @@ type AtomState = { e?: AnyError } -type GetAtomState = (atom: Atom) => AtomState - const isAtomStateInitialized = (atomState: AtomState) => 'v' in atomState || 'e' in atomState @@ -243,101 +241,6 @@ const flushPending = (pending: Pending) => { } } -const setAtomStateValueOrPromise = ( - getAtomState: GetAtomState, - atom: AnyAtom, - atomState: AtomState, - valueOrPromise: unknown, - abortPromise = () => {}, - completePromise = () => {}, -) => { - const hasPrevValue = 'v' in atomState - const prevValue = atomState.v - const pendingPromise = getPendingContinuablePromise(atomState) - if (isPromiseLike(valueOrPromise)) { - if (pendingPromise) { - if (pendingPromise !== valueOrPromise) { - pendingPromise[CONTINUE_PROMISE](valueOrPromise, abortPromise) - ++atomState.n - } - } else { - const continuablePromise = createContinuablePromise( - valueOrPromise, - abortPromise, - completePromise, - ) - if (continuablePromise.status === PENDING) { - for (const a of atomState.d.keys()) { - const aState = getAtomState(a) - addPendingContinuablePromiseToDependency( - atom, - continuablePromise, - aState, - ) - } - } - atomState.v = continuablePromise - delete atomState.e - } - } else { - if (pendingPromise) { - pendingPromise[CONTINUE_PROMISE]( - Promise.resolve(valueOrPromise), - abortPromise, - ) - } - atomState.v = valueOrPromise - delete atomState.e - } - if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { - ++atomState.n - } -} - -const addDependency = ( - pending: Pending | undefined, - getAtomState: GetAtomState, - atom: Atom, - a: AnyAtom, - aState: AtomState, -) => { - if (import.meta.env?.MODE !== 'production' && a === atom) { - throw new Error('[Bug] atom cannot depend on itself') - } - const atomState = getAtomState(atom) - atomState.d.set(a, aState.n) - const continuablePromise = getPendingContinuablePromise(atomState) - if (continuablePromise) { - addPendingContinuablePromiseToDependency(atom, continuablePromise, aState) - } - aState.m?.t.add(atom) - if (pending) { - addPendingDependent(pending, a, atom) - } -} - -const getDependents = ( - pending: Pending, - getAtomState: GetAtomState, - a: AnyAtom, -): Map => { - const aState = getAtomState(a) - const dependents = new Map() - for (const a of aState.m?.t || []) { - dependents.set(a, getAtomState(a)) - } - for (const atomWithPendingContinuablePromise of aState.p) { - dependents.set( - atomWithPendingContinuablePromise, - getAtomState(atomWithPendingContinuablePromise), - ) - } - getPendingDependents(pending, a)?.forEach((dependent) => { - dependents.set(dependent, getAtomState(dependent)) - }) - return dependents -} - type AtomStateMap = { get(key: Atom): AtomState | undefined set(key: Atom, value: AtomState): void @@ -388,6 +291,77 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { return atomState } + const setAtomStateValueOrPromise = ( + atom: AnyAtom, + atomState: AtomState, + valueOrPromise: unknown, + abortPromise = () => {}, + completePromise = () => {}, + ) => { + const hasPrevValue = 'v' in atomState + const prevValue = atomState.v + const pendingPromise = getPendingContinuablePromise(atomState) + if (isPromiseLike(valueOrPromise)) { + if (pendingPromise) { + if (pendingPromise !== valueOrPromise) { + pendingPromise[CONTINUE_PROMISE](valueOrPromise, abortPromise) + ++atomState.n + } + } else { + const continuablePromise = createContinuablePromise( + valueOrPromise, + abortPromise, + completePromise, + ) + if (continuablePromise.status === PENDING) { + for (const a of atomState.d.keys()) { + const aState = getAtomState(a) + addPendingContinuablePromiseToDependency( + atom, + continuablePromise, + aState, + ) + } + } + atomState.v = continuablePromise + delete atomState.e + } + } else { + if (pendingPromise) { + pendingPromise[CONTINUE_PROMISE]( + Promise.resolve(valueOrPromise), + abortPromise, + ) + } + atomState.v = valueOrPromise + delete atomState.e + } + if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { + ++atomState.n + } + } + + const addDependency = ( + pending: Pending | undefined, + atom: Atom, + a: AnyAtom, + aState: AtomState, + ) => { + if (import.meta.env?.MODE !== 'production' && a === atom) { + throw new Error('[Bug] atom cannot depend on itself') + } + const atomState = getAtomState(atom) + atomState.d.set(a, aState.n) + const continuablePromise = getPendingContinuablePromise(atomState) + if (continuablePromise) { + addPendingContinuablePromiseToDependency(atom, continuablePromise, aState) + } + aState.m?.t.add(atom) + if (pending) { + addPendingDependent(pending, a, atom) + } + } + const readAtomState = ( pending: Pending | undefined, atom: Atom, @@ -422,7 +396,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const aState = getAtomState(a) if (!isAtomStateInitialized(aState)) { if (hasInitialValue(a)) { - setAtomStateValueOrPromise(getAtomState, a, aState, a.init) + setAtomStateValueOrPromise(a, aState, a.init) } else { // NOTE invalid derived atoms can reach here throw new Error('no atom init') @@ -433,10 +407,10 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { // a !== atom const aState = readAtomState(pending, a, force) if (isSync) { - addDependency(pending, getAtomState, atom, a, aState) + addDependency(pending, atom, a, aState) } else { const pending = createPending() - addDependency(pending, getAtomState, atom, a, aState) + addDependency(pending, atom, a, aState) mountDependencies(pending, atom, atomState) flushPending(pending) } @@ -474,7 +448,6 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { try { const valueOrPromise = atom.read(getter, options as never) setAtomStateValueOrPromise( - getAtomState, atom, atomState, valueOrPromise, @@ -501,6 +474,27 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const readAtom = (atom: Atom): Value => returnAtomValue(readAtomState(undefined, atom)) + const getDependents = ( + pending: Pending, + a: AnyAtom, + ): Map => { + const aState = getAtomState(a) + const dependents = new Map() + for (const a of aState.m?.t || []) { + dependents.set(a, getAtomState(a)) + } + for (const atomWithPendingContinuablePromise of aState.p) { + dependents.set( + atomWithPendingContinuablePromise, + getAtomState(atomWithPendingContinuablePromise), + ) + } + getPendingDependents(pending, a)?.forEach((dependent) => { + dependents.set(dependent, getAtomState(dependent)) + }) + return dependents + } + const recomputeDependents = (pending: Pending, atom: AnyAtom) => { // This is a topological sort via depth-first search, slightly modified from // what's described here for simplicity and performance reasons: @@ -515,7 +509,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { return } markedAtoms.add(a) - for (const [m, s] of getDependents(pending, getAtomState, a)) { + for (const [m, s] of getDependents(pending, a)) { if (a !== m) { visit(m, s) } @@ -575,7 +569,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const hasPrevValue = 'v' in aState const prevValue = aState.v const v = args[0] as V - setAtomStateValueOrPromise(getAtomState, a, aState, v) + setAtomStateValueOrPromise(a, aState, v) mountDependencies(pending, a, aState) if (!hasPrevValue || !Object.is(prevValue, aState.v)) { addPendingAtom(pending, a, aState) @@ -729,7 +723,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const atomState = getAtomState(atom) const hasPrevValue = 'v' in atomState const prevValue = atomState.v - setAtomStateValueOrPromise(getAtomState, atom, atomState, value) + setAtomStateValueOrPromise(atom, atomState, value) mountDependencies(pending, atom, atomState) if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { addPendingAtom(pending, atom, atomState) From 19423bb2a47b8cfd3e4c11860df5eb5fca41cf0b Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 17 Jul 2024 10:41:17 +0900 Subject: [PATCH 22/31] wip: hm, how is bundle size? --- src/vanilla/store.ts | 192 ++++++++++++++++++++++++------------------- 1 file changed, 109 insertions(+), 83 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 9100984447..9c3eb12a07 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -188,6 +188,27 @@ const addPendingContinuablePromiseToDependency = ( } } +const addDependency = ( + pending: Pending | undefined, + atom: Atom, + atomState: AtomState, + a: AnyAtom, + aState: AtomState, +) => { + if (import.meta.env?.MODE !== 'production' && a === atom) { + throw new Error('[Bug] atom cannot depend on itself') + } + atomState.d.set(a, aState.n) + const continuablePromise = getPendingContinuablePromise(atomState) + if (continuablePromise) { + addPendingContinuablePromiseToDependency(atom, continuablePromise, aState) + } + aState.m?.t.add(atom) + if (pending) { + addPendingDependent(pending, a, atom) + } +} + // // Pending // @@ -241,20 +262,22 @@ const flushPending = (pending: Pending) => { } } -type AtomStateMap = { - get(key: Atom): AtomState | undefined - set(key: Atom, value: AtomState): void -} +type GetAtomState = ( + atom: Atom, + originAtomState?: AtomState, +) => AtomState // internal & unstable type type StoreArgs = readonly [ - atomStateMap: AtomStateMap, + getAtomState: GetAtomState, // possible other arguments in the future ] // for debugging purpose only type DevStoreRev4 = { - dev4_get_internal_weak_map: () => AtomStateMap + dev4_get_internal_weak_map: () => { + get: (atom: AnyAtom) => AtomState | undefined + } dev4_get_mounted_atoms: () => Set dev4_restore_atoms: (values: Iterable) => void } @@ -274,7 +297,7 @@ type Store = PrdStore | (PrdStore & DevStoreRev4) export type INTERNAL_DevStoreRev4 = DevStoreRev4 export type INTERNAL_PrdStore = PrdStore -const buildStore = (atomStateMap: StoreArgs[0]): Store => { +const buildStore = (getAtomState: StoreArgs[0]): Store => { // for debugging purpose only let debugMountedAtoms: Set @@ -282,15 +305,6 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { debugMountedAtoms = new Set() } - const getAtomState = (atom: Atom) => { - let atomState = atomStateMap.get(atom) - if (!atomState) { - atomState = { d: new Map(), p: new Set(), n: 0 } - atomStateMap.set(atom, atomState) - } - return atomState - } - const setAtomStateValueOrPromise = ( atom: AnyAtom, atomState: AtomState, @@ -315,11 +329,10 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { ) if (continuablePromise.status === PENDING) { for (const a of atomState.d.keys()) { - const aState = getAtomState(a) addPendingContinuablePromiseToDependency( atom, continuablePromise, - aState, + getAtomState(a, atomState), ) } } @@ -341,34 +354,13 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { } } - const addDependency = ( - pending: Pending | undefined, - atom: Atom, - a: AnyAtom, - aState: AtomState, - ) => { - if (import.meta.env?.MODE !== 'production' && a === atom) { - throw new Error('[Bug] atom cannot depend on itself') - } - const atomState = getAtomState(atom) - atomState.d.set(a, aState.n) - const continuablePromise = getPendingContinuablePromise(atomState) - if (continuablePromise) { - addPendingContinuablePromiseToDependency(atom, continuablePromise, aState) - } - aState.m?.t.add(atom) - if (pending) { - addPendingDependent(pending, a, atom) - } - } - const readAtomState = ( pending: Pending | undefined, atom: Atom, + atomState: AtomState, force?: (a: AnyAtom) => boolean, ): AtomState => { // See if we can skip recomputing this atom. - const atomState = getAtomState(atom) if (!force?.(atom) && isAtomStateInitialized(atomState)) { // If the atom is mounted, we can use the cache. // because it should have been updated by dependencies. @@ -382,7 +374,8 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { ([a, n]) => // Recursively, read the atom state of the dependency, and // check if the atom epoch number is unchanged - readAtomState(pending, a, force).n === n, + readAtomState(pending, a, getAtomState(a, atomState), force).n === + n, ) ) { return atomState @@ -393,7 +386,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { let isSync = true const getter: Getter = (a: Atom) => { if (a === (atom as AnyAtom)) { - const aState = getAtomState(a) + const aState = getAtomState(a, atomState) if (!isAtomStateInitialized(aState)) { if (hasInitialValue(a)) { setAtomStateValueOrPromise(a, aState, a.init) @@ -405,12 +398,17 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { return returnAtomValue(aState) } // a !== atom - const aState = readAtomState(pending, a, force) + const aState = readAtomState( + pending, + a, + getAtomState(a, atomState), + force, + ) if (isSync) { - addDependency(pending, atom, a, aState) + addDependency(pending, atom, atomState, a, aState) } else { const pending = createPending() - addDependency(pending, atom, a, aState) + addDependency(pending, atom, atomState, a, aState) mountDependencies(pending, atom, atomState) flushPending(pending) } @@ -472,30 +470,34 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { } const readAtom = (atom: Atom): Value => - returnAtomValue(readAtomState(undefined, atom)) + returnAtomValue(readAtomState(undefined, atom, getAtomState(atom))) - const getDependents = ( + const getDependents = ( pending: Pending, - a: AnyAtom, + atom: Atom, + atomState: AtomState, ): Map => { - const aState = getAtomState(a) const dependents = new Map() - for (const a of aState.m?.t || []) { - dependents.set(a, getAtomState(a)) + for (const a of atomState.m?.t || []) { + dependents.set(a, getAtomState(a, atomState)) } - for (const atomWithPendingContinuablePromise of aState.p) { + for (const atomWithPendingContinuablePromise of atomState.p) { dependents.set( atomWithPendingContinuablePromise, - getAtomState(atomWithPendingContinuablePromise), + getAtomState(atomWithPendingContinuablePromise, atomState), ) } - getPendingDependents(pending, a)?.forEach((dependent) => { - dependents.set(dependent, getAtomState(dependent)) + getPendingDependents(pending, atom)?.forEach((dependent) => { + dependents.set(dependent, getAtomState(dependent, atomState)) }) return dependents } - const recomputeDependents = (pending: Pending, atom: AnyAtom) => { + const recomputeDependents = ( + pending: Pending, + atom: Atom, + atomState: AtomState, + ) => { // This is a topological sort via depth-first search, slightly modified from // what's described here for simplicity and performance reasons: // https://en.wikipedia.org/wiki/Topological_sorting#Depth-first_search @@ -509,7 +511,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { return } markedAtoms.add(a) - for (const [m, s] of getDependents(pending, a)) { + for (const [m, s] of getDependents(pending, a, aState)) { if (a !== m) { visit(m, s) } @@ -521,7 +523,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { } // Visit the root atom. This is the only atom in the dependency graph // without incoming edges, which is one reason we can simplify the algorithm - visit(atom, getAtomState(atom)) + visit(atom, atomState) // Step 2: use the topsorted atom list to recompute all affected atoms // Track what's changed, so that we can short circuit when possible const changedAtoms = new Set([atom]) @@ -537,7 +539,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { } } if (hasChangedDeps) { - readAtomState(pending, a, isMarked) + readAtomState(pending, a, aState, isMarked) mountDependencies(pending, a, aState) if (prevEpochNumber !== aState.n) { addPendingAtom(pending, a, aState) @@ -551,21 +553,22 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const writeAtomState = ( pending: Pending, atom: WritableAtom, + atomState: AtomState, ...args: Args ): Result => { const getter: Getter = (a: Atom) => - returnAtomValue(readAtomState(pending, a)) + returnAtomValue(readAtomState(pending, a, getAtomState(a, atomState))) const setter: Setter = ( a: WritableAtom, ...args: As ) => { + const aState = getAtomState(a, atomState) let r: R | undefined if (a === (atom as AnyAtom)) { if (!hasInitialValue(a)) { // NOTE technically possible but restricted as it may cause bugs throw new Error('atom not writable') } - const aState = getAtomState(a) const hasPrevValue = 'v' in aState const prevValue = aState.v const v = args[0] as V @@ -573,10 +576,10 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { mountDependencies(pending, a, aState) if (!hasPrevValue || !Object.is(prevValue, aState.v)) { addPendingAtom(pending, a, aState) - recomputeDependents(pending, a) + recomputeDependents(pending, a, aState) } } else { - r = writeAtomState(pending, a, ...args) as R + r = writeAtomState(pending, a, aState, ...args) as R } flushPending(pending) return r as R @@ -590,7 +593,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { ...args: Args ): Result => { const pending = createPending() - const result = writeAtomState(pending, atom, ...args) + const result = writeAtomState(pending, atom, getAtomState(atom), ...args) flushPending(pending) return result } @@ -603,14 +606,14 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { if (atomState.m && !getPendingContinuablePromise(atomState)) { for (const a of atomState.d.keys()) { if (!atomState.m.d.has(a)) { - const aMounted = mountAtom(pending, a) + const aMounted = mountAtom(pending, a, getAtomState(a, atomState)) aMounted.t.add(atom) atomState.m.d.add(a) } } for (const a of atomState.m.d || []) { if (!atomState.d.has(a)) { - const aMounted = unmountAtom(pending, a) + const aMounted = unmountAtom(pending, a, getAtomState(a, atomState)) aMounted?.t.delete(atom) atomState.m.d.delete(a) } @@ -618,14 +621,17 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { } } - const mountAtom = (pending: Pending, atom: AnyAtom): Mounted => { - const atomState = getAtomState(atom) + const mountAtom = ( + pending: Pending, + atom: Atom, + atomState: AtomState, + ): Mounted => { if (!atomState.m) { // recompute atom state - readAtomState(pending, atom) + readAtomState(pending, atom, atomState) // mount dependencies first for (const a of atomState.d.keys()) { - const aMounted = mountAtom(pending, a) + const aMounted = mountAtom(pending, a, getAtomState(a, atomState)) aMounted.t.add(atom) } // mount self @@ -642,7 +648,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const { onMount } = atom addPendingFunction(pending, () => { const onUnmount = onMount((...args) => - writeAtomState(pending, atom, ...args), + writeAtomState(pending, atom, atomState, ...args), ) if (onUnmount) { mounted.u = onUnmount @@ -653,15 +659,15 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { return atomState.m } - const unmountAtom = ( + const unmountAtom = ( pending: Pending, - atom: AnyAtom, + atom: Atom, + atomState: AtomState, ): Mounted | undefined => { - const atomState = getAtomState(atom) if ( atomState.m && !atomState.m.l.size && - !Array.from(atomState.m.t).some((a) => getAtomState(a).m) + !Array.from(atomState.m.t).some((a) => getAtomState(a, atomState).m) ) { // unmount self const onUnmount = atomState.m.u @@ -674,7 +680,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { } // unmount dependencies for (const a of atomState.d.keys()) { - const aMounted = unmountAtom(pending, a) + const aMounted = unmountAtom(pending, a, getAtomState(a, atomState)) aMounted?.t.delete(atom) } // abort pending promise @@ -690,20 +696,21 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { const subscribeAtom = (atom: AnyAtom, listener: () => void) => { const pending = createPending() - const mounted = mountAtom(pending, atom) + const atomState = getAtomState(atom) + const mounted = mountAtom(pending, atom, atomState) flushPending(pending) const listeners = mounted.l listeners.add(listener) return () => { listeners.delete(listener) const pending = createPending() - unmountAtom(pending, atom) + unmountAtom(pending, atom, atomState) flushPending(pending) } } const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) => - buildStore(...fn(atomStateMap)) + buildStore(...fn(getAtomState)) const store: Store = { get: readAtom, @@ -714,7 +721,16 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { if (import.meta.env?.MODE !== 'production') { const devStore: DevStoreRev4 = { // store dev methods (these are tentative and subject to change without notice) - dev4_get_internal_weak_map: () => atomStateMap, + dev4_get_internal_weak_map: () => ({ + get: (atom) => { + const atomState = getAtomState(atom) + if (atomState.n === 0) { + // for backward compatibility + return undefined + } + return atomState + }, + }), dev4_get_mounted_atoms: () => debugMountedAtoms, dev4_restore_atoms: (values) => { const pending = createPending() @@ -727,7 +743,7 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { mountDependencies(pending, atom, atomState) if (!hasPrevValue || !Object.is(prevValue, atomState.v)) { addPendingAtom(pending, atom, atomState) - recomputeDependents(pending, atom) + recomputeDependents(pending, atom, atomState) } } } @@ -739,8 +755,18 @@ const buildStore = (atomStateMap: StoreArgs[0]): Store => { return store } -export const createStore = (): Store => - buildStore(new WeakMap() as AtomStateMap) +export const createStore = (): Store => { + const atomStateMap = new WeakMap() + const getAtomState = (atom: Atom) => { + let atomState = atomStateMap.get(atom) + if (!atomState) { + atomState = { d: new Map(), p: new Set(), n: 0 } + atomStateMap.set(atom, atomState) + } + return atomState + } + return buildStore(getAtomState) +} let defaultStore: Store | undefined From 2fc0c9773548b87a203821f4176fa9247c38b1f4 Mon Sep 17 00:00:00 2001 From: daishi Date: Wed, 17 Jul 2024 10:50:59 +0900 Subject: [PATCH 23/31] fix tests --- tests/vanilla/store.test.tsx | 117 +++++++++++++++-------------------- 1 file changed, 51 insertions(+), 66 deletions(-) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 1320aeb540..a1f693a0d2 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -1,7 +1,7 @@ import { waitFor } from '@testing-library/dom' import { assert, describe, expect, it, vi } from 'vitest' import { atom, createStore } from 'jotai/vanilla' -import type { Getter } from 'jotai/vanilla' +import type { Atom, Getter } from 'jotai/vanilla' it('should not fire on subscribe', async () => { const store = createStore() @@ -556,7 +556,6 @@ describe('aborting atoms', () => { }) }) -/* skip for now describe('unstable_derive for scoping atoms', () => { it('primitive atom', async () => { const a = atom('a') @@ -564,25 +563,22 @@ describe('unstable_derive for scoping atoms', () => { const scopedAtoms = new Set>([a]) const store = createStore() - const derivedStore = store.unstable_derive( - (getAtomState, getAtomContext) => { - const scopedAtomStateMap = new WeakMap() - return [ - (atom: Atom, context: unknown) => { - if (scopedAtoms.has(atom)) { - let atomState = scopedAtomStateMap.get(atom) - if (!atomState) { - atomState = { d: new Map(), p: new Set(), n: 0 } - scopedAtomStateMap.set(atom, atomState) - } - return atomState + const derivedStore = store.unstable_derive((getAtomState) => { + const scopedAtomStateMap = new WeakMap() + return [ + (atom, originAtomState) => { + if (scopedAtoms.has(atom)) { + let atomState = scopedAtomStateMap.get(atom) + if (!atomState) { + atomState = { d: new Map(), p: new Set(), n: 0 } + scopedAtomStateMap.set(atom, atomState) } - return getAtomState(atom, context) - }, - getAtomContext, - ] - }, - ) + return atomState + } + return getAtomState(atom, originAtomState) + }, + ] + }) expect(store.get(a)).toBe('a') expect(derivedStore.get(a)).toBe('a') @@ -604,25 +600,22 @@ describe('unstable_derive for scoping atoms', () => { const scopedAtoms = new Set>([a]) const store = createStore() - const derivedStore = store.unstable_derive( - (getAtomState, getAtomContext) => { - const scopedAtomStateMap = new WeakMap() - return [ - (atom: Atom, context: unknown) => { - if (scopedAtoms.has(atom)) { - let atomState = scopedAtomStateMap.get(atom) - if (!atomState) { - atomState = { d: new Map(), p: new Set(), n: 0 } - scopedAtomStateMap.set(atom, atomState) - } - return atomState + const derivedStore = store.unstable_derive((getAtomState) => { + const scopedAtomStateMap = new WeakMap() + return [ + (atom, originAtomState) => { + if (scopedAtoms.has(atom)) { + let atomState = scopedAtomStateMap.get(atom) + if (!atomState) { + atomState = { d: new Map(), p: new Set(), n: 0 } + scopedAtomStateMap.set(atom, atomState) } - return getAtomState(atom, context) - }, - getAtomContext, - ] - }, - ) + return atomState + } + return getAtomState(atom, originAtomState) + }, + ] + }) expect(store.get(b)).toBe('a') expect(derivedStore.get(b)).toBe('a') @@ -643,34 +636,27 @@ describe('unstable_derive for scoping atoms', () => { const scopedAtoms = new Set>([b]) const store = createStore() - const derivedStore = store.unstable_derive( - (getAtomState, getAtomContext) => { - const scopedAtomStateMap = new WeakMap() - return [ - (atom: Atom, context: unknown) => { - if ( - (context as { scoped: unknown } | undefined)?.scoped || - scopedAtoms.has(atom) - ) { - let atomState = scopedAtomStateMap.get(atom) - if (!atomState) { - atomState = { d: new Map(), p: new Set(), n: 0 } - scopedAtomStateMap.set(atom, atomState) - } - return atomState - } - return getAtomState(atom, context) - }, - (atom, context) => { - context = getAtomContext(atom, context) - if (scopedAtoms.has(atom)) { - return { ...(context as object), scoped: true } + const derivedStore = store.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 context - }, - ] - }, - ) + return atomState + } + return getAtomState(atom, originAtomState) + }, + ] + }) expect(store.get(a)).toBe('a') expect(store.get(b)).toBe('a') @@ -706,4 +692,3 @@ describe('unstable_derive for scoping atoms', () => { expect(derivedStore.get(b)).toBe('a5') }) }) -*/ From b30da262aa0668b707f86d34f55f05d8c968e364 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Fri, 19 Jul 2024 04:09:52 -0700 Subject: [PATCH 24/31] add test for unstable_derive (#2665) --- tests/vanilla/store.test.tsx | 175 +++++++++++++++++++++++++++++++++-- 1 file changed, 169 insertions(+), 6 deletions(-) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index a1f693a0d2..e60b547776 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -557,6 +557,10 @@ describe('aborting atoms', () => { }) describe('unstable_derive for scoping atoms', () => { + /** + * a + * S1[a]: a1 + */ it('primitive atom', async () => { const a = atom('a') a.onMount = (setSelf) => setSelf((v) => v + ':mounted') @@ -594,9 +598,14 @@ describe('unstable_derive for scoping atoms', () => { expect(derivedStore.get(a)).toBe('a:mounted:updated') }) + /** + * a, b, c(a + b) + * S1[a]: a1, b0, c0(a1 + b0) + */ it('derived atom (scoping primitive)', async () => { const a = atom('a') - const b = atom((get) => get(a)) + const b = atom('b') + const c = atom((get) => get(a) + get(b)) const scopedAtoms = new Set>([a]) const store = createStore() @@ -617,14 +626,18 @@ describe('unstable_derive for scoping atoms', () => { ] }) - expect(store.get(b)).toBe('a') - expect(derivedStore.get(b)).toBe('a') - derivedStore.set(a, 'b') + expect(store.get(c)).toBe('ab') + expect(derivedStore.get(c)).toBe('ab') + derivedStore.set(a, 'a2') await new Promise((resolve) => setTimeout(resolve)) - expect(store.get(b)).toBe('a') - expect(derivedStore.get(b)).toBe('b') + expect(store.get(c)).toBe('ab') + expect(derivedStore.get(c)).toBe('a2b') }) + /** + * a, b(a) + * S1[b]: a0, b1(a1) + */ it('derived atom (scoping derived)', async () => { const a = atom('a') const b = atom( @@ -691,4 +704,154 @@ describe('unstable_derive for scoping atoms', () => { expect(derivedStore.get(a)).toBe('a4') expect(derivedStore.get(b)).toBe('a5') }) + + /** + * a, b, c(a), d(c), e(d + b) + * S1[d]: a0, b0, c0(a0), d1(c1(a1)), e0(d1(c1(a1)) + b0) + */ + it('derived atom (scoping derived chain)', async () => { + const a = atom('a') + const b = atom('b') + const c = atom( + (get) => get(a), + (_get, set, v: string) => set(a, v), + ) + const d = atom( + (get) => get(c), + (_get, set, v: string) => set(c, v), + ) + const e = atom( + (get) => get(d) + get(b), + (_get, set, av: string, bv: string) => { + set(d, av) + set(b, bv) + }, + ) + const scopedAtoms = new Set>([d]) + + function makeStores() { + const baseStore = createStore() + const deriStore = 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) + }, + ] + }) + expect(getAtoms(baseStore)).toEqual(['a', 'b', 'a', 'a', 'ab']) + expect(getAtoms(deriStore)).toEqual(['a', 'b', 'a', 'a', 'ab']) + return { baseStore, deriStore } + } + type Store = ReturnType + function getAtoms(store: Store) { + return [ + store.get(a), + store.get(b), + store.get(c), + store.get(d), + store.get(e), + ] + } + + /** + * base[d]: a0, b0, c0(a0), d0(c0(a0)), e0(d0(c0(a0)) + b0) + * deri[d]: a0, b0, c0(a0), d1(c1(a1)), e0(d1(c1(a1)) + b0) + */ + { + // UPDATE a0 + // NOCHGE b0 and a1 + const { baseStore, deriStore } = makeStores() + baseStore.set(a, '*') + expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) + expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) + } + { + // UPDATE b0 + // NOCHGE a0 and a1 + const { baseStore, deriStore } = makeStores() + baseStore.set(b, '*') + expect(getAtoms(baseStore)).toEqual(['a', '*', 'a', 'a', 'a*']) + expect(getAtoms(deriStore)).toEqual(['a', '*', 'a', 'a', 'a*']) + } + { + // UPDATE c0, c0 -> a0 + // NOCHGE b0 and a1 + const { baseStore, deriStore } = makeStores() + baseStore.set(c, '*') + expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) + expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) + } + { + // UPDATE d0, d0 -> c0 -> a0 + // NOCHGE b0 and a1 + const { baseStore, deriStore } = makeStores() + baseStore.set(d, '*') + expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) + expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) + } + { + // UPDATE e0, e0 -> d0 -> c0 -> a0 + // └--------------> b0 + // NOCHGE a1 + const { baseStore, deriStore } = makeStores() + baseStore.set(e, '*', '*') + expect(getAtoms(baseStore)).toEqual(['*', '*', '*', '*', '**']) + expect(getAtoms(deriStore)).toEqual(['*', '*', '*', 'a', 'a*']) + } + { + // UPDATE a0 + // NOCHGE b0 and a1 + const { baseStore, deriStore } = makeStores() + deriStore.set(a, '*') + expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) + expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) + } + { + // UPDATE b0 + // NOCHGE a0 and a1 + const { baseStore, deriStore } = makeStores() + deriStore.set(b, '*') + expect(getAtoms(baseStore)).toEqual(['a', '*', 'a', 'a', 'a*']) + expect(getAtoms(deriStore)).toEqual(['a', '*', 'a', 'a', 'a*']) + } + { + // UPDATE c0, c0 -> a0 + // NOCHGE b0 and a1 + const { baseStore, deriStore } = makeStores() + deriStore.set(c, '*') + expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) + expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) + } + { + // UPDATE d1, d1 -> c1 -> a1 + // NOCHGE b0 and a0 + const { baseStore, deriStore } = makeStores() + deriStore.set(d, '*') + expect(getAtoms(baseStore)).toEqual(['a', 'b', 'a', 'a', 'ab']) + expect(getAtoms(deriStore)).toEqual(['a', 'b', 'a', '*', '*b']) + } + { + // UPDATE e0, e0 -> d1 -> c1 -> a1 + // └--------------> b0 + // NOCHGE a0 + const { baseStore, deriStore } = makeStores() + deriStore.set(e, '*', '*') + expect(getAtoms(baseStore)).toEqual(['a', '*', 'a', 'a', 'a*']) + expect(getAtoms(deriStore)).toEqual(['a', '*', 'a', '*', '**']) + } + }) }) From 09e907f196029e1bf9474b6ad42ac7b0684e1f3e Mon Sep 17 00:00:00 2001 From: daishi Date: Sat, 20 Jul 2024 22:49:01 +0900 Subject: [PATCH 25/31] add a failing test --- tests/vanilla/store.test.tsx | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index e60b547776..73e946d8f2 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -628,10 +628,17 @@ describe('unstable_derive for scoping atoms', () => { expect(store.get(c)).toBe('ab') expect(derivedStore.get(c)).toBe('ab') + derivedStore.set(a, 'a2') await new Promise((resolve) => setTimeout(resolve)) expect(store.get(c)).toBe('ab') expect(derivedStore.get(c)).toBe('a2b') + + derivedStore.sub(c, vi.fn()) + derivedStore.set(b, 'b2') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(c)).toBe('ab2') + expect(derivedStore.get(c)).toBe('a2b2') }) /** From 49549127887e7939f63380f20a43ae47127329e9 Mon Sep 17 00:00:00 2001 From: daishi Date: Sun, 21 Jul 2024 10:16:32 +0900 Subject: [PATCH 26/31] wip: failed attempt --- src/vanilla/store.ts | 2 +- tests/vanilla/store.test.tsx | 55 +++++++++++++++++++++++++++++------- 2 files changed, 46 insertions(+), 11 deletions(-) diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 9c3eb12a07..d3bd07efde 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -758,7 +758,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { export const createStore = (): Store => { const atomStateMap = new WeakMap() const getAtomState = (atom: Atom) => { - let atomState = atomStateMap.get(atom) + let atomState = atomStateMap.get(atom) as AtomState | undefined if (!atomState) { atomState = { d: new Map(), p: new Set(), n: 0 } atomStateMap.set(atom, atomState) diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index 73e946d8f2..a9d33a6d44 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -1,7 +1,7 @@ import { waitFor } from '@testing-library/dom' import { assert, describe, expect, it, vi } from 'vitest' import { atom, createStore } from 'jotai/vanilla' -import type { Atom, Getter } from 'jotai/vanilla' +import type { Atom, Getter, ExtractAtomValue } from 'jotai/vanilla' it('should not fire on subscribe', async () => { const store = createStore() @@ -610,18 +610,53 @@ describe('unstable_derive for scoping atoms', () => { const store = createStore() const derivedStore = store.unstable_derive((getAtomState) => { - const scopedAtomStateMap = new WeakMap() + type AnyAtom = Parameters[0] + type AtomState = ReturnType + const scopedAtomStateMap = new WeakMap() + const scopedInvertAtomStateMap = new WeakMap() + const copyMounted = ( + mounted: NonNullable, + ): NonNullable => ({ + ...mounted, + d: new Set(mounted.d), + t: new Set(mounted.t), + }) + const copyAtomState = (atomState: AtomState): AtomState => ({ + ...atomState, + d: new Map(atomState.d), + p: new Set(atomState.p), + ...('m' in atomState ? { m: copyMounted(atomState.m) } : {}), + }) return [ (atom, originAtomState) => { - if (scopedAtoms.has(atom)) { - let atomState = scopedAtomStateMap.get(atom) - if (!atomState) { - atomState = { d: new Map(), p: new Set(), n: 0 } - scopedAtomStateMap.set(atom, atomState) - } - return atomState + type TheAtomState = ReturnType< + typeof getAtomState> + > + let atomState = scopedAtomStateMap.get(atom) + if (atomState) { + return atomState as TheAtomState } - return getAtomState(atom, originAtomState) + if ( + scopedInvertAtomStateMap.has(originAtomState as never) || + scopedAtoms.has(atom) + ) { + atomState = { d: new Map(), p: new Set(), n: 0 } + scopedAtomStateMap.set(atom, atomState) + scopedInvertAtomStateMap.set(atomState, atom) + return atomState as TheAtomState + } + const originalAtomState = getAtomState(atom, originAtomState) + if ( + Array.from(originalAtomState.d).some(([a]) => + scopedAtomStateMap.has(a), + ) + ) { + atomState = copyAtomState(originalAtomState) + scopedAtomStateMap.set(atom, atomState) + scopedInvertAtomStateMap.set(atomState, atom) + return atomState as TheAtomState + } + return originalAtomState }, ] }) From 6fca9157c3384932ca996af5000cb31a4b6ee28b Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 23 Jul 2024 18:46:35 +0900 Subject: [PATCH 27/31] split unstable_derive test --- tests/vanilla/store.test.tsx | 344 +----------------------- tests/vanilla/unstable_derive.test.tsx | 345 +++++++++++++++++++++++++ 2 files changed, 346 insertions(+), 343 deletions(-) create mode 100644 tests/vanilla/unstable_derive.test.tsx diff --git a/tests/vanilla/store.test.tsx b/tests/vanilla/store.test.tsx index df8ff3fa64..1c9e6d212f 100644 --- a/tests/vanilla/store.test.tsx +++ b/tests/vanilla/store.test.tsx @@ -1,7 +1,7 @@ import { waitFor } from '@testing-library/dom' import { assert, describe, expect, it, vi } from 'vitest' import { atom, createStore } from 'jotai/vanilla' -import type { Atom, ExtractAtomValue, Getter } from 'jotai/vanilla' +import type { Getter } from 'jotai/vanilla' it('should not fire on subscribe', async () => { const store = createStore() @@ -572,345 +572,3 @@ it('Unmount an atom that is no longer dependent within a derived atom (#2658)', store.set(condAtom, false) expect(onUnmount).toHaveBeenCalledTimes(1) }) - -describe('unstable_derive for scoping atoms', () => { - /** - * a - * S1[a]: a1 - */ - it('primitive atom', async () => { - const a = atom('a') - a.onMount = (setSelf) => setSelf((v) => v + ':mounted') - const scopedAtoms = new Set>([a]) - - const store = createStore() - const derivedStore = store.unstable_derive((getAtomState) => { - const scopedAtomStateMap = new WeakMap() - return [ - (atom, originAtomState) => { - if (scopedAtoms.has(atom)) { - let atomState = scopedAtomStateMap.get(atom) - if (!atomState) { - atomState = { d: new Map(), p: new Set(), n: 0 } - scopedAtomStateMap.set(atom, atomState) - } - return atomState - } - return getAtomState(atom, originAtomState) - }, - ] - }) - - expect(store.get(a)).toBe('a') - expect(derivedStore.get(a)).toBe('a') - - derivedStore.sub(a, vi.fn()) - await new Promise((resolve) => setTimeout(resolve)) - expect(store.get(a)).toBe('a') - expect(derivedStore.get(a)).toBe('a:mounted') - - derivedStore.set(a, (v) => v + ':updated') - await new Promise((resolve) => setTimeout(resolve)) - expect(store.get(a)).toBe('a') - expect(derivedStore.get(a)).toBe('a:mounted:updated') - }) - - /** - * a, b, c(a + b) - * S1[a]: a1, b0, c0(a1 + b0) - */ - it('derived atom (scoping primitive)', async () => { - const a = atom('a') - const b = atom('b') - const c = atom((get) => get(a) + get(b)) - const scopedAtoms = new Set>([a]) - - const store = createStore() - const derivedStore = store.unstable_derive((getAtomState) => { - type AnyAtom = Parameters[0] - type AtomState = ReturnType - const scopedAtomStateMap = new WeakMap() - const scopedInvertAtomStateMap = new WeakMap() - const copyMounted = ( - mounted: NonNullable, - ): NonNullable => ({ - ...mounted, - d: new Set(mounted.d), - t: new Set(mounted.t), - }) - const copyAtomState = (atomState: AtomState): AtomState => ({ - ...atomState, - d: new Map(atomState.d), - p: new Set(atomState.p), - ...('m' in atomState ? { m: copyMounted(atomState.m) } : {}), - }) - return [ - (atom, originAtomState) => { - type TheAtomState = ReturnType< - typeof getAtomState> - > - let atomState = scopedAtomStateMap.get(atom) - if (atomState) { - return atomState as TheAtomState - } - if ( - scopedInvertAtomStateMap.has(originAtomState as never) || - scopedAtoms.has(atom) - ) { - atomState = { d: new Map(), p: new Set(), n: 0 } - scopedAtomStateMap.set(atom, atomState) - scopedInvertAtomStateMap.set(atomState, atom) - return atomState as TheAtomState - } - const originalAtomState = getAtomState(atom, originAtomState) - if ( - Array.from(originalAtomState.d).some(([a]) => - scopedAtomStateMap.has(a), - ) - ) { - atomState = copyAtomState(originalAtomState) - scopedAtomStateMap.set(atom, atomState) - scopedInvertAtomStateMap.set(atomState, atom) - return atomState as TheAtomState - } - return originalAtomState - }, - ] - }) - - expect(store.get(c)).toBe('ab') - expect(derivedStore.get(c)).toBe('ab') - - derivedStore.set(a, 'a2') - await new Promise((resolve) => setTimeout(resolve)) - expect(store.get(c)).toBe('ab') - expect(derivedStore.get(c)).toBe('a2b') - - derivedStore.sub(c, vi.fn()) - derivedStore.set(b, 'b2') - await new Promise((resolve) => setTimeout(resolve)) - expect(store.get(c)).toBe('ab2') - expect(derivedStore.get(c)).toBe('a2b2') - }) - - /** - * a, b(a) - * S1[b]: a0, b1(a1) - */ - it('derived atom (scoping derived)', async () => { - const a = atom('a') - const b = atom( - (get) => get(a), - (_get, set, v: string) => { - set(a, v) - }, - ) - const scopedAtoms = new Set>([b]) - - const store = createStore() - const derivedStore = store.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) - }, - ] - }) - - expect(store.get(a)).toBe('a') - expect(store.get(b)).toBe('a') - expect(derivedStore.get(a)).toBe('a') - expect(derivedStore.get(b)).toBe('a') - - store.set(a, 'a2') - await new Promise((resolve) => setTimeout(resolve)) - expect(store.get(a)).toBe('a2') - expect(store.get(b)).toBe('a2') - expect(derivedStore.get(a)).toBe('a2') - expect(derivedStore.get(b)).toBe('a') - - store.set(b, 'a3') - await new Promise((resolve) => setTimeout(resolve)) - expect(store.get(a)).toBe('a3') - expect(store.get(b)).toBe('a3') - expect(derivedStore.get(a)).toBe('a3') - expect(derivedStore.get(b)).toBe('a') - - derivedStore.set(a, 'a4') - await new Promise((resolve) => setTimeout(resolve)) - expect(store.get(a)).toBe('a4') - expect(store.get(b)).toBe('a4') - expect(derivedStore.get(a)).toBe('a4') - expect(derivedStore.get(b)).toBe('a') - - derivedStore.set(b, 'a5') - await new Promise((resolve) => setTimeout(resolve)) - expect(store.get(a)).toBe('a4') - expect(store.get(b)).toBe('a4') - expect(derivedStore.get(a)).toBe('a4') - expect(derivedStore.get(b)).toBe('a5') - }) - - /** - * a, b, c(a), d(c), e(d + b) - * S1[d]: a0, b0, c0(a0), d1(c1(a1)), e0(d1(c1(a1)) + b0) - */ - it('derived atom (scoping derived chain)', async () => { - const a = atom('a') - const b = atom('b') - const c = atom( - (get) => get(a), - (_get, set, v: string) => set(a, v), - ) - const d = atom( - (get) => get(c), - (_get, set, v: string) => set(c, v), - ) - const e = atom( - (get) => get(d) + get(b), - (_get, set, av: string, bv: string) => { - set(d, av) - set(b, bv) - }, - ) - const scopedAtoms = new Set>([d]) - - function makeStores() { - const baseStore = createStore() - const deriStore = 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) - }, - ] - }) - expect(getAtoms(baseStore)).toEqual(['a', 'b', 'a', 'a', 'ab']) - expect(getAtoms(deriStore)).toEqual(['a', 'b', 'a', 'a', 'ab']) - return { baseStore, deriStore } - } - type Store = ReturnType - function getAtoms(store: Store) { - return [ - store.get(a), - store.get(b), - store.get(c), - store.get(d), - store.get(e), - ] - } - - /** - * base[d]: a0, b0, c0(a0), d0(c0(a0)), e0(d0(c0(a0)) + b0) - * deri[d]: a0, b0, c0(a0), d1(c1(a1)), e0(d1(c1(a1)) + b0) - */ - { - // UPDATE a0 - // NOCHGE b0 and a1 - const { baseStore, deriStore } = makeStores() - baseStore.set(a, '*') - expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) - expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) - } - { - // UPDATE b0 - // NOCHGE a0 and a1 - const { baseStore, deriStore } = makeStores() - baseStore.set(b, '*') - expect(getAtoms(baseStore)).toEqual(['a', '*', 'a', 'a', 'a*']) - expect(getAtoms(deriStore)).toEqual(['a', '*', 'a', 'a', 'a*']) - } - { - // UPDATE c0, c0 -> a0 - // NOCHGE b0 and a1 - const { baseStore, deriStore } = makeStores() - baseStore.set(c, '*') - expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) - expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) - } - { - // UPDATE d0, d0 -> c0 -> a0 - // NOCHGE b0 and a1 - const { baseStore, deriStore } = makeStores() - baseStore.set(d, '*') - expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) - expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) - } - { - // UPDATE e0, e0 -> d0 -> c0 -> a0 - // └--------------> b0 - // NOCHGE a1 - const { baseStore, deriStore } = makeStores() - baseStore.set(e, '*', '*') - expect(getAtoms(baseStore)).toEqual(['*', '*', '*', '*', '**']) - expect(getAtoms(deriStore)).toEqual(['*', '*', '*', 'a', 'a*']) - } - { - // UPDATE a0 - // NOCHGE b0 and a1 - const { baseStore, deriStore } = makeStores() - deriStore.set(a, '*') - expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) - expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) - } - { - // UPDATE b0 - // NOCHGE a0 and a1 - const { baseStore, deriStore } = makeStores() - deriStore.set(b, '*') - expect(getAtoms(baseStore)).toEqual(['a', '*', 'a', 'a', 'a*']) - expect(getAtoms(deriStore)).toEqual(['a', '*', 'a', 'a', 'a*']) - } - { - // UPDATE c0, c0 -> a0 - // NOCHGE b0 and a1 - const { baseStore, deriStore } = makeStores() - deriStore.set(c, '*') - expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) - expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) - } - { - // UPDATE d1, d1 -> c1 -> a1 - // NOCHGE b0 and a0 - const { baseStore, deriStore } = makeStores() - deriStore.set(d, '*') - expect(getAtoms(baseStore)).toEqual(['a', 'b', 'a', 'a', 'ab']) - expect(getAtoms(deriStore)).toEqual(['a', 'b', 'a', '*', '*b']) - } - { - // UPDATE e0, e0 -> d1 -> c1 -> a1 - // └--------------> b0 - // NOCHGE a0 - const { baseStore, deriStore } = makeStores() - deriStore.set(e, '*', '*') - expect(getAtoms(baseStore)).toEqual(['a', '*', 'a', 'a', 'a*']) - expect(getAtoms(deriStore)).toEqual(['a', '*', 'a', '*', '**']) - } - }) -}) diff --git a/tests/vanilla/unstable_derive.test.tsx b/tests/vanilla/unstable_derive.test.tsx new file mode 100644 index 0000000000..03e8d73daf --- /dev/null +++ b/tests/vanilla/unstable_derive.test.tsx @@ -0,0 +1,345 @@ +import { describe, expect, it, vi } from 'vitest' +import { atom, createStore } from 'jotai/vanilla' +import type { Atom, ExtractAtomValue } from 'jotai/vanilla' + +describe('unstable_derive for scoping atoms', () => { + /** + * a + * S1[a]: a1 + */ + it('primitive atom', async () => { + const a = atom('a') + a.onMount = (setSelf) => setSelf((v) => v + ':mounted') + const scopedAtoms = new Set>([a]) + + const store = createStore() + const derivedStore = store.unstable_derive((getAtomState) => { + const scopedAtomStateMap = new WeakMap() + return [ + (atom, originAtomState) => { + if (scopedAtoms.has(atom)) { + let atomState = scopedAtomStateMap.get(atom) + if (!atomState) { + atomState = { d: new Map(), p: new Set(), n: 0 } + scopedAtomStateMap.set(atom, atomState) + } + return atomState + } + return getAtomState(atom, originAtomState) + }, + ] + }) + + expect(store.get(a)).toBe('a') + expect(derivedStore.get(a)).toBe('a') + + derivedStore.sub(a, vi.fn()) + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(a)).toBe('a') + expect(derivedStore.get(a)).toBe('a:mounted') + + derivedStore.set(a, (v) => v + ':updated') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(a)).toBe('a') + expect(derivedStore.get(a)).toBe('a:mounted:updated') + }) + + /** + * a, b, c(a + b) + * S1[a]: a1, b0, c0(a1 + b0) + */ + it('derived atom (scoping primitive)', async () => { + const a = atom('a') + const b = atom('b') + const c = atom((get) => get(a) + get(b)) + const scopedAtoms = new Set>([a]) + + const store = createStore() + const derivedStore = store.unstable_derive((getAtomState) => { + type AnyAtom = Parameters[0] + type AtomState = ReturnType + const scopedAtomStateMap = new WeakMap() + const scopedInvertAtomStateMap = new WeakMap() + const copyMounted = ( + mounted: NonNullable, + ): NonNullable => ({ + ...mounted, + d: new Set(mounted.d), + t: new Set(mounted.t), + }) + const copyAtomState = (atomState: AtomState): AtomState => ({ + ...atomState, + d: new Map(atomState.d), + p: new Set(atomState.p), + ...('m' in atomState ? { m: copyMounted(atomState.m) } : {}), + }) + return [ + (atom, originAtomState) => { + type TheAtomState = ReturnType< + typeof getAtomState> + > + let atomState = scopedAtomStateMap.get(atom) + if (atomState) { + return atomState as TheAtomState + } + if ( + scopedInvertAtomStateMap.has(originAtomState as never) || + scopedAtoms.has(atom) + ) { + atomState = { d: new Map(), p: new Set(), n: 0 } + scopedAtomStateMap.set(atom, atomState) + scopedInvertAtomStateMap.set(atomState, atom) + return atomState as TheAtomState + } + const originalAtomState = getAtomState(atom, originAtomState) + if ( + Array.from(originalAtomState.d).some(([a]) => + scopedAtomStateMap.has(a), + ) + ) { + atomState = copyAtomState(originalAtomState) + scopedAtomStateMap.set(atom, atomState) + scopedInvertAtomStateMap.set(atomState, atom) + return atomState as TheAtomState + } + return originalAtomState + }, + ] + }) + + expect(store.get(c)).toBe('ab') + expect(derivedStore.get(c)).toBe('ab') + + derivedStore.set(a, 'a2') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(c)).toBe('ab') + expect(derivedStore.get(c)).toBe('a2b') + + derivedStore.sub(c, vi.fn()) + derivedStore.set(b, 'b2') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(c)).toBe('ab2') + expect(derivedStore.get(c)).toBe('a2b2') + }) + + /** + * a, b(a) + * S1[b]: a0, b1(a1) + */ + it('derived atom (scoping derived)', async () => { + const a = atom('a') + const b = atom( + (get) => get(a), + (_get, set, v: string) => { + set(a, v) + }, + ) + const scopedAtoms = new Set>([b]) + + const store = createStore() + const derivedStore = store.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) + }, + ] + }) + + expect(store.get(a)).toBe('a') + expect(store.get(b)).toBe('a') + expect(derivedStore.get(a)).toBe('a') + expect(derivedStore.get(b)).toBe('a') + + store.set(a, 'a2') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(a)).toBe('a2') + expect(store.get(b)).toBe('a2') + expect(derivedStore.get(a)).toBe('a2') + expect(derivedStore.get(b)).toBe('a') + + store.set(b, 'a3') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(a)).toBe('a3') + expect(store.get(b)).toBe('a3') + expect(derivedStore.get(a)).toBe('a3') + expect(derivedStore.get(b)).toBe('a') + + derivedStore.set(a, 'a4') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(a)).toBe('a4') + expect(store.get(b)).toBe('a4') + expect(derivedStore.get(a)).toBe('a4') + expect(derivedStore.get(b)).toBe('a') + + derivedStore.set(b, 'a5') + await new Promise((resolve) => setTimeout(resolve)) + expect(store.get(a)).toBe('a4') + expect(store.get(b)).toBe('a4') + expect(derivedStore.get(a)).toBe('a4') + expect(derivedStore.get(b)).toBe('a5') + }) + + /** + * a, b, c(a), d(c), e(d + b) + * S1[d]: a0, b0, c0(a0), d1(c1(a1)), e0(d1(c1(a1)) + b0) + */ + it('derived atom (scoping derived chain)', async () => { + const a = atom('a') + const b = atom('b') + const c = atom( + (get) => get(a), + (_get, set, v: string) => set(a, v), + ) + const d = atom( + (get) => get(c), + (_get, set, v: string) => set(c, v), + ) + const e = atom( + (get) => get(d) + get(b), + (_get, set, av: string, bv: string) => { + set(d, av) + set(b, bv) + }, + ) + const scopedAtoms = new Set>([d]) + + function makeStores() { + const baseStore = createStore() + const deriStore = 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) + }, + ] + }) + expect(getAtoms(baseStore)).toEqual(['a', 'b', 'a', 'a', 'ab']) + expect(getAtoms(deriStore)).toEqual(['a', 'b', 'a', 'a', 'ab']) + return { baseStore, deriStore } + } + type Store = ReturnType + function getAtoms(store: Store) { + return [ + store.get(a), + store.get(b), + store.get(c), + store.get(d), + store.get(e), + ] + } + + /** + * base[d]: a0, b0, c0(a0), d0(c0(a0)), e0(d0(c0(a0)) + b0) + * deri[d]: a0, b0, c0(a0), d1(c1(a1)), e0(d1(c1(a1)) + b0) + */ + { + // UPDATE a0 + // NOCHGE b0 and a1 + const { baseStore, deriStore } = makeStores() + baseStore.set(a, '*') + expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) + expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) + } + { + // UPDATE b0 + // NOCHGE a0 and a1 + const { baseStore, deriStore } = makeStores() + baseStore.set(b, '*') + expect(getAtoms(baseStore)).toEqual(['a', '*', 'a', 'a', 'a*']) + expect(getAtoms(deriStore)).toEqual(['a', '*', 'a', 'a', 'a*']) + } + { + // UPDATE c0, c0 -> a0 + // NOCHGE b0 and a1 + const { baseStore, deriStore } = makeStores() + baseStore.set(c, '*') + expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) + expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) + } + { + // UPDATE d0, d0 -> c0 -> a0 + // NOCHGE b0 and a1 + const { baseStore, deriStore } = makeStores() + baseStore.set(d, '*') + expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) + expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) + } + { + // UPDATE e0, e0 -> d0 -> c0 -> a0 + // └--------------> b0 + // NOCHGE a1 + const { baseStore, deriStore } = makeStores() + baseStore.set(e, '*', '*') + expect(getAtoms(baseStore)).toEqual(['*', '*', '*', '*', '**']) + expect(getAtoms(deriStore)).toEqual(['*', '*', '*', 'a', 'a*']) + } + { + // UPDATE a0 + // NOCHGE b0 and a1 + const { baseStore, deriStore } = makeStores() + deriStore.set(a, '*') + expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) + expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) + } + { + // UPDATE b0 + // NOCHGE a0 and a1 + const { baseStore, deriStore } = makeStores() + deriStore.set(b, '*') + expect(getAtoms(baseStore)).toEqual(['a', '*', 'a', 'a', 'a*']) + expect(getAtoms(deriStore)).toEqual(['a', '*', 'a', 'a', 'a*']) + } + { + // UPDATE c0, c0 -> a0 + // NOCHGE b0 and a1 + const { baseStore, deriStore } = makeStores() + deriStore.set(c, '*') + expect(getAtoms(baseStore)).toEqual(['*', 'b', '*', '*', '*b']) + expect(getAtoms(deriStore)).toEqual(['*', 'b', '*', 'a', 'ab']) + } + { + // UPDATE d1, d1 -> c1 -> a1 + // NOCHGE b0 and a0 + const { baseStore, deriStore } = makeStores() + deriStore.set(d, '*') + expect(getAtoms(baseStore)).toEqual(['a', 'b', 'a', 'a', 'ab']) + expect(getAtoms(deriStore)).toEqual(['a', 'b', 'a', '*', '*b']) + } + { + // UPDATE e0, e0 -> d1 -> c1 -> a1 + // └--------------> b0 + // NOCHGE a0 + const { baseStore, deriStore } = makeStores() + deriStore.set(e, '*', '*') + expect(getAtoms(baseStore)).toEqual(['a', '*', 'a', 'a', 'a*']) + expect(getAtoms(deriStore)).toEqual(['a', '*', 'a', '*', '**']) + } + }) +}) From 39e8ddb38816930153bc7ff683831420cb735418 Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 23 Jul 2024 21:43:11 +0900 Subject: [PATCH 28/31] wip: failed attempt 2 --- tests/vanilla/unstable_derive.test.tsx | 124 +++++++++++++++---------- 1 file changed, 76 insertions(+), 48 deletions(-) diff --git a/tests/vanilla/unstable_derive.test.tsx b/tests/vanilla/unstable_derive.test.tsx index 03e8d73daf..80e1cca6ab 100644 --- a/tests/vanilla/unstable_derive.test.tsx +++ b/tests/vanilla/unstable_derive.test.tsx @@ -55,57 +55,85 @@ describe('unstable_derive for scoping atoms', () => { const scopedAtoms = new Set>([a]) const store = createStore() - const derivedStore = store.unstable_derive((getAtomState) => { - type AnyAtom = Parameters[0] - type AtomState = ReturnType - const scopedAtomStateMap = new WeakMap() - const scopedInvertAtomStateMap = new WeakMap() - const copyMounted = ( - mounted: NonNullable, - ): NonNullable => ({ - ...mounted, - d: new Set(mounted.d), - t: new Set(mounted.t), - }) - const copyAtomState = (atomState: AtomState): AtomState => ({ - ...atomState, - d: new Map(atomState.d), - p: new Set(atomState.p), - ...('m' in atomState ? { m: copyMounted(atomState.m) } : {}), - }) - return [ - (atom, originAtomState) => { - type TheAtomState = ReturnType< - typeof getAtomState> - > - let atomState = scopedAtomStateMap.get(atom) - if (atomState) { - return atomState as TheAtomState + + const scopedAtomStateMap = new WeakMap, unknown>() + const wipAtoms = new Set>() + let propagateScopedAtom: ((atoms: Set>) => void) | undefined + const wrapStore = (origStore: typeof store) => { + const wrapMethod = unknown>( + method: T, + ): T => + ((...args: Parameters) => { + try { + return method(...args) + } finally { + propagateScopedAtom?.(wipAtoms) + wipAtoms.clear() } - if ( - scopedInvertAtomStateMap.has(originAtomState as never) || - scopedAtoms.has(atom) - ) { - atomState = { d: new Map(), p: new Set(), n: 0 } - scopedAtomStateMap.set(atom, atomState) - scopedInvertAtomStateMap.set(atomState, atom) - return atomState as TheAtomState + }) as never + return { + ...origStore, + get: wrapMethod(origStore.get), + set: wrapMethod(origStore.set), + sub: wrapMethod(origStore.sub), + } + } + const derivedStore = wrapStore( + store.unstable_derive((getAtomState) => { + type AtomState = ReturnType + const copyMounted = ( + mounted: NonNullable, + ): NonNullable => ({ + ...mounted, + d: new Set(mounted.d), + t: new Set(mounted.t), + }) + const copyAtomState = (atomState: AtomState): AtomState => ({ + ...atomState, + d: new Map(atomState.d), + p: new Set(atomState.p), + ...('m' in atomState ? { m: copyMounted(atomState.m) } : {}), + }) + propagateScopedAtom ||= (atoms) => { + const nextAtoms = new Set>() + for (const atom of atoms) { + const atomState = scopedAtomStateMap.get(atom) as + | AtomState + | undefined + if (atomState && atomState.m) { + for (const dependent of atomState.m.t) { + if (!scopedAtomStateMap.has(dependent)) { + scopedAtomStateMap.set( + dependent, + copyAtomState(getAtomState(dependent)), + ) + nextAtoms.add(dependent) + } + } + } } - const originalAtomState = getAtomState(atom, originAtomState) - if ( - Array.from(originalAtomState.d).some(([a]) => - scopedAtomStateMap.has(a), - ) - ) { - atomState = copyAtomState(originalAtomState) - scopedAtomStateMap.set(atom, atomState) - scopedInvertAtomStateMap.set(atomState, atom) - return atomState as TheAtomState + if (nextAtoms.size) { + propagateScopedAtom?.(nextAtoms) } - return originalAtomState - }, - ] - }) + } + return [ + (atom, originAtomState) => { + type TheAtomState = ReturnType< + typeof getAtomState> + > + let atomState = scopedAtomStateMap.get(atom) + if (!atomState && scopedAtoms.has(atom)) { + atomState = { d: new Map(), p: new Set(), n: 0 } + scopedAtomStateMap.set(atom, atomState) + } + if (atomState) { + return atomState as TheAtomState + } + return getAtomState(atom, originAtomState) + }, + ] + }), + ) expect(store.get(c)).toBe('ab') expect(derivedStore.get(c)).toBe('ab') From 2e76a9930674cb9c207ba045c2c9fe04bebbb3da Mon Sep 17 00:00:00 2001 From: daishi Date: Mon, 12 Aug 2024 22:30:09 +0900 Subject: [PATCH 29/31] revert failing tests --- tests/vanilla/unstable_derive.test.tsx | 95 ++++---------------------- 1 file changed, 13 insertions(+), 82 deletions(-) diff --git a/tests/vanilla/unstable_derive.test.tsx b/tests/vanilla/unstable_derive.test.tsx index 80e1cca6ab..e7bfa468e6 100644 --- a/tests/vanilla/unstable_derive.test.tsx +++ b/tests/vanilla/unstable_derive.test.tsx @@ -1,6 +1,6 @@ import { describe, expect, it, vi } from 'vitest' import { atom, createStore } from 'jotai/vanilla' -import type { Atom, ExtractAtomValue } from 'jotai/vanilla' +import type { Atom } from 'jotai/vanilla' describe('unstable_derive for scoping atoms', () => { /** @@ -55,85 +55,22 @@ describe('unstable_derive for scoping atoms', () => { const scopedAtoms = new Set>([a]) const store = createStore() - - const scopedAtomStateMap = new WeakMap, unknown>() - const wipAtoms = new Set>() - let propagateScopedAtom: ((atoms: Set>) => void) | undefined - const wrapStore = (origStore: typeof store) => { - const wrapMethod = unknown>( - method: T, - ): T => - ((...args: Parameters) => { - try { - return method(...args) - } finally { - propagateScopedAtom?.(wipAtoms) - wipAtoms.clear() - } - }) as never - return { - ...origStore, - get: wrapMethod(origStore.get), - set: wrapMethod(origStore.set), - sub: wrapMethod(origStore.sub), - } - } - const derivedStore = wrapStore( - store.unstable_derive((getAtomState) => { - type AtomState = ReturnType - const copyMounted = ( - mounted: NonNullable, - ): NonNullable => ({ - ...mounted, - d: new Set(mounted.d), - t: new Set(mounted.t), - }) - const copyAtomState = (atomState: AtomState): AtomState => ({ - ...atomState, - d: new Map(atomState.d), - p: new Set(atomState.p), - ...('m' in atomState ? { m: copyMounted(atomState.m) } : {}), - }) - propagateScopedAtom ||= (atoms) => { - const nextAtoms = new Set>() - for (const atom of atoms) { - const atomState = scopedAtomStateMap.get(atom) as - | AtomState - | undefined - if (atomState && atomState.m) { - for (const dependent of atomState.m.t) { - if (!scopedAtomStateMap.has(dependent)) { - scopedAtomStateMap.set( - dependent, - copyAtomState(getAtomState(dependent)), - ) - nextAtoms.add(dependent) - } - } - } - } - if (nextAtoms.size) { - propagateScopedAtom?.(nextAtoms) - } - } - return [ - (atom, originAtomState) => { - type TheAtomState = ReturnType< - typeof getAtomState> - > + const derivedStore = store.unstable_derive((getAtomState) => { + const scopedAtomStateMap = new WeakMap() + return [ + (atom, originAtomState) => { + if (scopedAtoms.has(atom)) { let atomState = scopedAtomStateMap.get(atom) - if (!atomState && scopedAtoms.has(atom)) { + if (!atomState) { atomState = { d: new Map(), p: new Set(), n: 0 } scopedAtomStateMap.set(atom, atomState) } - if (atomState) { - return atomState as TheAtomState - } - return getAtomState(atom, originAtomState) - }, - ] - }), - ) + return atomState + } + return getAtomState(atom, originAtomState) + }, + ] + }) expect(store.get(c)).toBe('ab') expect(derivedStore.get(c)).toBe('ab') @@ -142,12 +79,6 @@ describe('unstable_derive for scoping atoms', () => { await new Promise((resolve) => setTimeout(resolve)) expect(store.get(c)).toBe('ab') expect(derivedStore.get(c)).toBe('a2b') - - derivedStore.sub(c, vi.fn()) - derivedStore.set(b, 'b2') - await new Promise((resolve) => setTimeout(resolve)) - expect(store.get(c)).toBe('ab2') - expect(derivedStore.get(c)).toBe('a2b2') }) /** From 2b7b67a2f3c8945873f4ae667d673a5a77f075a8 Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 13 Aug 2024 07:54:39 +0900 Subject: [PATCH 30/31] empty commit From 91d24702ea3ec71d722f21a61230fc03292adf5e Mon Sep 17 00:00:00 2001 From: daishi Date: Tue, 13 Aug 2024 09:33:53 +0900 Subject: [PATCH 31/31] restore unstable_is of now (we will remove it in the later version.) --- src/vanilla/atom.ts | 1 + src/vanilla/store.ts | 7 +++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/vanilla/atom.ts b/src/vanilla/atom.ts index af2a68d98f..d0907329ea 100644 --- a/src/vanilla/atom.ts +++ b/src/vanilla/atom.ts @@ -40,6 +40,7 @@ type OnMount = < export interface Atom { toString: () => string read: Read + unstable_is?(a: Atom): boolean debugLabel?: string /** * To ONLY be used by Jotai libraries to mark atoms as private. Subject to change. diff --git a/src/vanilla/store.ts b/src/vanilla/store.ts index 0828750e05..7b04acdbae 100644 --- a/src/vanilla/store.ts +++ b/src/vanilla/store.ts @@ -8,6 +8,9 @@ type OnUnmount = () => void type Getter = Parameters[0] type Setter = Parameters[1] +const isSelfAtom = (atom: AnyAtom, a: AnyAtom): boolean => + atom.unstable_is ? atom.unstable_is(a) : a === atom + const hasInitialValue = >( atom: T, ): atom is T & (T extends Atom ? { init: Value } : never) => @@ -385,7 +388,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { atomState.d.clear() let isSync = true const getter: Getter = (a: Atom) => { - if (a === (atom as AnyAtom)) { + if (isSelfAtom(atom, a)) { const aState = getAtomState(a, atomState) if (!isAtomStateInitialized(aState)) { if (hasInitialValue(a)) { @@ -567,7 +570,7 @@ const buildStore = (getAtomState: StoreArgs[0]): Store => { ) => { const aState = getAtomState(a, atomState) let r: R | undefined - if (a === (atom as AnyAtom)) { + if (isSelfAtom(atom, a)) { if (!hasInitialValue(a)) { // NOTE technically possible but restricted as it may cause bugs throw new Error('atom not writable')