diff --git a/packages/toolkit/src/entities/sorted_state_adapter.ts b/packages/toolkit/src/entities/sorted_state_adapter.ts index dc7301451..8f16ed294 100644 --- a/packages/toolkit/src/entities/sorted_state_adapter.ts +++ b/packages/toolkit/src/entities/sorted_state_adapter.ts @@ -1,3 +1,4 @@ +import { current, isDraft } from 'immer' import type { IdSelector, Comparer, @@ -12,11 +13,46 @@ import { selectIdValue, ensureEntitiesArray, splitAddedUpdatedEntities, + getCurrent, } from './utils' +// Borrowed from Replay +export function findInsertIndex( + sortedItems: T[], + item: T, + comparisonFunction: Comparer, +): number { + let lowIndex = 0 + let highIndex = sortedItems.length + while (lowIndex < highIndex) { + let middleIndex = (lowIndex + highIndex) >>> 1 + const currentItem = sortedItems[middleIndex] + const res = comparisonFunction(item, currentItem) + if (res >= 0) { + lowIndex = middleIndex + 1 + } else { + highIndex = middleIndex + } + } + + return lowIndex +} + +export function insert( + sortedItems: T[], + item: T, + comparisonFunction: Comparer, +): T[] { + const insertAtIndex = findInsertIndex(sortedItems, item, comparisonFunction) + + sortedItems.splice(insertAtIndex, 0, item) + + return sortedItems +} + export function createSortedStateAdapter( selectId: IdSelector, - sort: Comparer, + comparer: Comparer, ): EntityStateAdapter { type R = DraftableEntityState @@ -30,15 +66,20 @@ export function createSortedStateAdapter( function addManyMutably( newEntities: readonly T[] | Record, state: R, + existingIds?: Id[], ): void { newEntities = ensureEntitiesArray(newEntities) + const existingKeys = new Set( + existingIds ?? (current(state.ids) as Id[]), + ) + const models = newEntities.filter( - (model) => !(selectIdValue(model, selectId) in state.entities), + (model) => !existingKeys.has(selectIdValue(model, selectId)), ) if (models.length !== 0) { - merge(models, state) + mergeFunction(state, models) } } @@ -52,7 +93,10 @@ export function createSortedStateAdapter( ): void { newEntities = ensureEntitiesArray(newEntities) if (newEntities.length !== 0) { - merge(newEntities, state) + for (const item of newEntities) { + delete (state.entities as Record)[selectId(item)] + } + mergeFunction(state, newEntities) } } @@ -64,7 +108,7 @@ export function createSortedStateAdapter( state.entities = {} as Record state.ids = [] - addManyMutably(newEntities, state) + addManyMutably(newEntities, state, []) } function updateOneMutably(update: Update, state: R): void { @@ -76,6 +120,7 @@ export function createSortedStateAdapter( state: R, ): void { let appliedUpdates = false + let replacedIds = false for (let update of updates) { const entity: T | undefined = (state.entities as Record)[update.id] @@ -87,14 +132,20 @@ export function createSortedStateAdapter( Object.assign(entity, update.changes) const newId = selectId(entity) + if (update.id !== newId) { + // We do support the case where updates can change an item's ID. + // This makes things trickier - go ahead and swap the IDs in state now. + replacedIds = true delete (state.entities as Record)[update.id] + const oldIndex = (state.ids as Id[]).indexOf(update.id) + state.ids[oldIndex] = newId ;(state.entities as Record)[newId] = entity } } if (appliedUpdates) { - resortEntities(state) + mergeFunction(state, [], appliedUpdates, replacedIds) } } @@ -106,14 +157,18 @@ export function createSortedStateAdapter( newEntities: readonly T[] | Record, state: R, ): void { - const [added, updated] = splitAddedUpdatedEntities( + const [added, updated, existingIdsArray] = splitAddedUpdatedEntities( newEntities, selectId, state, ) - updateManyMutably(updated, state) - addManyMutably(added, state) + if (updated.length) { + updateManyMutably(updated, state) + } + if (added.length) { + addManyMutably(added, state, existingIdsArray) + } } function areArraysEqual(a: readonly unknown[], b: readonly unknown[]) { @@ -130,27 +185,65 @@ export function createSortedStateAdapter( return true } - function merge(models: readonly T[], state: R): void { + type MergeFunction = ( + state: R, + addedItems: readonly T[], + appliedUpdates?: boolean, + replacedIds?: boolean, + ) => void + + const mergeInsertion: MergeFunction = ( + state, + addedItems, + appliedUpdates, + replacedIds, + ) => { + const currentEntities = getCurrent(state.entities) as Record + const currentIds = getCurrent(state.ids) as Id[] + + const stateEntities = state.entities as Record + + let ids = currentIds + if (replacedIds) { + ids = Array.from(new Set(currentIds)) + } + + let sortedEntities: T[] = [] + for (const id of ids) { + const entity = currentEntities[id] + if (entity) { + sortedEntities.push(entity) + } + } + const wasPreviouslyEmpty = sortedEntities.length === 0 + // Insert/overwrite all new/updated - models.forEach((model) => { - ;(state.entities as Record)[selectId(model)] = model - }) + for (const item of addedItems) { + stateEntities[selectId(item)] = item - resortEntities(state) - } + if (!wasPreviouslyEmpty) { + // Binary search insertion generally requires fewer comparisons + insert(sortedEntities, item, comparer) + } + } - function resortEntities(state: R) { - const allEntities = Object.values(state.entities) as T[] - allEntities.sort(sort) + if (wasPreviouslyEmpty) { + // All we have is the incoming values, sort them + sortedEntities = addedItems.slice().sort(comparer) + } else if (appliedUpdates) { + // We should have a _mostly_-sorted array already + sortedEntities.sort(comparer) + } - const newSortedIds = allEntities.map(selectId) - const { ids } = state + const newSortedIds = sortedEntities.map(selectId) - if (!areArraysEqual(ids, newSortedIds)) { + if (!areArraysEqual(currentIds, newSortedIds)) { state.ids = newSortedIds } } + const mergeFunction: MergeFunction = mergeInsertion + return { removeOne, removeMany, diff --git a/packages/toolkit/src/entities/tests/sorted_state_adapter.test.ts b/packages/toolkit/src/entities/tests/sorted_state_adapter.test.ts index 9683715fb..cc1bf568d 100644 --- a/packages/toolkit/src/entities/tests/sorted_state_adapter.test.ts +++ b/packages/toolkit/src/entities/tests/sorted_state_adapter.test.ts @@ -1,6 +1,11 @@ import type { EntityAdapter, EntityState } from '../models' import { createEntityAdapter } from '../create_adapter' -import { createAction, createSlice, configureStore } from '@reduxjs/toolkit' +import { + createAction, + createSlice, + configureStore, + nanoid, +} from '@reduxjs/toolkit' import type { BookModel } from './fixtures/book' import { TheGreatGatsby, @@ -232,11 +237,11 @@ describe('Sorted State Adapter', () => { }) it('Replaces an existing entity if you change the ID while updating', () => { - const withAdded = adapter.setAll(state, [ - { id: 'a', title: 'First' }, - { id: 'b', title: 'Second' }, - { id: 'c', title: 'Third' }, - ]) + const a = { id: 'a', title: 'First' } + const b = { id: 'b', title: 'Second' } + const c = { id: 'c', title: 'Third' } + const d = { id: 'd', title: 'Fourth' } + const withAdded = adapter.setAll(state, [a, b, c]) const withUpdated = adapter.updateOne(withAdded, { id: 'b', @@ -247,7 +252,7 @@ describe('Sorted State Adapter', () => { const { ids, entities } = withUpdated - expect(ids.length).toBe(2) + expect(ids).toEqual(['a', 'c']) expect(entities.a).toBeTruthy() expect(entities.b).not.toBeTruthy() expect(entities.c).toBeTruthy() @@ -355,20 +360,30 @@ describe('Sorted State Adapter', () => { const withInitialItems = sortedItemsAdapter.setAll( sortedItemsAdapter.getInitialState(), [ + { id: 'C', order: 3, ts: 0 }, { id: 'A', order: 1, ts: 0 }, + { id: 'F', order: 4, ts: 0 }, { id: 'B', order: 2, ts: 0 }, - { id: 'C', order: 3, ts: 0 }, { id: 'D', order: 3, ts: 0 }, { id: 'E', order: 3, ts: 0 }, ], ) + expect(withInitialItems.ids).toEqual(['A', 'B', 'C', 'D', 'E', 'F']) + const updated = sortedItemsAdapter.updateOne(withInitialItems, { id: 'C', changes: { ts: 5 }, }) - expect(updated.ids).toEqual(['A', 'B', 'C', 'D', 'E']) + expect(updated.ids).toEqual(['A', 'B', 'C', 'D', 'E', 'F']) + + const updated2 = sortedItemsAdapter.updateOne(withInitialItems, { + id: 'D', + changes: { ts: 6 }, + }) + + expect(updated2.ids).toEqual(['A', 'B', 'C', 'D', 'E', 'F']) }) it('should let you update many entities by id in the state', () => { @@ -584,6 +599,190 @@ describe('Sorted State Adapter', () => { expect(withUpdate.entities['b']!.title).toBe(book1.title) }) + it('should minimize the amount of sorting work needed', () => { + const INITIAL_ITEMS = 10_000 + const ADDED_ITEMS = 1_000 + + type Entity = { id: string; name: string; position: number } + + let numSorts = 0 + + const adaptor = createEntityAdapter({ + selectId: (entity: Entity) => entity.id, + sortComparer: (a, b) => { + numSorts++ + if (a.position < b.position) return -1 + else if (a.position > b.position) return 1 + return 0 + }, + }) + + function generateItems(count: number) { + const items: readonly Entity[] = new Array(count) + .fill(undefined) + .map((x, i) => ({ + name: `${i}`, + position: Math.random(), + id: nanoid(), + })) + return items + } + + const entitySlice = createSlice({ + name: 'entity', + initialState: adaptor.getInitialState(), + reducers: { + updateOne: adaptor.updateOne, + upsertOne: adaptor.upsertOne, + upsertMany: adaptor.upsertMany, + addMany: adaptor.addMany, + }, + }) + + const store = configureStore({ + reducer: { + entity: entitySlice.reducer, + }, + middleware: (getDefaultMiddleware) => { + return getDefaultMiddleware({ + serializableCheck: false, + immutableCheck: false, + }) + }, + }) + + numSorts = 0 + + const logComparisons = false + + function measureComparisons(name: string, cb: () => void) { + numSorts = 0 + const start = new Date().getTime() + cb() + const end = new Date().getTime() + const duration = end - start + + if (logComparisons) { + console.log( + `${name}: sortComparer called ${numSorts.toLocaleString()} times in ${duration.toLocaleString()}ms`, + ) + } + } + + const initialItems = generateItems(INITIAL_ITEMS) + + measureComparisons('Original Setup', () => { + store.dispatch(entitySlice.actions.upsertMany(initialItems)) + }) + + expect(numSorts).toBeLessThan(INITIAL_ITEMS * 20) + + measureComparisons('Insert One (random)', () => { + store.dispatch( + entitySlice.actions.upsertOne({ + id: nanoid(), + position: Math.random(), + name: 'test', + }), + ) + }) + + expect(numSorts).toBeLessThan(50) + + measureComparisons('Insert One (middle)', () => { + store.dispatch( + entitySlice.actions.upsertOne({ + id: nanoid(), + position: 0.5, + name: 'test', + }), + ) + }) + + expect(numSorts).toBeLessThan(50) + + measureComparisons('Insert One (end)', () => { + store.dispatch( + entitySlice.actions.upsertOne({ + id: nanoid(), + position: 0.9998, + name: 'test', + }), + ) + }) + + expect(numSorts).toBeLessThan(50) + + const addedItems = generateItems(ADDED_ITEMS) + measureComparisons('Add Many', () => { + store.dispatch(entitySlice.actions.addMany(addedItems)) + }) + + expect(numSorts).toBeLessThan(ADDED_ITEMS * 20) + + // These numbers will vary because of the randomness, but generally + // with 10K items the old code had 200K+ sort calls, while the new code + // is around 13K sort calls. + expect(numSorts).toBeLessThan(20_000) + + const { ids } = store.getState().entity + const middleItemId = ids[(ids.length / 2) | 0] + + measureComparisons('Update One (end)', () => { + store.dispatch( + // Move this middle item near the end + entitySlice.actions.updateOne({ + id: middleItemId, + changes: { + position: 0.99999, + }, + }), + ) + }) + + const SORTING_COUNT_BUFFER = 100 + + expect(numSorts).toBeLessThan( + INITIAL_ITEMS + ADDED_ITEMS + SORTING_COUNT_BUFFER, + ) + + measureComparisons('Update One (middle)', () => { + store.dispatch( + // Move this middle item near the end + entitySlice.actions.updateOne({ + id: middleItemId, + changes: { + position: 0.42, + }, + }), + ) + }) + + expect(numSorts).toBeLessThan( + INITIAL_ITEMS + ADDED_ITEMS + SORTING_COUNT_BUFFER, + ) + + measureComparisons('Update One (replace)', () => { + store.dispatch( + // Move this middle item near the end + entitySlice.actions.updateOne({ + id: middleItemId, + changes: { + id: nanoid(), + position: 0.98, + }, + }), + ) + }) + + expect(numSorts).toBeLessThan( + INITIAL_ITEMS + ADDED_ITEMS + SORTING_COUNT_BUFFER, + ) + + // The old code was around 120K, the new code is around 10K. + //expect(numSorts).toBeLessThan(25_000) + }) + describe('can be used mutably when wrapped in createNextState', () => { test('removeAll', () => { const withTwo = adapter.addMany(state, [TheGreatGatsby, AnimalFarm]) diff --git a/packages/toolkit/src/entities/utils.ts b/packages/toolkit/src/entities/utils.ts index 5e3fb92fc..895ef2b35 100644 --- a/packages/toolkit/src/entities/utils.ts +++ b/packages/toolkit/src/entities/utils.ts @@ -1,3 +1,4 @@ +import { current, isDraft } from 'immer' import type { IdSelector, Update, @@ -35,23 +36,30 @@ export function ensureEntitiesArray( return entities } +export function getCurrent(value: T): T { + return isDraft(value) ? current(value) : value +} + export function splitAddedUpdatedEntities( newEntities: readonly T[] | Record, selectId: IdSelector, state: DraftableEntityState, -): [T[], Update[]] { +): [T[], Update[], Id[]] { newEntities = ensureEntitiesArray(newEntities) + const existingIdsArray = current(state.ids) as Id[] + const existingIds = new Set(existingIdsArray) + const added: T[] = [] const updated: Update[] = [] for (const entity of newEntities) { const id = selectIdValue(entity, selectId) - if (id in state.entities) { + if (existingIds.has(id)) { updated.push({ id, changes: entity }) } else { added.push(entity) } } - return [added, updated] + return [added, updated, existingIdsArray] }