From 914d987f777971cbe29c1343e369751ccaa76638 Mon Sep 17 00:00:00 2001 From: David Maskasky Date: Wed, 18 Sep 2024 19:37:22 -0700 Subject: [PATCH] add atom shims and export atomState type --- __tests__/derive/types.ts | 76 ++----------------- .../derive/understanding/atomState.test.ts | 2 +- .../understanding/unstable_derive.test.ts | 2 +- jotai/vanilla/store.ts | 2 +- src/ScopeProvider2/ScopeProvider2.tsx | 65 +++++++++------- 5 files changed, 47 insertions(+), 100 deletions(-) diff --git a/__tests__/derive/types.ts b/__tests__/derive/types.ts index 31e6e2a..5431188 100644 --- a/__tests__/derive/types.ts +++ b/__tests__/derive/types.ts @@ -1,71 +1,7 @@ -import type { Atom, WritableAtom } from '../../jotai'; +import type { Atom, WritableAtom } from '../../jotai' -export type AnyValue = unknown; -export type AnyError = unknown; -export type AnyAtom = Atom; -export type AnyWritableAtom = WritableAtom; -export type OnUnmount = () => void; - -/** - * Mutable atom state, - * tracked for both mounted and unmounted atoms in a store. - */ -export type AtomState = { - /** - * Map of atoms that the atom depends on. - * The map value is the epoch number of the dependency. - */ - readonly d: Map; - - /** - * Set of atoms with pending promise that depend on the atom. - * - * This may cause memory leaks, but it's for the capability to continue promises - */ - readonly p: Set; - - /** The epoch number of the atom. */ - n: number; - - /** - * Object to store mounted state of the atom. - * - * State tracked for mounted atoms. An atom is considered "mounted" if it has a - * subscriber, or is a transitive dependency of another atom that has a - * subscriber. - * - * The mounted state of an atom is freed once it is no longer mounted. - * - * only available if the atom is mounted - */ - m?: { - /** Set of listeners to notify when the atom value changes. */ - readonly l: Set<() => void>; - - /** Set of mounted atoms that the atom depends on. */ - readonly d: Set; - - /** Set of mounted atoms that depends on the atom. */ - readonly t: Set; - - /** Function to run when the atom is unmounted. */ - u?: OnUnmount; - }; - - /** Atom value */ - v?: Value; - - /** Atom error */ - e?: AnyError; -}; - -export type GetAtomState = ( - atom: Atom, - originAtomState?: AtomState, -) => AtomState; - -// internal & unstable type -export type StoreArgs = readonly [ - getAtomState: GetAtomState, - // possible other arguments in the future -]; +export type AnyValue = unknown +export type AnyError = unknown +export type AnyAtom = Atom +export type AnyWritableAtom = WritableAtom +export type OnUnmount = () => void diff --git a/__tests__/derive/understanding/atomState.test.ts b/__tests__/derive/understanding/atomState.test.ts index a01b18c..88048c5 100644 --- a/__tests__/derive/understanding/atomState.test.ts +++ b/__tests__/derive/understanding/atomState.test.ts @@ -1,10 +1,10 @@ +import type { AtomState } from '../../../jotai/vanilla/store' import { atom, createStore } from '../../../jotai' import { assertIsDevStore } from '../../utils' const store = createStore() assertIsDevStore(store) const stateMap = store.dev4_get_internal_weak_map() -type AtomState = NonNullable> const atomA = atom(0) atomA.debugLabel = 'atomA' diff --git a/__tests__/derive/understanding/unstable_derive.test.ts b/__tests__/derive/understanding/unstable_derive.test.ts index edefd85..e24da06 100644 --- a/__tests__/derive/understanding/unstable_derive.test.ts +++ b/__tests__/derive/understanding/unstable_derive.test.ts @@ -1,3 +1,4 @@ +import type { AtomState } from '../../../jotai/vanilla/store' import { atom, createStore, SetStateAction } from '../../../jotai' import { assertIsDevStore } from '../../utils' @@ -15,7 +16,6 @@ const deriveCallback: DeriveCallack = jest.fn((baseGetAtomState) => { const store = createStore().unstable_derive(deriveCallback) assertIsDevStore(store) const stateMap = store.dev4_get_internal_weak_map() -type AtomState = NonNullable> type AtomStateWithLabel = AtomState & { label?: string } const atomA = atom(0) diff --git a/jotai/vanilla/store.ts b/jotai/vanilla/store.ts index 9e5a8d1..982bbc6 100644 --- a/jotai/vanilla/store.ts +++ b/jotai/vanilla/store.ts @@ -85,7 +85,7 @@ type Mounted = { * Mutable atom state, * tracked for both mounted and unmounted atoms in a store. */ -type AtomState = { +export type AtomState = { /** * Map of atoms that the atom depends on. * The map value is the epoch number of the dependency. diff --git a/src/ScopeProvider2/ScopeProvider2.tsx b/src/ScopeProvider2/ScopeProvider2.tsx index 67abe8f..69d056e 100644 --- a/src/ScopeProvider2/ScopeProvider2.tsx +++ b/src/ScopeProvider2/ScopeProvider2.tsx @@ -7,6 +7,7 @@ import { type WritableAtom, } from '../../jotai/vanilla' import type { AtomFamily } from '../../jotai/vanilla/utils/atomFamily' +import type { AtomState } from '../../jotai/vanilla/store' import { MapProxy } from './stateProxy' type Store = ReturnType @@ -97,10 +98,10 @@ export function createScopedStore( // -------------------------------------------- /** map of scoped atoms to their atomState states */ - const scopedAtomStateMap = new WeakMap() + const scopedAtomStateMap = new WeakMap>() /** set of scoped atom states */ - const scopedAtomStateSet = new WeakSet() + const scopedAtomStateSet = new WeakSet>() /** set of computed atoms already proxied */ const proxiedAtomStateSet = new WeakSet() @@ -144,34 +145,46 @@ export function createScopedStore( atomState.l.add(() => {}) atomState.s ??= new Set() } - const derivedStore: NamedStore = baseStore.unstable_derive((getAtomState) => { - return [ - (atom, originAtomState) => { - if (isComputedAtom(atom)) { - emplace(atom, proxiedAtomStateSet, () => - proxyAtomState(atom, getAtomState(atom, originAtomState)), - ) - if (computedConsumer.has(atom)) { + const derivedStore: NamedStore = baseStore.unstable_derive( + (baseGetAtomState) => { + return [ + function getAtomState(atom, originAtomState) { + if (isComputedAtom(atom)) { + emplace(atom, proxiedAtomStateSet, () => + proxyAtomState(atom, baseGetAtomState(atom, originAtomState)), + ) + if (computedConsumer.has(atom)) { + emplace(atom, scopedAtomStateMap, () => { + return cloneAtomState(baseGetAtomState(atom))! + }) + return scopedAtomStateMap.get(atom)! + } + } + + if (scopedAtomStateSet.has(originAtomState!) || explicit.has(atom)) { emplace(atom, scopedAtomStateMap, () => { - return cloneAtomState(getAtomState(atom)) + const atomState: AtomState = { + d: new Map(), + p: new Set(), + n: 0, + } + scopedAtomStateSet.add(atomState) + return atomState }) return scopedAtomStateMap.get(atom)! } - } - if (scopedAtomStateSet.has(originAtomState!) || explicit.has(atom)) { - emplace(atom, scopedAtomStateMap, () => { - const atomState: AtomState = { d: new Map(), p: new Set(), n: 0 } - scopedAtomStateSet.add(atomState) - return atomState - }) - return scopedAtomStateMap.get(atom) - } - - return getAtomState(atom, originAtomState) - }, - ] - }) + return baseGetAtomState(atom, originAtomState)! + }, + function atomRead(atom, getter, options) { + return atom.read(getter, options) + }, + function atomWrite(atom, getter, setter, ...args) { + return atom.write(getter, setter, ...args) + }, + ] + }, + ) if (debugName) { derivedStore.name = debugName } @@ -188,8 +201,6 @@ function isWritableAtom(anAtom: AnyAtom): anAtom is AnyWritableAtom { const { read: defaultRead, write: defaultWrite } = atom(null) -type GetAtomState = Parameters[0]>[0] -type AtomState = ReturnType type Mutable = { -readonly [P in keyof T]: T[P] }