Skip to content

Commit

Permalink
add batch to internal callback
Browse files Browse the repository at this point in the history
  • Loading branch information
dmaskasky committed Dec 21, 2024
1 parent 4ebc654 commit ada9115
Show file tree
Hide file tree
Showing 6 changed files with 178 additions and 74 deletions.
20 changes: 11 additions & 9 deletions src/vanilla/store.ts
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,8 @@ const isPromiseLike = (
): x is PromiseLike<unknown> & { onCancel?: (fn: CancelHandler) => void } =>
typeof (x as any)?.then === 'function'

type BatchListener = (batch: Batch) => void

/**
* 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
Expand All @@ -77,7 +79,7 @@ type Mounted = {
/** Set of mounted atoms that depends on the atom. */
readonly t: Set<AnyAtom>
/** Function to run when the atom is unmounted. */
u?: (batch: Batch) => void
u?: BatchListener
}

/**
Expand All @@ -91,7 +93,7 @@ export type AtomState<Value = AnyValue> = {
*/
readonly d: Map<AnyAtom, number>
/** Set of priority listeners to run when the atom value changes. */
l?: Set<readonly [listener: () => void, priority?: BatchPriority]>
l?: Set<readonly [listener: BatchListener, priority?: BatchPriority]>
/**
* Set of atoms with pending promise that depend on the atom.
*
Expand Down Expand Up @@ -171,11 +173,11 @@ type Batch = Readonly<{
/** Atom dependents map */
D: Map<AnyAtom, Set<AnyAtom>>
/** High priority functions */
H: Set<() => void>
H: Set<BatchListener>
/** Medium priority functions */
M: Set<() => void>
M: Set<BatchListener>
/** Low priority functions */
L: Set<() => void>
L: Set<BatchListener>
}>

const createBatch = (): Batch => ({
Expand All @@ -187,7 +189,7 @@ const createBatch = (): Batch => ({

const addBatchFunc = (
batch: Batch,
fn: () => void,
fn: BatchListener,
priority: BatchPriority,
) => {
batch[priority].add(fn)
Expand All @@ -205,7 +207,7 @@ const registerBatchAtom = (
addBatchFunc(batch, listener, priority)
}
for (const listener of atomState.m?.l || []) {
addBatchFunc(batch, listener, 'M')
addBatchFunc(batch, () => listener(), 'M')
}
}
addBatchFunc(batch, scheduleListeners, 'H')
Expand All @@ -229,9 +231,9 @@ const getBatchAtomDependents = (batch: Batch, atom: AnyAtom) =>
const flushBatch = (batch: Batch) => {
let error: AnyError
let hasError = false
const call = (fn: () => void) => {
const call = (fn: BatchListener) => {
try {
fn()
fn(batch)
} catch (e) {
if (!hasError) {
error = e
Expand Down
14 changes: 14 additions & 0 deletions tests/setup.ts
Original file line number Diff line number Diff line change
@@ -1 +1,15 @@
import '@testing-library/jest-dom/vitest'
import { expect } from 'vitest'

expect.extend({
toHaveBeenCalledBefore(received, expected) {

Check failure on line 5 in tests/setup.ts

View workflow job for this annotation

GitHub Actions / test_matrix (3.9.7)

Parameter 'received' implicitly has an 'any' type.

Check failure on line 5 in tests/setup.ts

View workflow job for this annotation

GitHub Actions / test_matrix (3.9.7)

Parameter 'expected' implicitly has an 'any' type.

Check failure on line 5 in tests/setup.ts

View workflow job for this annotation

GitHub Actions / test_matrix (4.0.5)

Parameter 'received' implicitly has an 'any' type.

Check failure on line 5 in tests/setup.ts

View workflow job for this annotation

GitHub Actions / test_matrix (4.0.5)

Parameter 'expected' implicitly has an 'any' type.

Check failure on line 5 in tests/setup.ts

View workflow job for this annotation

GitHub Actions / test_matrix (3.8.3)

Parameter 'received' implicitly has an 'any' type.

Check failure on line 5 in tests/setup.ts

View workflow job for this annotation

GitHub Actions / test_matrix (3.8.3)

Parameter 'expected' implicitly has an 'any' type.
const pass =
received.mock.invocationCallOrder[0] <
expected.mock.invocationCallOrder[0]
return {
pass,
message: () =>
`expected ${received} to have been called before ${expected}`,
}
},
})
51 changes: 30 additions & 21 deletions tests/vanilla/effect.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,13 @@ type Ref = {
epoch: number
}

type INTERNAL_onInit = NonNullable<AnyAtom['INTERNAL_onInit']>
type AtomState = Parameters<INTERNAL_onInit>[1]
type BatchListeners = NonNullable<AtomState['l']>
type BatchEntry = BatchListeners extends Set<infer U> ? U : never
type BatchListener = BatchEntry[0]
type BatchPriority = NonNullable<BatchEntry[1]>

function atomSyncEffect(effect: Effect) {
const refAtom = atom(
() => ({ deps: new Set(), inProgress: 0, epoch: 0 }) as Ref,
Expand Down Expand Up @@ -58,31 +65,32 @@ function atomSyncEffect(effect: Effect) {
)
refAtom.onMount = (mount) => mount()
const refreshAtom = atom(0)
const internalAtom = atom((get) => {
get(refreshAtom)
const ref = get(refAtom)
if (!ref.get) {
ref.get = ((a) => {
ref.deps.add(a)
return get(a)
}) as Getter & { peak: Getter }
}
ref.deps.forEach(get)
ref.isPending = true
return ++ref.epoch
})
const bridgeAtom = atom(
(get) => get(internalAtom),
const internalAtom = atom(
(get) => {
get(refreshAtom)
const ref = get(refAtom)
if (!ref.get) {
ref.get = ((a) => {
ref.deps.add(a)
return get(a)
}) as Getter & { peak: Getter }
}
ref.deps.forEach(get)
ref.isPending = true
return ++ref.epoch
},
(get, set) => {
set(refreshAtom, (v) => ++v)
return get(refAtom).sub()
},
)
bridgeAtom.onMount = (mount) => mount()
bridgeAtom.INTERNAL_onInit = (store, atomState) => {
atomState.l ||= new Set()
internalAtom.onMount = (mount) => mount()
internalAtom.INTERNAL_onInit = (store, atomState) => {
if (!('l' in atomState)) {
atomState.l = new Set()
}
store.get(refAtom).sub = function subscribe() {
function listener() {
const batchListener: BatchListener = (_batch) => {
const ref = store.get(refAtom)
if (!ref.isPending || ref.inProgress > 0) {
return
Expand All @@ -102,13 +110,14 @@ function atomSyncEffect(effect: Effect) {
}
: null
}
const entry = [listener, 'H'] as const
const priority: BatchPriority = 'H'
const entry = [batchListener, priority] as const
atomState.l!.add(entry)
return () => atomState.l!.delete(entry)
}
}
const effectAtom = Object.assign(
atom((get) => void get(bridgeAtom)),
atom((get) => void get(internalAtom)),
{ effect },
)
return effectAtom
Expand Down
52 changes: 50 additions & 2 deletions tests/vanilla/store.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -1034,13 +1034,13 @@ it('should call onInit only once per store', () => {
const a = atom(0)
type AtomState = Parameters<NonNullable<Atom<unknown>['INTERNAL_onInit']>>[1]
let aAtomState: AtomState
const aOnInit = vi.fn((_store, atomState) => {
const aOnInit = vi.fn((_store: Store, atomState: AtomState) => {
aAtomState = atomState
})
a.INTERNAL_onInit = aOnInit
const b = atom(0)
let bAtomState: AtomState
const bOnInit = vi.fn((_store, atomState) => {
const bOnInit = vi.fn((_store: Store, atomState: AtomState) => {
bAtomState = atomState
})
b.INTERNAL_onInit = bOnInit
Expand Down Expand Up @@ -1086,3 +1086,51 @@ it('should call onInit only once per store', () => {
) as Store,
)
})

it('should pass store and atomState to the atom initializer', () => {
expect.assertions(2)
const store = createStore()
const a = atom(null)
a.INTERNAL_onInit = (store, atomState) => {
expect(store).toBe(store)
expect(atomState).toEqual(expect.objectContaining({}))
}
store.get(a)
})

it('should call the batch listener with batch and respect the priority', () => {
type INTERNAL_onInit = NonNullable<Atom<unknown>['INTERNAL_onInit']>
type AtomState = Parameters<INTERNAL_onInit>[1]
type BatchListeners = NonNullable<AtomState['l']>
type BatchEntry = BatchListeners extends Set<infer U> ? U : never
type BatchListener = BatchEntry[0]

const a = atom(0)
const highPriorityBatchListener = vi.fn() as BatchListener
const mediumPriorityBatchListener = vi.fn() as BatchListener
const defaultPriorityBatchListener = vi.fn() as BatchListener // medium
const lowPriorityBatchListener = vi.fn() as BatchListener
a.INTERNAL_onInit = (_store, atomState) => {
atomState.l = new Set([

Check failure on line 1114 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.7.4)

Type 'Set<[BatchListener, "L"] | [BatchListener] | [BatchListener, "H"] | [BatchListener, "M"]>' is not assignable to type 'Set<readonly [BatchListener, BatchPriority]>'.

Check failure on line 1114 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.5.5)

Type 'Set<[BatchListener, "L"] | [BatchListener] | [BatchListener, "H"] | [BatchListener, "M"]>' is not assignable to type 'Set<readonly [BatchListener, BatchPriority]>'.

Check failure on line 1114 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.6.4)

Type 'Set<[BatchListener, "L"] | [BatchListener] | [BatchListener, "H"] | [BatchListener, "M"]>' is not assignable to type 'Set<readonly [BatchListener, BatchPriority]>'.

Check failure on line 1114 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.4.4)

Type 'Set<[BatchListener, "L"] | [BatchListener] | [BatchListener, "H"] | [BatchListener, "M"]>' is not assignable to type 'Set<readonly [BatchListener, BatchPriority]>'.

Check failure on line 1114 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.3.5)

Type 'Set<[BatchListener, "L"] | [BatchListener] | [BatchListener, "H"] | [BatchListener, "M"]>' is not assignable to type 'Set<readonly [BatchListener, BatchPriority]>'.
[lowPriorityBatchListener, 'L'],

Check failure on line 1115 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (3.9.7)

Type 'Set<[BatchListener, "L"] | [BatchListener] | [BatchListener, "H"] | [BatchListener, "M"]>' is not assignable to type 'Set<readonly [BatchListener, BatchPriority]>'.

Check failure on line 1115 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.0.5)

Type 'Set<[BatchListener, "L"] | [BatchListener] | [BatchListener, "H"] | [BatchListener, "M"]>' is not assignable to type 'Set<readonly [BatchListener, BatchPriority]>'.

Check failure on line 1115 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (3.8.3)

Type 'Set<[BatchListener, "L"] | [BatchListener] | [BatchListener, "H"] | [BatchListener, "M"]>' is not assignable to type 'Set<readonly [BatchListener, BatchPriority]>'.

Check failure on line 1115 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.2.3)

Type 'Set<[BatchListener, "L"] | [BatchListener] | [BatchListener, "H"] | [BatchListener, "M"]>' is not assignable to type 'Set<readonly [BatchListener, BatchPriority]>'.

Check failure on line 1115 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.1.5)

Type 'Set<[BatchListener, "L"] | [BatchListener] | [BatchListener, "H"] | [BatchListener, "M"]>' is not assignable to type 'Set<readonly [BatchListener, BatchPriority]>'.
[defaultPriorityBatchListener],
[highPriorityBatchListener, 'H'],
[mediumPriorityBatchListener, 'M'],
])
}
const store = createStore()
store.set(a, 1)
const mockBatch = expect.objectContaining({})
expect(highPriorityBatchListener).toHaveBeenCalledWith(mockBatch)
expect(mediumPriorityBatchListener).toHaveBeenCalledWith(mockBatch)
expect(lowPriorityBatchListener).toHaveBeenCalledWith(mockBatch)
expect(highPriorityBatchListener).toHaveBeenCalledBefore(

Check failure on line 1127 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.7.4)

Property 'toHaveBeenCalledBefore' does not exist on type 'Assertion<BatchListener, any>'. Did you mean 'toHaveBeenCalledOnce'?

Check failure on line 1127 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.5.5)

Property 'toHaveBeenCalledBefore' does not exist on type 'Assertion<BatchListener, any>'. Did you mean 'toHaveBeenCalledOnce'?

Check failure on line 1127 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.8.4)

Property 'toHaveBeenCalledBefore' does not exist on type 'Assertion<BatchListener, any>'. Did you mean 'toHaveBeenCalled'?

Check failure on line 1127 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.6.4)

Property 'toHaveBeenCalledBefore' does not exist on type 'Assertion<BatchListener, any>'. Did you mean 'toHaveBeenCalledOnce'?
mediumPriorityBatchListener,

Check failure on line 1128 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.4.4)

Argument of type 'BatchListener' is not assignable to parameter of type 'SpyInstance<any[], any>'.

Check failure on line 1128 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.3.5)

Argument of type 'BatchListener' is not assignable to parameter of type 'SpyInstance<any[], any>'.
)

Check failure on line 1129 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.2.3)

Argument of type 'BatchListener' is not assignable to parameter of type 'SpyInstance<any[], any>'.

Check failure on line 1129 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.1.5)

Argument of type 'BatchListener' is not assignable to parameter of type 'SpyInstance<any[], any>'.
expect(mediumPriorityBatchListener).toHaveBeenCalledBefore(

Check failure on line 1130 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.7.4)

Property 'toHaveBeenCalledBefore' does not exist on type 'Assertion<BatchListener, any>'. Did you mean 'toHaveBeenCalledOnce'?

Check failure on line 1130 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.5.5)

Property 'toHaveBeenCalledBefore' does not exist on type 'Assertion<BatchListener, any>'. Did you mean 'toHaveBeenCalledOnce'?

Check failure on line 1130 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.8.4)

Property 'toHaveBeenCalledBefore' does not exist on type 'Assertion<BatchListener, any>'. Did you mean 'toHaveBeenCalled'?

Check failure on line 1130 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.6.4)

Property 'toHaveBeenCalledBefore' does not exist on type 'Assertion<BatchListener, any>'. Did you mean 'toHaveBeenCalledOnce'?
lowPriorityBatchListener,

Check failure on line 1131 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.4.4)

Argument of type 'BatchListener' is not assignable to parameter of type 'SpyInstance<any[], any>'.

Check failure on line 1131 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.3.5)

Argument of type 'BatchListener' is not assignable to parameter of type 'SpyInstance<any[], any>'.
)

Check failure on line 1132 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.2.3)

Argument of type 'BatchListener' is not assignable to parameter of type 'SpyInstance<any[], any>'.

Check failure on line 1132 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.1.5)

Argument of type 'BatchListener' is not assignable to parameter of type 'SpyInstance<any[], any>'.
expect(defaultPriorityBatchListener).toHaveBeenCalledBefore(

Check failure on line 1133 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.7.4)

Property 'toHaveBeenCalledBefore' does not exist on type 'Assertion<BatchListener, any>'. Did you mean 'toHaveBeenCalledOnce'?

Check failure on line 1133 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.5.5)

Property 'toHaveBeenCalledBefore' does not exist on type 'Assertion<BatchListener, any>'. Did you mean 'toHaveBeenCalledOnce'?

Check failure on line 1133 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.8.4)

Property 'toHaveBeenCalledBefore' does not exist on type 'Assertion<BatchListener, any>'. Did you mean 'toHaveBeenCalled'?

Check failure on line 1133 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.6.4)

Property 'toHaveBeenCalledBefore' does not exist on type 'Assertion<BatchListener, any>'. Did you mean 'toHaveBeenCalledOnce'?
lowPriorityBatchListener,

Check failure on line 1134 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.4.4)

Argument of type 'BatchListener' is not assignable to parameter of type 'SpyInstance<any[], any>'.

Check failure on line 1134 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.3.5)

Argument of type 'BatchListener' is not assignable to parameter of type 'SpyInstance<any[], any>'.
)

Check failure on line 1135 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.2.3)

Argument of type 'BatchListener' is not assignable to parameter of type 'SpyInstance<any[], any>'.

Check failure on line 1135 in tests/vanilla/store.test.tsx

View workflow job for this annotation

GitHub Actions / test_matrix (4.1.5)

Argument of type 'BatchListener' is not assignable to parameter of type 'SpyInstance<any[], any>'.
})
101 changes: 59 additions & 42 deletions tests/vanilla/unstable_derive.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -13,23 +13,28 @@ describe('unstable_derive for scoping atoms', () => {
const scopedAtoms = new Set<Atom<unknown>>([a])

const store = createStore()
const derivedStore = store.unstable_derive((getAtomState, ...args) => {
const scopedAtomStateMap = new WeakMap()
return [
(atom) => {
if (scopedAtoms.has(atom)) {
let atomState = scopedAtomStateMap.get(atom)
if (!atomState) {
atomState = { d: new Map(), p: new Set(), n: 0 }
scopedAtomStateMap.set(atom, atomState)
const derivedStore = store.unstable_derive(
(getAtomState, atomRead, atomWrite, atomOnMount, atomOnInit) => {
const scopedAtomStateMap = new WeakMap()
return [
(atom) => {
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 atomState
}
return getAtomState(atom)
},
...args,
]
})
return getAtomState(atom)
},
atomRead,
atomWrite,
atomOnMount,
atomOnInit,
]
},
)

expect(store.get(a)).toBe('a')
expect(derivedStore.get(a)).toBe('a')
Expand All @@ -54,23 +59,28 @@ describe('unstable_derive for scoping atoms', () => {
const scopedAtoms = new Set<Atom<unknown>>([a])

const store = createStore()
const derivedStore = store.unstable_derive((getAtomState, ...args) => {
const scopedAtomStateMap = new WeakMap()
return [
(atom) => {
if (scopedAtoms.has(atom)) {
let atomState = scopedAtomStateMap.get(atom)
if (!atomState) {
atomState = { d: new Map(), p: new Set(), n: 0 }
scopedAtomStateMap.set(atom, atomState)
const derivedStore = store.unstable_derive(
(getAtomState, atomRead, atomWrite, atomOnMount, atomOnInit) => {
const scopedAtomStateMap = new WeakMap()
return [
(atom) => {
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 atomState
}
return getAtomState(atom)
},
...args,
]
})
return getAtomState(atom)
},
atomRead,
atomWrite,
atomOnMount,
atomOnInit,
]
},
)

expect(store.get(c)).toBe('ab')
expect(derivedStore.get(c)).toBe('ab')
Expand All @@ -95,7 +105,7 @@ describe('unstable_derive for scoping atoms', () => {
function makeStores() {
const store = createStore()
const derivedStore = store.unstable_derive(
(getAtomState, atomRead, ...args) => {
(getAtomState, atomRead, atomWrite, atomOnMount, atomOnInit) => {
const scopedAtomStateMap = new WeakMap()
return [
(atom) => {
Expand All @@ -118,7 +128,9 @@ describe('unstable_derive for scoping atoms', () => {
}
return atomRead(a, myGet, options)
},
...args,
atomWrite,
atomOnMount,
atomOnInit,
]
},
)
Expand Down Expand Up @@ -173,14 +185,19 @@ describe('unstable_derive for scoping atoms', () => {
it('should pass the correct store instance to the atom initializer', () => {
expect.assertions(2)
const baseStore = createStore()
const derivedStore = baseStore.unstable_derive((getAtomState, ...args) => [
(a, atomOnInit) => {
const atomState = getAtomState(a)
atomOnInit?.(a, atomState)
return atomState
},
...args,
])
const derivedStore = baseStore.unstable_derive(
(getAtomState, atomRead, atomWrite, atomOnMount, atomOnInit) => [
(a, atomOnInit) => {
const atomState = getAtomState(a)
atomOnInit?.(a, atomState)
return atomState
},
atomRead,
atomWrite,
atomOnMount,
atomOnInit,
],
)
const a = atom(null)
a.INTERNAL_onInit = (store) => {
expect(store).toBe(baseStore)
Expand Down
14 changes: 14 additions & 0 deletions tests/vitest-extend.d.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
// vitest-extend.d.ts
import type { SpyInstance } from 'vitest'

// 1. Augment the 'vitest' module:
declare module 'vitest' {
interface Assertion<T = any> {
/**
* Custom matcher that asserts that one mock was called before another.
* Usage:
* expect(mockA).toHaveBeenCalledBefore(mockB)
*/
toHaveBeenCalledBefore(expected: SpyInstance): void
}
}

0 comments on commit ada9115

Please sign in to comment.