Skip to content

Commit

Permalink
initial commit
Browse files Browse the repository at this point in the history
  • Loading branch information
@dmaskasky
dmaskasky committed Jan 4, 2025
1 parent 84dc6b6 commit b903c61
Show file tree
Hide file tree
Showing 6 changed files with 528 additions and 106 deletions.
7 changes: 7 additions & 0 deletions src/vanilla/atom.ts
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import type { INTERNAL_PrdStore as Store } from './store'

type Getter = <Value>(atom: Atom<Value>) => Value

type Setter = <Value, Args extends unknown[], Result>(
Expand Down Expand Up @@ -47,6 +49,11 @@ export interface Atom<Value> {
* @private
*/
debugPrivate?: boolean
/**
* Fires after atom is referenced by the store for the first time
* For advanced use only and subject to change without notice.
*/
unstable_onInit?: (store: Store) => void
}

export interface WritableAtom<Value, Args extends unknown[], Result>
Expand Down
179 changes: 94 additions & 85 deletions src/vanilla/store.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,8 +100,10 @@ type AtomState<Value = AnyValue> = {
n: number
/** Object to store mounted state of the atom. */
m?: Mounted // only available if the atom is mounted
/** Listener to notify when the atom value is updated. */
u?: (batch: Batch) => void
/** Listener to notify when the atom is mounted or unmounted. */
h?: () => void
h?: (batch: Batch) => void
/** Atom value */
v?: Value
/** Atom error */
Expand Down Expand Up @@ -174,28 +176,25 @@ type BatchPriority =
| typeof BATCH_PRIORITY_MEDIUM
| typeof BATCH_PRIORITY_LOW

type Batch = Readonly<{
/** Atom dependents map */
D: Map<AnyAtom, Set<AnyAtom>>
type Batch = Set<() => void>[] & {
/** High priority functions */
[BATCH_PRIORITY_HIGH]: Set<() => void>
/** Medium priority functions */
[BATCH_PRIORITY_MEDIUM]: Set<() => void>
/** Low priority functions */
[BATCH_PRIORITY_LOW]: Set<() => void>
}>
/** Atom dependents map */
D: Map<AnyAtom, Set<AnyAtom>>
}

const createBatch = (): Batch => {
const batch = { D: new Map() } as {
D: Map<AnyAtom, Set<AnyAtom>>
[BATCH_PRIORITY_HIGH]: Set<() => void>
[BATCH_PRIORITY_MEDIUM]: Set<() => void>
[BATCH_PRIORITY_LOW]: Set<() => void>
}
batch[BATCH_PRIORITY_HIGH] = new Set()
batch[BATCH_PRIORITY_MEDIUM] = new Set()
batch[BATCH_PRIORITY_LOW] = new Set()
return batch
const [H, M, L] = [new Set(), new Set(), new Set()]
return Object.assign([H, M, L], {
[BATCH_PRIORITY_HIGH]: H,
[BATCH_PRIORITY_MEDIUM]: M,
[BATCH_PRIORITY_LOW]: L,
D: new Map(),
}) as Batch
}

const addBatchFunc = (
Expand All @@ -213,6 +212,7 @@ const registerBatchAtom = (
) => {
if (!batch.D.has(atom)) {
batch.D.set(atom, new Set())
atomState.u?.(batch)
const scheduleListeners = () => {
atomState.m?.l.forEach((listener) =>
addBatchFunc(batch, BATCH_PRIORITY_MEDIUM, listener),
Expand Down Expand Up @@ -249,14 +249,12 @@ const flushBatch = (batch: Batch) => {
}
}
}
while (batch.H.size || batch.M.size || batch.L.size) {
while (batch.some((channel) => channel.size)) {
batch.D.clear()
batch.H.forEach(call)
batch.H.clear()
batch.M.forEach(call)
batch.M.clear()
batch.L.forEach(call)
batch.L.clear()
for (const channel of batch) {
channel.forEach(call)
channel.clear()
}
}
if (hasError) {
throw error
Expand All @@ -265,7 +263,11 @@ const flushBatch = (batch: Batch) => {

// internal & unstable type
type StoreArgs = readonly [
getAtomState: <Value>(atom: Atom<Value>) => AtomState<Value>,
getAtomState: <Value>(atom: Atom<Value>) => AtomState<Value> | undefined,
setAtomState: <Value>(
atom: Atom<Value>,
atomState: AtomState<Value>,
) => AtomState<Value>,
atomRead: <Value>(
atom: Atom<Value>,
...params: Parameters<Atom<Value>['read']>
Expand All @@ -274,6 +276,7 @@ type StoreArgs = readonly [
atom: WritableAtom<Value, Args, Result>,
...params: Parameters<WritableAtom<Value, Args, Result>['write']>
) => Result,
atomOnInit: <Value>(atom: Atom<Value>, store: Store) => void,
atomOnMount: <Value, Args extends unknown[], Result>(
atom: WritableAtom<Value, Args, Result>,
setAtom: (...args: Args) => Result,
Expand Down Expand Up @@ -302,9 +305,28 @@ type Store = {
export type INTERNAL_DevStoreRev4 = DevStoreRev4
export type INTERNAL_PrdStore = Store

const buildStore = (
...[getAtomState, atomRead, atomWrite, atomOnMount]: StoreArgs
): Store => {
const buildStore = (...storeArgs: StoreArgs): Store => {
const [
getAtomState,
setAtomState,
atomRead,
atomWrite,
atomOnInit,
atomOnMount,
] = storeArgs
const ensureAtomState = <Value>(atom: Atom<Value>) => {
if (import.meta.env?.MODE !== 'production' && !atom) {
throw new Error('Atom is undefined or null')
}
let atomState = getAtomState(atom)
if (!atomState) {
atomState = { d: new Map(), p: new Set(), n: 0 }
atomState = setAtomState(atom, atomState)
atomOnInit?.(atom, store)
}
return atomState
}

const setAtomStateValueOrPromise = (
atom: AnyAtom,
atomState: AtomState,
Expand All @@ -316,7 +338,7 @@ const buildStore = (
if (isPromiseLike(valueOrPromise)) {
patchPromiseForCancelability(valueOrPromise)
for (const a of atomState.d.keys()) {
addPendingPromiseToDependency(atom, valueOrPromise, getAtomState(a))
addPendingPromiseToDependency(atom, valueOrPromise, ensureAtomState(a))
}
atomState.v = valueOrPromise
} else {
Expand All @@ -336,7 +358,7 @@ const buildStore = (
batch: Batch | undefined,
atom: Atom<Value>,
): AtomState<Value> => {
const atomState = getAtomState(atom)
const atomState = ensureAtomState(atom)
// See if we can skip recomputing this atom.
if (isAtomStateInitialized(atomState)) {
// If the atom is mounted, we can use cached atom state.
Expand All @@ -363,7 +385,7 @@ const buildStore = (
let isSync = true
const getter: Getter = <V>(a: Atom<V>) => {
if (isSelfAtom(atom, a)) {
const aState = getAtomState(a)
const aState = ensureAtomState(a)
if (!isAtomStateInitialized(aState)) {
if (hasInitialValue(a)) {
setAtomStateValueOrPromise(a, aState, a.init)
Expand Down Expand Up @@ -454,19 +476,19 @@ const buildStore = (
): Map<AnyAtom, AtomState> => {
const dependents = new Map<AnyAtom, AtomState>()
for (const a of atomState.m?.t || []) {
const aState = getAtomState(a)
const aState = ensureAtomState(a)
if (aState.m) {
dependents.set(a, aState)
}
}
for (const atomWithPendingPromise of atomState.p) {
dependents.set(
atomWithPendingPromise,
getAtomState(atomWithPendingPromise),
ensureAtomState(atomWithPendingPromise),
)
}
getBatchAtomDependents(batch, atom)?.forEach((dependent) => {
dependents.set(dependent, getAtomState(dependent))
dependents.set(dependent, ensureAtomState(dependent))
})
return dependents
}
Expand Down Expand Up @@ -558,7 +580,7 @@ const buildStore = (
a: WritableAtom<V, As, R>,
...args: As
) => {
const aState = getAtomState(a)
const aState = ensureAtomState(a)
try {
if (isSelfAtom(atom, a)) {
if (!hasInitialValue(a)) {
Expand Down Expand Up @@ -610,15 +632,15 @@ const buildStore = (
if (atomState.m && !isPendingPromise(atomState.v)) {
for (const a of atomState.d.keys()) {
if (!atomState.m.d.has(a)) {
const aMounted = mountAtom(batch, a, getAtomState(a))
const aMounted = mountAtom(batch, a, ensureAtomState(a))
aMounted.t.add(atom)
atomState.m.d.add(a)
}
}
for (const a of atomState.m.d || []) {
if (!atomState.d.has(a)) {
atomState.m.d.delete(a)
const aMounted = unmountAtom(batch, a, getAtomState(a))
const aMounted = unmountAtom(batch, a, ensureAtomState(a))
aMounted?.t.delete(atom)
}
}
Expand All @@ -635,7 +657,7 @@ const buildStore = (
readAtomState(batch, atom)
// mount dependencies first
for (const a of atomState.d.keys()) {
const aMounted = mountAtom(batch, a, getAtomState(a))
const aMounted = mountAtom(batch, a, ensureAtomState(a))
aMounted.t.add(atom)
}
// mount self
Expand All @@ -644,7 +666,7 @@ const buildStore = (
d: new Set(atomState.d.keys()),
t: new Set(),
}
atomState.h?.()
atomState.h?.(batch)
if (isActuallyWritableAtom(atom)) {
const mounted = atomState.m
let setAtom: (...args: unknown[]) => unknown
Expand Down Expand Up @@ -687,18 +709,18 @@ const buildStore = (
if (
atomState.m &&
!atomState.m.l.size &&
!Array.from(atomState.m.t).some((a) => getAtomState(a).m?.d.has(atom))
!Array.from(atomState.m.t).some((a) => ensureAtomState(a).m?.d.has(atom))
) {
// unmount self
const onUnmount = atomState.m.u
if (onUnmount) {
addBatchFunc(batch, 'L', () => onUnmount(batch))
addBatchFunc(batch, BATCH_PRIORITY_LOW, () => onUnmount(batch))
}
delete atomState.m
atomState.h?.()
atomState.h?.(batch)
// unmount dependencies
for (const a of atomState.d.keys()) {
const aMounted = unmountAtom(batch, a, getAtomState(a))
const aMounted = unmountAtom(batch, a, ensureAtomState(a))
aMounted?.t.delete(atom)
}
return undefined
Expand All @@ -708,7 +730,7 @@ const buildStore = (

const subscribeAtom = (atom: AnyAtom, listener: () => void) => {
const batch = createBatch()
const atomState = getAtomState(atom)
const atomState = ensureAtomState(atom)
const mounted = mountAtom(batch, atom, atomState)
const listeners = mounted.l
listeners.add(listener)
Expand All @@ -721,8 +743,8 @@ const buildStore = (
}
}

const unstable_derive = (fn: (...args: StoreArgs) => StoreArgs) =>
buildStore(...fn(getAtomState, atomRead, atomWrite, atomOnMount))
const unstable_derive: Store['unstable_derive'] = (fn) =>
buildStore(...fn(...storeArgs))

const store: Store = {
get: readAtom,
Expand All @@ -737,41 +759,37 @@ const deriveDevStoreRev4 = (store: Store): Store & DevStoreRev4 => {
const debugMountedAtoms = new Set<AnyAtom>()
let savedGetAtomState: StoreArgs[0]
let inRestoreAtom = 0
const derivedStore = store.unstable_derive(
(getAtomState, atomRead, atomWrite, atomOnMount) => {
savedGetAtomState = getAtomState
return [
(atom) => {
const atomState = getAtomState(atom)
const originalMounted = atomState.h
atomState.h = () => {
originalMounted?.()
if (atomState.m) {
debugMountedAtoms.add(atom)
} else {
debugMountedAtoms.delete(atom)
}
}
return atomState
},
atomRead,
(atom, getter, setter, ...args) => {
if (inRestoreAtom) {
return setter(atom, ...args)
}
return atomWrite(atom, getter, setter, ...args)
},
atomOnMount,
]
},
)
const derivedStore = store.unstable_derive((...storeArgs: [...StoreArgs]) => {
const [getAtomState, setAtomState, , atomWrite] = storeArgs
savedGetAtomState = getAtomState
storeArgs[1] = (atom, atomState) => {
const newAtomState = setAtomState(atom, atomState)
const originalMounted = newAtomState.h
newAtomState.h = (batch) => {
originalMounted?.(batch)
if (newAtomState.m) {
debugMountedAtoms.add(atom)
} else {
debugMountedAtoms.delete(atom)
}
}
return newAtomState
}
storeArgs[3] = (atom, getter, setter, ...args) => {
if (inRestoreAtom) {
return setter(atom, ...args)
}
return atomWrite(atom, getter, setter, ...args)
}
return storeArgs
})
const savedStoreSet = derivedStore.set
const devStore: DevStoreRev4 = {
// store dev methods (these are tentative and subject to change without notice)
dev4_get_internal_weak_map: () => ({
get: (atom) => {
const atomState = savedGetAtomState(atom)
if (atomState.n === 0) {
if (!atomState || atomState.n === 0) {
// for backward compatibility
return undefined
}
Expand Down Expand Up @@ -805,21 +823,12 @@ type PrdOrDevStore = Store | (Store & DevStoreRev4)

export const createStore = (): PrdOrDevStore => {
const atomStateMap = new WeakMap()
const getAtomState = <Value>(atom: Atom<Value>) => {
if (import.meta.env?.MODE !== 'production' && !atom) {
throw new Error('Atom is undefined or null')
}
let atomState = atomStateMap.get(atom) as AtomState<Value> | undefined
if (!atomState) {
atomState = { d: new Map(), p: new Set(), n: 0 }
atomStateMap.set(atom, atomState)
}
return atomState
}
const store = buildStore(
getAtomState,
(atom) => atomStateMap.get(atom),
(atom, atomState) => atomStateMap.set(atom, atomState).get(atom),
(atom, ...params) => atom.read(...params),
(atom, ...params) => atom.write(...params),
(atom, ...params) => atom.unstable_onInit?.(...params),
(atom, ...params) => atom.onMount?.(...params),
)
if (import.meta.env?.MODE !== 'production') {
Expand Down
Loading

0 comments on commit b903c61

Please sign in to comment.