Ensure we have a mostly-sorted items array to speed up sorting

Author: markerikson

packages/toolkit/src/entities/sorted_state_adapter.ts

@@ -76,6 +76,7 @@ export function createSortedStateAdapter<T, Id extends EntityId>(
     state: R,
   ): void {
     let appliedUpdates = false
+    let replacedIds = false
 
     for (let update of updates) {
       const entity: T | undefined = (state.entities as Record<Id, T>)[update.id]
@@ -88,13 +89,16 @@ export function createSortedStateAdapter<T, Id extends EntityId>(
       Object.assign(entity, update.changes)
       const newId = selectId(entity)
       if (update.id !== newId) {
+        replacedIds = true
         delete (state.entities as Record<Id, T>)[update.id]
+        const oldIndex = (state.ids as Id[]).indexOf(update.id)
+        state.ids[oldIndex] = newId
         ;(state.entities as Record<Id, T>)[newId] = entity
       }
     }
 
     if (appliedUpdates) {
-      resortEntities(state)
+      resortEntities(state, [], replacedIds)
     }
   }
 
@@ -136,11 +140,44 @@ export function createSortedStateAdapter<T, Id extends EntityId>(
       ;(state.entities as Record<Id, T>)[selectId(model)] = model
     })
 
-    resortEntities(state)
+    resortEntities(state, models)
   }
 
-  function resortEntities(state: R) {
-    const allEntities = Object.values(state.entities) as T[]
+  function resortEntities(
+    state: R,
+    addedItems: readonly T[] = [],
+    replacedIds = false,
+  ) {
+    let allEntities: T[]
+
+    allEntities = Object.values(state.entities) as T[]
+    if (replacedIds) {
+      // This is a really annoying edge case. Just figure this out from scratch
+      // rather than try to be clever. This will be more expensive since it isn't sorted right.
+      allEntities = Object.values(state.entities) as T[]
+    } else {
+      // We're starting with an already-sorted list.
+      let existingIds = state.ids
+
+      if (addedItems.length) {
+        // There's a couple edge cases where we can have duplicate item IDs.
+        // Ensure we don't have duplicates.
+        const uniqueIds = new Set(existingIds as Id[])
+
+        addedItems.forEach((item) => {
+          uniqueIds.add(selectId(item))
+        })
+        existingIds = Array.from(uniqueIds)
+      }
+
+      // By this point `ids` and `entities` should be 1:1, but not necessarily sorted.
+      // Make this a sorta-mostly-sorted array.
+      allEntities = existingIds.map(
+        (id) => (state.entities as Record<Id, T>)[id as Id],
+      )
+    }
+
+    // Now when we sort, things should be _close_ already, and fewer comparisons are needed.
     allEntities.sort(sort)
 
     const newSortedIds = allEntities.map(selectId)

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,
@@ -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()
@@ -584,6 +589,86 @@ describe('Sorted State Adapter', () => {
     expect(withUpdate.entities['b']!.title).toBe(book1.title)
   })
 
+  it('should minimize the amount of sorting work needed', () => {
+    const PARAMETERS = {
+      NUM_ITEMS: 10_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
+      },
+    })
+
+    const initialState: Entity[] = new Array(PARAMETERS.NUM_ITEMS)
+      .fill(undefined)
+      .map((x, i) => ({
+        name: `${i}`,
+        position: Math.random(),
+        id: nanoid(),
+      }))
+
+    const entitySlice = createSlice({
+      name: 'entity',
+      initialState: adaptor.getInitialState(undefined, initialState),
+      reducers: {
+        updateOne: adaptor.updateOne,
+        upsertOne: adaptor.upsertOne,
+        upsertMany: adaptor.upsertMany,
+      },
+    })
+
+    const store = configureStore({
+      reducer: {
+        entity: entitySlice.reducer,
+      },
+      middleware: (getDefaultMiddleware) => {
+        return getDefaultMiddleware({
+          serializableCheck: false,
+          immutableCheck: false,
+        })
+      },
+    })
+
+    store.dispatch(
+      entitySlice.actions.upsertOne({
+        id: nanoid(),
+        position: Math.random(),
+        name: 'test',
+      }),
+    )
+
+    // 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 130K sort calls.
+    expect(numSorts).toBeLessThan(200_000)
+
+    const { ids } = store.getState().entity
+    const middleItemId = ids[(ids.length / 2) | 0]
+
+    numSorts = 0
+
+    store.dispatch(
+      // Move this middle item near the end
+      entitySlice.actions.updateOne({
+        id: middleItemId,
+        changes: {
+          position: 0.99999,
+        },
+      }),
+    )
+    // 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])