Improve perf for internal map/filter ops (#5126)

* Add filterMap util

* Use filterMap instead of multiple filters

* Copy getCurrent util

* Iterate over current arrays where possible

* Track pending requests counter

* Fix typo in build config name

* Remove flatten tests

* Use reduce+flat for smaller filterMap size

Author: markerikson

packages/toolkit/src/query/core/buildMiddleware/cacheLifecycle.ts

@@ -4,8 +4,11 @@ import type {
   BaseQueryMeta,
   BaseQueryResult,
 } from '../../baseQueryTypes'
-import type { BaseEndpointDefinition } from '../../endpointDefinitions'
-import { DefinitionType, isAnyQueryDefinition } from '../../endpointDefinitions'
+import type {
+  BaseEndpointDefinition,
+  DefinitionType,
+} from '../../endpointDefinitions'
+import { isAnyQueryDefinition } from '../../endpointDefinitions'
 import type { QueryCacheKey, RootState } from '../apiState'
 import type {
   MutationResultSelectorResult,

packages/toolkit/src/query/core/buildMiddleware/invalidationByTags.ts

@@ -38,13 +38,25 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
   )
 
   const isQueryEnd = isAnyOf(
-    isFulfilled(mutationThunk, queryThunk),
-    isRejected(mutationThunk, queryThunk),
+    isFulfilled(queryThunk, mutationThunk),
+    isRejected(queryThunk, mutationThunk),
   )
-
   let pendingTagInvalidations: FullTagDescription<string>[] = []
+  // Track via counter so we can avoid iterating over state every time
+  let pendingRequestCount = 0
 
   const handler: ApiMiddlewareInternalHandler = (action, mwApi) => {
+    if (
+      queryThunk.pending.match(action) ||
+      mutationThunk.pending.match(action)
+    ) {
+      pendingRequestCount++
+    }
+
+    if (isQueryEnd(action)) {
+      pendingRequestCount = Math.max(0, pendingRequestCount - 1)
+    }
+
     if (isThunkActionWithTags(action)) {
       invalidateTags(
         calculateProvidedByThunk(
@@ -72,16 +84,8 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
     }
   }
 
-  function hasPendingRequests(
-    state: CombinedState<EndpointDefinitions, string, string>,
-  ) {
-    const { queries, mutations } = state
-    for (const cacheRecord of [queries, mutations]) {
-      for (const key in cacheRecord) {
-        if (cacheRecord[key]?.status === QueryStatus.pending) return true
-      }
-    }
-    return false
+  function hasPendingRequests() {
+    return pendingRequestCount > 0
   }
 
   function invalidateTags(
@@ -95,7 +99,7 @@ export const buildInvalidationByTagsHandler: InternalHandlerBuilder = ({
 
     if (
       state.config.invalidationBehavior === 'delayed' &&
-      hasPendingRequests(state)
+      hasPendingRequests()
     ) {
       return
     }

packages/toolkit/src/query/core/buildSelectors.ts

@@ -13,7 +13,7 @@ import type {
   TagTypesFrom,
 } from '../endpointDefinitions'
 import { expandTagDescription } from '../endpointDefinitions'
-import { flatten, isNotNullish } from '../utils'
+import { filterMap, isNotNullish } from '../utils'
 import type {
   InfiniteData,
   InfiniteQueryConfigOptions,
@@ -342,7 +342,8 @@ export function buildSelectors<
   }> {
     const apiState = state[reducerPath]
     const toInvalidate = new Set<QueryCacheKey>()
-    for (const tag of tags.filter(isNotNullish).map(expandTagDescription)) {
+    const finalTags = filterMap(tags, isNotNullish, expandTagDescription)
+    for (const tag of finalTags) {
       const provided = apiState.provided.tags[tag.type]
       if (!provided) {
         continue
@@ -353,27 +354,23 @@ export function buildSelectors<
           ? // id given: invalidate all queries that provide this type & id
             provided[tag.id]
           : // no id: invalidate all queries that provide this type
-            flatten(Object.values(provided))) ?? []
+            Object.values(provided).flat()) ?? []
 
       for (const invalidate of invalidateSubscriptions) {
         toInvalidate.add(invalidate)
       }
     }
 
-    return flatten(
-      Array.from(toInvalidate.values()).map((queryCacheKey) => {
-        const querySubState = apiState.queries[queryCacheKey]
-        return querySubState
-          ? [
-              {
-                queryCacheKey,
-                endpointName: querySubState.endpointName!,
-                originalArgs: querySubState.originalArgs,
-              },
-            ]
-          : []
-      }),
-    )
+    return Array.from(toInvalidate.values()).flatMap((queryCacheKey) => {
+      const querySubState = apiState.queries[queryCacheKey]
+      return querySubState
+        ? {
+            queryCacheKey,
+            endpointName: querySubState.endpointName!,
+            originalArgs: querySubState.originalArgs,
+          }
+        : []
+    })
   }
 
   function selectCachedArgsForQuery<
@@ -382,18 +379,18 @@ export function buildSelectors<
     state: RootState,
     queryName: QueryName,
   ): Array<QueryArgFromAnyQuery<Definitions[QueryName]>> {
-    return Object.values(selectQueries(state) as QueryState<any>)
-      .filter(
-        (
-          entry,
-        ): entry is Exclude<
-          QuerySubState<Definitions[QueryName]>,
-          { status: QueryStatus.uninitialized }
-        > =>
-          entry?.endpointName === queryName &&
-          entry.status !== QueryStatus.uninitialized,
-      )
-      .map((entry) => entry.originalArgs)
+    return filterMap(
+      Object.values(selectQueries(state) as QueryState<any>),
+      (
+        entry,
+      ): entry is Exclude<
+        QuerySubState<Definitions[QueryName]>,
+        { status: QueryStatus.uninitialized }
+      > =>
+        entry?.endpointName === queryName &&
+        entry.status !== QueryStatus.uninitialized,
+      (entry) => entry.originalArgs,
+    )
   }
 
   function getHasNextPage(

packages/toolkit/src/query/core/buildSlice.ts

@@ -57,6 +57,7 @@ import type { ApiContext } from '../apiTypes'
 import { isUpsertQuery } from './buildInitiate'
 import type { InternalSerializeQueryArgs } from '../defaultSerializeQueryArgs'
 import type { UnwrapPromise } from '../tsHelpers'
+import { getCurrent } from '../utils/getCurrent'
 
 /**
  * A typesafe single entry to be upserted into the cache
@@ -587,7 +588,7 @@ export function buildSlice({
     draft: InvalidationState<any>,
     queryCacheKey: QueryCacheKey,
   ) {
-    const existingTags = draft.keys[queryCacheKey] ?? []
+    const existingTags = getCurrent(draft.keys[queryCacheKey] ?? [])
 
     // Delete this cache key from any existing tags that may have provided it
     for (const tag of existingTags) {
@@ -596,7 +597,7 @@ export function buildSlice({
       const tagSubscriptions = draft.tags[tagType]?.[tagId]
 
       if (tagSubscriptions) {
-        draft.tags[tagType][tagId] = tagSubscriptions.filter(
+        draft.tags[tagType][tagId] = getCurrent(tagSubscriptions).filter(
           (qc) => qc !== queryCacheKey,
         )
       }

packages/toolkit/src/query/endpointDefinitions.ts

@@ -39,6 +39,7 @@ import type {
 } from './tsHelpers'
 import { isNotNullish } from './utils'
 import type { NamedSchemaError } from './standardSchema'
+import { filterMap } from './utils/filterMap'
 
 const rawResultType = /* @__PURE__ */ Symbol()
 const resultType = /* @__PURE__ */ Symbol()
@@ -1406,20 +1407,21 @@ export function calculateProvidedBy<ResultType, QueryArg, ErrorType, MetaType>(
   meta: MetaType | undefined,
   assertTagTypes: AssertTagTypes,
 ): readonly FullTagDescription<string>[] {
-  if (isFunction(description)) {
-    return description(
-      result as ResultType,
-      error as undefined,
-      queryArg,
-      meta as MetaType,
+  const finalDescription = isFunction(description)
+    ? description(
+        result as ResultType,
+        error as undefined,
+        queryArg,
+        meta as MetaType,
+      )
+    : description
+
+  if (finalDescription) {
+    return filterMap(finalDescription, isNotNullish, (tag) =>
+      assertTagTypes(expandTagDescription(tag)),
     )
-      .filter(isNotNullish)
-      .map(expandTagDescription)
-      .map(assertTagTypes)
-  }
-  if (Array.isArray(description)) {
-    return description.map(expandTagDescription).map(assertTagTypes)
   }
+
   return []
 }
 

packages/toolkit/src/query/tests/utils.test.ts

@@ -1,10 +1,5 @@
 import { vi } from 'vitest'
-import {
-  isOnline,
-  isDocumentVisible,
-  flatten,
-  joinUrls,
-} from '@internal/query/utils'
+import { isOnline, isDocumentVisible, joinUrls } from '@internal/query/utils'
 
 afterAll(() => {
   vi.restoreAllMocks()
@@ -96,14 +91,3 @@ describe('joinUrls', () => {
     expect(joinUrls(base, url)).toBe(expected)
   })
 })
-
-describe('flatten', () => {
-  test('flattens an array to a depth of 1', () => {
-    expect(flatten([1, 2, [3, 4]])).toEqual([1, 2, 3, 4])
-  })
-  test('does not flatten to a depth of 2', () => {
-    const flattenResult = flatten([1, 2, [3, 4, [5, 6]]])
-    expect(flattenResult).not.toEqual([1, 2, 3, 4, 5, 6])
-    expect(flattenResult).toEqual([1, 2, 3, 4, [5, 6]])
-  })
-})

packages/toolkit/src/query/utils/filterMap.ts

@@ -0,0 +1,27 @@
+// Preserve type guard predicate behavior when passing to mapper
+export function filterMap<T, U, S extends T = T>(
+  array: readonly T[],
+  predicate: (item: T, index: number) => item is S,
+  mapper: (item: S, index: number) => U | U[],
+): U[]
+
+export function filterMap<T, U>(
+  array: readonly T[],
+  predicate: (item: T, index: number) => boolean,
+  mapper: (item: T, index: number) => U | U[],
+): U[]
+
+export function filterMap<T, U>(
+  array: readonly T[],
+  predicate: (item: T, index: number) => boolean,
+  mapper: (item: T, index: number) => U | U[],
+): U[] {
+  return array
+    .reduce<(U | U[])[]>((acc, item, i) => {
+      if (predicate(item as any, i)) {
+        acc.push(mapper(item as any, i))
+      }
+      return acc
+    }, [])
+    .flat() as U[]
+}

packages/toolkit/src/query/utils/flatten.ts

@@ -0,0 +1,6 @@
+import type { Draft } from 'immer'
+import { current, isDraft } from 'immer'
+
+export function getCurrent<T>(value: T | Draft<T>): T {
+  return (isDraft(value) ? current(value) : value) as T
+}

packages/toolkit/src/query/utils/getCurrent.ts

@@ -1,7 +1,7 @@
 export * from './capitalize'
 export * from './copyWithStructuralSharing'
 export * from './countObjectKeys'
-export * from './flatten'
+export * from './filterMap'
 export * from './isAbsoluteUrl'
 export * from './isDocumentVisible'
 export * from './isNotNullish'