Skip to content

Commit

Permalink
add atom shims and export atomState type
Browse files Browse the repository at this point in the history
  • Loading branch information
David Maskasky committed Sep 24, 2024
1 parent 7008857 commit 914d987
Show file tree
Hide file tree
Showing 5 changed files with 47 additions and 100 deletions.
76 changes: 6 additions & 70 deletions __tests__/derive/types.ts
Original file line number Diff line number Diff line change
@@ -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<AnyValue>;
export type AnyWritableAtom = WritableAtom<AnyValue, unknown[], unknown>;
export type OnUnmount = () => void;

/**
* Mutable atom state,
* tracked for both mounted and unmounted atoms in a store.
*/
export type AtomState<Value = AnyValue> = {
/**
* Map of atoms that the atom depends on.
* The map value is the epoch number of the dependency.
*/
readonly d: Map<AnyAtom, number>;

/**
* 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<AnyAtom>;

/** 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<AnyAtom>;

/** Set of mounted atoms that depends on the atom. */
readonly t: Set<AnyAtom>;

/** Function to run when the atom is unmounted. */
u?: OnUnmount;
};

/** Atom value */
v?: Value;

/** Atom error */
e?: AnyError;
};

export type GetAtomState = <Value>(
atom: Atom<Value>,
originAtomState?: AtomState,
) => AtomState<Value>;

// 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<AnyValue>
export type AnyWritableAtom = WritableAtom<AnyValue, unknown[], unknown>
export type OnUnmount = () => void
2 changes: 1 addition & 1 deletion __tests__/derive/understanding/atomState.test.ts
Original file line number Diff line number Diff line change
@@ -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<ReturnType<typeof stateMap.get>>

const atomA = atom(0)
atomA.debugLabel = 'atomA'
Expand Down
2 changes: 1 addition & 1 deletion __tests__/derive/understanding/unstable_derive.test.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import type { AtomState } from '../../../jotai/vanilla/store'
import { atom, createStore, SetStateAction } from '../../../jotai'
import { assertIsDevStore } from '../../utils'

Expand All @@ -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<ReturnType<typeof stateMap.get>>
type AtomStateWithLabel = AtomState & { label?: string }

const atomA = atom(0)
Expand Down
2 changes: 1 addition & 1 deletion jotai/vanilla/store.ts
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ type Mounted = {
* Mutable atom state,
* tracked for both mounted and unmounted atoms in a store.
*/
type AtomState<Value = AnyValue> = {
export type AtomState<Value = AnyValue> = {
/**
* Map of atoms that the atom depends on.
* The map value is the epoch number of the dependency.
Expand Down
65 changes: 38 additions & 27 deletions src/ScopeProvider2/ScopeProvider2.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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<typeof getDefaultStore>
Expand Down Expand Up @@ -97,10 +98,10 @@ export function createScopedStore(
// --------------------------------------------

/** map of scoped atoms to their atomState states */
const scopedAtomStateMap = new WeakMap<AnyAtom, AtomState>()
const scopedAtomStateMap = new WeakMap<AnyAtom, AtomState<any>>()

/** set of scoped atom states */
const scopedAtomStateSet = new WeakSet<AtomState>()
const scopedAtomStateSet = new WeakSet<AtomState<any>>()

/** set of computed atoms already proxied */
const proxiedAtomStateSet = new WeakSet<AnyAtom>()
Expand Down Expand Up @@ -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<any> = {
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
}
Expand All @@ -188,8 +201,6 @@ function isWritableAtom(anAtom: AnyAtom): anAtom is AnyWritableAtom {

const { read: defaultRead, write: defaultWrite } = atom(null)

type GetAtomState = Parameters<Parameters<Store['unstable_derive']>[0]>[0]
type AtomState = ReturnType<GetAtomState>
type Mutable<T> = {
-readonly [P in keyof T]: T[P]
}
Expand Down

0 comments on commit 914d987

Please sign in to comment.