From 08533a7578f85925d9cc8f978f6e0f828659063a Mon Sep 17 00:00:00 2001 From: Jake Teton-Landis Date: Sun, 24 Aug 2025 17:30:42 -0400 Subject: [PATCH 1/8] speculative optimizations --- packages/query-core/src/query.ts | 47 +++++ packages/query-core/src/queryCache.ts | 283 +++++++++++++++++++++++++- 2 files changed, 324 insertions(+), 6 deletions(-) diff --git a/packages/query-core/src/query.ts b/packages/query-core/src/query.ts index df5b7c030e..1b33ff8cbc 100644 --- a/packages/query-core/src/query.ts +++ b/packages/query-core/src/query.ts @@ -1,5 +1,6 @@ import { ensureQueryFn, + hashKey, noop, replaceData, resolveEnabled, @@ -20,6 +21,7 @@ import type { OmitKeyof, QueryFunctionContext, QueryKey, + QueryKeyHashFunction, QueryMeta, QueryOptions, QueryStatus, @@ -154,6 +156,40 @@ export interface SetStateOptions { meta?: any } +export class RefCountSet { + #refcounts = new Map(); + + [Symbol.iterator]() { + return this.#refcounts.keys() + } + + get size() { + return this.#refcounts.size + } + + add(value: T) { + const n = this.#refcounts.get(value) ?? 0 + this.#refcounts.set(value, n + 1) + } + + remove(value: T) { + let n = this.#refcounts.get(value) + if (n === undefined) { + return + } + n-- + if (n === 0) { + this.#refcounts.delete(value) + } else { + this.#refcounts.set(value, n) + } + } +} + +// Somewhere; I'm not sure what the appropriate place to store this is. Maybe part of QueryCache? +// Making it global is easier in my example code. +export const allQueryKeyHashFns = new RefCountSet>() + // CLASS export class Query< @@ -202,8 +238,19 @@ export class Query< setOptions( options?: QueryOptions, ): void { + const oldHashFn = this.options?.queryKeyHashFn this.options = { ...this.#defaultOptions, ...options } + const newHashFn = this.options?.queryKeyHashFn + if (oldHashFn !== newHashFn) { + if (oldHashFn && oldHashFn !== hashKey) { + allQueryKeyHashFns.remove(oldHashFn) + } + if (newHashFn && newHashFn !== hashKey) { + allQueryKeyHashFns.add(newHashFn) + } + } + this.updateGcTime(this.options.gcTime) } diff --git a/packages/query-core/src/queryCache.ts b/packages/query-core/src/queryCache.ts index dd7123eaac..40b2b2bc60 100644 --- a/packages/query-core/src/queryCache.ts +++ b/packages/query-core/src/queryCache.ts @@ -1,13 +1,14 @@ -import { hashQueryKeyByOptions, matchQuery } from './utils' -import { Query } from './query' +import { hashKey, hashQueryKeyByOptions, matchQuery } from './utils' +import { Query, RefCountSet, allQueryKeyHashFns } from './query' import { notifyManager } from './notifyManager' import { Subscribable } from './subscribable' -import type { QueryFilters } from './utils' import type { Action, QueryState } from './query' +import type { QueryFilters } from './utils' import type { DefaultError, NotifyEvent, QueryKey, + QueryKeyHashFunction, QueryOptions, WithRequired, } from './types' @@ -87,10 +88,218 @@ export interface QueryStore { values: () => IterableIterator } +type Primitive = string | number | boolean | bigint | symbol | undefined | null +function isPrimitive(value: unknown): value is Primitive { + if (value === undefined || value === null) { + return true + } + const t = typeof value + switch (t) { + case 'object': + case 'function': + return false + case 'string': + case 'number': + case 'boolean': + case 'bigint': + case 'symbol': + case 'undefined': + return true + default: + t satisfies never + return false + } +} + +type MapTrieNode = { + key: Primitive + /** + * - Entries who's TKey path leads exactly to this node, + * therefor their path is all primitives. + * ourPath = [p1, p2, p3, ...pN] + * theirPath = [p1, p2, p3, ...pN] + */ + exact?: RefCountSet + /** + * - Entries who's path after this point contains non-primitive keys. + * Such entries cannot be looked up by value deeper in the trie. + * Implies their TKey path != this node's keyPath. + * ourPath = [p1, p2, p3, ...pN] + * theirPath = [p1, p2, p3, ...pN, nonPrimitive, ...] + */ + nonPrimitiveSuffix?: RefCountSet + /** Child nodes storing entries who's TKey path is prefixed with this node's path. */ + children?: Map> +} + +/** Path length is always 1 greater than the key length, as it includes the root node. */ +function traverse( + root: MapTrieNode, + key: QueryKey, + // May create a child node if needed + lookup: ( + parent: MapTrieNode, + key: Primitive, + ) => MapTrieNode | undefined, +): Array> { + const path: Array> = [root] + let node: MapTrieNode | undefined = root + for (let i = 0; i < key.length && node; i++) { + const keyPart = key[i] + if (isPrimitive(keyPart)) { + node = lookup(node, keyPart) + } else { + node = undefined + } + if (node) { + path.push(node) + } + } + return path +} + +function gcPath(path: Array>): void { + if (path.length === 0) { + return + } + for (let i = path.length - 1; i >= 0; i--) { + const node = path[i] + if (!node) { + throw new Error('Should never occur (bug in MapTrie)') + } + + if ( + node.exact?.size || + node.nonPrimitiveSuffix?.size || + node.children?.size + ) { + // Has data. Do not GC. + return + } + + const parent = path[i - 1] + parent?.children?.delete(node.key) + } +} + +class MapTrieSet { + #root: MapTrieNode = { + key: undefined, + } + + add(key: TKey, value: TValue): void { + const path = traverse(this.#root, key, (parent, keyPart) => { + parent.children ??= new Map() + let child = parent.children.get(keyPart) + if (!child) { + child = { key: keyPart } + parent.children.set(keyPart, child) + } + return child + }) + const lastPathNode = path[path.length - 1] + if (!lastPathNode) { + throw new Error('Should never occur (bug in MapTrie)') + } + + if (key.length === path.length - 1) { + lastPathNode.exact ??= new RefCountSet() + lastPathNode.exact.add(value) + } else { + lastPathNode.nonPrimitiveSuffix ??= new RefCountSet() + lastPathNode.nonPrimitiveSuffix.add(value) + } + } + + remove(key: TKey, value: TValue): void { + const path = traverse(this.#root, key, (parent, keyPart) => + parent.children?.get(keyPart), + ) + const lastPathNode = path[path.length - 1] + if (!lastPathNode) { + throw new Error('Should never occur (bug in MapTrie)') + } + if (key.length === path.length - 1) { + lastPathNode.exact?.remove(value) + gcPath(path) + } else if (!isPrimitive(key[path.length - 1])) { + lastPathNode.nonPrimitiveSuffix?.remove(value) + gcPath(path) + } + } + + /** + * Returns all values that match the given key: + * Either the value has the same key and is all primitives, + * Or the value's key is a suffix of the given key and contains a non-primitive key. + */ + getByPrefix(key: TKey): Iterable | undefined { + let miss = false + const path = traverse(this.#root, key, (parent, keyPart) => { + const child = parent.children?.get(keyPart) + if (!child) { + miss = true + return undefined + } + return child + }) + // Failed to look up one of the primitive keys in the path. + // This means there's no match at all. + // Appears to be incorrectly reported by @typescript-eslint as always false :\ + // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition + if (miss) { + return undefined + } + + const lastNode = path[path.length - 1] + if (!lastNode) { + throw new Error('Should never occur (bug in MapTrie)') + } + + // If the `key` is all primitives then we need to recurse to find all values + // that match the prefix, as these values will be stored deeper in the trie. + // + // If the `key` contains a non-primitive part after the returned path, + // then all possible values that have the suffix are stored in this node. + const isPrimitivePath = path.length - 1 === key.length + if (!isPrimitivePath) { + return lastNode.nonPrimitiveSuffix?.[Symbol.iterator]() + } + + // See if we can avoid instantiating a generator + if ( + !lastNode.children && + (lastNode.exact || lastNode.nonPrimitiveSuffix) && + !(lastNode.exact && lastNode.nonPrimitiveSuffix) + ) { + return lastNode.exact ?? lastNode.nonPrimitiveSuffix + } + + return (function* depthFirstPrefixIterator() { + const queue = [lastNode] + while (queue.length > 0) { + const node = queue.pop()! + if (node.exact) { + yield* node.exact + } + if (node.nonPrimitiveSuffix) { + yield* node.nonPrimitiveSuffix + } + if (node.children) { + for (const child of node.children.values()) { + queue.push(child) + } + } + } + })() + } +} + // CLASS export class QueryCache extends Subscribable { #queries: QueryStore + #keyIndex = new MapTrieSet() constructor(public config: QueryCacheConfig = {}) { super() @@ -133,6 +342,7 @@ export class QueryCache extends Subscribable { add(query: Query): void { if (!this.#queries.has(query.queryHash)) { this.#queries.set(query.queryHash, query) + this.#keyIndex.add(query.queryKey, query) this.notify({ type: 'added', @@ -149,6 +359,7 @@ export class QueryCache extends Subscribable { if (queryInMap === query) { this.#queries.delete(query.queryHash) + this.#keyIndex.remove(query.queryKey, query) } this.notify({ type: 'removed', query }) @@ -184,19 +395,79 @@ export class QueryCache extends Subscribable { filters: WithRequired, ): Query | undefined { const defaultedFilters = { exact: true, ...filters } + if (defaultedFilters.exact) { + return this.findExact(filters) + } + + const candidates = this.#keyIndex.getByPrefix(filters.queryKey) + if (!candidates) { + return undefined + } + + for (const query of candidates) { + if (matchQuery(defaultedFilters, query)) { + return query as unknown as + | Query + | undefined + } + } - return this.getAll().find((query) => - matchQuery(defaultedFilters, query), - ) as Query | undefined + return undefined } findAll(filters: QueryFilters = {}): Array { + if (filters.exact && filters.queryKey) { + const query = this.findExact(filters) + return query ? [query] : [] + } + + if (filters.queryKey) { + const withPrefix = this.#keyIndex.getByPrefix(filters.queryKey) + const candidates = withPrefix ? Array.from(withPrefix) : [] + return candidates.filter((query) => matchQuery(filters, query)) + } + const queries = this.getAll() return Object.keys(filters).length > 0 ? queries.filter((query) => matchQuery(filters, query)) : queries } + private findExact< + TQueryFnData = unknown, + TError = DefaultError, + TData = TQueryFnData, + >( + filters: QueryFilters, + ): Query | undefined { + const tryHashFn = (hashFn: QueryKeyHashFunction) => { + try { + const query = this.get(hashFn(filters.queryKey)) + if (query && matchQuery(filters, query)) { + // Confirmed the query actually uses the hash function we tried + // and matches the non-queryKey filters + return query + } else { + return undefined + } + } catch (error) { + return undefined + } + } + + let query = tryHashFn(hashKey) + if (!query) { + for (const hashFn of allQueryKeyHashFns) { + query = tryHashFn(hashFn) + if (query) { + break + } + } + } + + return query as unknown as Query | undefined + } + notify(event: QueryCacheNotifyEvent): void { notifyManager.batch(() => { this.listeners.forEach((listener) => { From c0a661c0e56835d06aa10e35201600baf00a2da3 Mon Sep 17 00:00:00 2001 From: Jake Teton-Landis Date: Sun, 24 Aug 2025 18:15:18 -0400 Subject: [PATCH 2/8] dont leak hash functions --- packages/query-core/src/query.ts | 56 ++++++--------------------- packages/query-core/src/queryCache.ts | 55 +++++++++++++++++++++++++- 2 files changed, 65 insertions(+), 46 deletions(-) diff --git a/packages/query-core/src/query.ts b/packages/query-core/src/query.ts index 1b33ff8cbc..ba6f0c6174 100644 --- a/packages/query-core/src/query.ts +++ b/packages/query-core/src/query.ts @@ -1,6 +1,5 @@ import { ensureQueryFn, - hashKey, noop, replaceData, resolveEnabled, @@ -21,7 +20,6 @@ import type { OmitKeyof, QueryFunctionContext, QueryKey, - QueryKeyHashFunction, QueryMeta, QueryOptions, QueryStatus, @@ -156,40 +154,6 @@ export interface SetStateOptions { meta?: any } -export class RefCountSet { - #refcounts = new Map(); - - [Symbol.iterator]() { - return this.#refcounts.keys() - } - - get size() { - return this.#refcounts.size - } - - add(value: T) { - const n = this.#refcounts.get(value) ?? 0 - this.#refcounts.set(value, n + 1) - } - - remove(value: T) { - let n = this.#refcounts.get(value) - if (n === undefined) { - return - } - n-- - if (n === 0) { - this.#refcounts.delete(value) - } else { - this.#refcounts.set(value, n) - } - } -} - -// Somewhere; I'm not sure what the appropriate place to store this is. Maybe part of QueryCache? -// Making it global is easier in my example code. -export const allQueryKeyHashFns = new RefCountSet>() - // CLASS export class Query< @@ -238,16 +202,20 @@ export class Query< setOptions( options?: QueryOptions, ): void { - const oldHashFn = this.options?.queryKeyHashFn + const oldOptions = (this.options as typeof this.options | undefined) + ? this.options + : undefined + this.options = { ...this.#defaultOptions, ...options } - const newHashFn = this.options?.queryKeyHashFn - if (oldHashFn !== newHashFn) { - if (oldHashFn && oldHashFn !== hashKey) { - allQueryKeyHashFns.remove(oldHashFn) - } - if (newHashFn && newHashFn !== hashKey) { - allQueryKeyHashFns.add(newHashFn) + if (oldOptions) { + // Do not do this update when first created. + // The QueryCache manages admission / removal itself. + // We only need to keep it up-to-date here. + const prevHashFn = oldOptions.queryKeyHashFn + const newHashFn = this.options.queryKeyHashFn + if (prevHashFn !== newHashFn) { + this.#cache.onQueryKeyHashFunctionChanged(prevHashFn, newHashFn) } } diff --git a/packages/query-core/src/queryCache.ts b/packages/query-core/src/queryCache.ts index 40b2b2bc60..3bb922a923 100644 --- a/packages/query-core/src/queryCache.ts +++ b/packages/query-core/src/queryCache.ts @@ -1,5 +1,5 @@ import { hashKey, hashQueryKeyByOptions, matchQuery } from './utils' -import { Query, RefCountSet, allQueryKeyHashFns } from './query' +import { Query } from './query' import { notifyManager } from './notifyManager' import { Subscribable } from './subscribable' import type { Action, QueryState } from './query' @@ -88,6 +88,36 @@ export interface QueryStore { values: () => IterableIterator } +class RefCountSet { + #refcounts = new Map(); + + [Symbol.iterator]() { + return this.#refcounts.keys() + } + + get size() { + return this.#refcounts.size + } + + add(value: T) { + const n = this.#refcounts.get(value) ?? 0 + this.#refcounts.set(value, n + 1) + } + + remove(value: T) { + let n = this.#refcounts.get(value) + if (n === undefined) { + return + } + n-- + if (n === 0) { + this.#refcounts.delete(value) + } else { + this.#refcounts.set(value, n) + } + } +} + type Primitive = string | number | boolean | bigint | symbol | undefined | null function isPrimitive(value: unknown): value is Primitive { if (value === undefined || value === null) { @@ -300,6 +330,7 @@ class MapTrieSet { export class QueryCache extends Subscribable { #queries: QueryStore #keyIndex = new MapTrieSet() + #knownHashFns = new RefCountSet>() constructor(public config: QueryCacheConfig = {}) { super() @@ -343,6 +374,10 @@ export class QueryCache extends Subscribable { if (!this.#queries.has(query.queryHash)) { this.#queries.set(query.queryHash, query) this.#keyIndex.add(query.queryKey, query) + const hashFn = query.options.queryKeyHashFn + if (hashFn) { + this.#knownHashFns.add(hashFn) + } this.notify({ type: 'added', @@ -360,6 +395,10 @@ export class QueryCache extends Subscribable { if (queryInMap === query) { this.#queries.delete(query.queryHash) this.#keyIndex.remove(query.queryKey, query) + const hashFn = query.options.queryKeyHashFn + if (hashFn) { + this.#knownHashFns.remove(hashFn) + } } this.notify({ type: 'removed', query }) @@ -457,7 +496,7 @@ export class QueryCache extends Subscribable { let query = tryHashFn(hashKey) if (!query) { - for (const hashFn of allQueryKeyHashFns) { + for (const hashFn of this.#knownHashFns) { query = tryHashFn(hashFn) if (query) { break @@ -491,4 +530,16 @@ export class QueryCache extends Subscribable { }) }) } + + onQueryKeyHashFunctionChanged( + before: QueryKeyHashFunction | undefined, + after: QueryKeyHashFunction | undefined, + ): void { + if (before) { + this.#knownHashFns.remove(before) + } + if (after) { + this.#knownHashFns.add(after) + } + } } From 06039658f84625d83bbe417cfea860dc2f6787a2 Mon Sep 17 00:00:00 2001 From: Jake Teton-Landis Date: Sun, 24 Aug 2025 18:24:04 -0400 Subject: [PATCH 3/8] test that prefix matches are not O(n) --- .../src/__tests__/queryCache.test.tsx | 78 ++++++++++++++++++- 1 file changed, 77 insertions(+), 1 deletion(-) diff --git a/packages/query-core/src/__tests__/queryCache.test.tsx b/packages/query-core/src/__tests__/queryCache.test.tsx index 6edb5598f7..42422a6510 100644 --- a/packages/query-core/src/__tests__/queryCache.test.tsx +++ b/packages/query-core/src/__tests__/queryCache.test.tsx @@ -1,6 +1,6 @@ import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest' import { queryKey, sleep } from '@tanstack/query-test-utils' -import { QueryCache, QueryClient, QueryObserver } from '..' +import { Query, QueryCache, QueryClient, QueryObserver } from '..' describe('queryCache', () => { let queryClient: QueryClient @@ -293,6 +293,82 @@ describe('queryCache', () => { expect(queryCache.findAll({ fetchStatus: 'fetching' })).toEqual([]) }) + test('only visits queries with the given prefix', async () => { + const prefix = ['posts'] + queryClient.prefetchQuery({ + queryKey: ['other'], + queryFn: () => sleep(100).then(() => 'other'), + }) + queryClient.prefetchQuery({ + queryKey: ['other', {}], + queryFn: () => sleep(100).then(() => 'otherObjectSuffix'), + }) + queryClient.prefetchQuery({ + queryKey: [{}, 'objectPrefix'], + queryFn: () => sleep(100).then(() => 'otherObjectPrefix'), + }) + + queryClient.prefetchQuery({ + queryKey: prefix, + queryFn: () => sleep(100).then(() => 'exactMatch'), + }) + const exactMatchQuery = queryCache.find({ queryKey: prefix }) + + queryClient.prefetchQuery({ + queryKey: [...prefix, 1], + queryFn: () => sleep(100).then(() => 'primitiveSuffix'), + }) + const primitiveSuffixQuery = queryCache.find({ queryKey: [...prefix, 1] }) + + queryClient.prefetchQuery({ + queryKey: [...prefix, 2, 'primitiveSuffix'], + queryFn: () => sleep(100).then(() => 'primitiveSuffixLength2'), + }) + const primitiveSuffixLength2Query = queryCache.find({ + queryKey: [...prefix, 2, 'primitiveSuffix'], + }) + + queryClient.prefetchQuery({ + queryKey: [...prefix, 3, {}, 'obj'], + queryFn: () => sleep(100).then(() => 'matchingObjectSuffix'), + }) + const matchingObjectSuffixQuery = queryCache.find({ + queryKey: [...prefix, 3, {}, 'obj'], + }) + + await vi.advanceTimersByTimeAsync(100) + + let predicateCallCount = 0 + const results = queryCache.findAll({ + queryKey: prefix, + predicate: () => { + predicateCallCount++ + return true + }, + }) + + expect(predicateCallCount).toBe(results.length) + + const sortByHash = (a: Query | undefined, b: Query | undefined) => { + if (!a && !b) return 0 + if (!a) return -1 + if (!b) return 1 + if (a.queryHash < b.queryHash) return -1 + if (a.queryHash > b.queryHash) return 1 + return 0 + } + results.sort(sortByHash) + + expect(results).toEqual( + [ + exactMatchQuery, + primitiveSuffixQuery, + primitiveSuffixLength2Query, + matchingObjectSuffixQuery, + ].sort(sortByHash), + ) + }) + test('should return all the queries when no filters are defined', async () => { const key1 = queryKey() const key2 = queryKey() From 5b4d0fdee8dc6b1872a73207392753536397ed74 Mon Sep 17 00:00:00 2001 From: Jake Teton-Landis Date: Sun, 24 Aug 2025 18:40:20 -0400 Subject: [PATCH 4/8] iterate in correct order (oldest to newest) --- packages/query-core/src/queryCache.ts | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/packages/query-core/src/queryCache.ts b/packages/query-core/src/queryCache.ts index 3bb922a923..c003a81bc8 100644 --- a/packages/query-core/src/queryCache.ts +++ b/packages/query-core/src/queryCache.ts @@ -316,9 +316,8 @@ class MapTrieSet { yield* node.nonPrimitiveSuffix } if (node.children) { - for (const child of node.children.values()) { - queue.push(child) - } + const children = Array.from(node.children.values()).reverse() + children.forEach((child) => queue.push(child)) } } })() From 6e6e8a7e25139bbf826d5a716ba4774e4b62c38c Mon Sep 17 00:00:00 2001 From: Jake Teton-Landis Date: Sun, 24 Aug 2025 19:47:27 -0400 Subject: [PATCH 5/8] fixup: always use defaultedFilters in QueryCache#find --- packages/query-core/src/queryCache.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/query-core/src/queryCache.ts b/packages/query-core/src/queryCache.ts index c003a81bc8..e4eee512a9 100644 --- a/packages/query-core/src/queryCache.ts +++ b/packages/query-core/src/queryCache.ts @@ -434,10 +434,10 @@ export class QueryCache extends Subscribable { ): Query | undefined { const defaultedFilters = { exact: true, ...filters } if (defaultedFilters.exact) { - return this.findExact(filters) + return this.findExact(defaultedFilters) } - const candidates = this.#keyIndex.getByPrefix(filters.queryKey) + const candidates = this.#keyIndex.getByPrefix(defaultedFilters.queryKey) if (!candidates) { return undefined } From f102be74062a692d1e8dbe06db8f44e1e6901039 Mon Sep 17 00:00:00 2001 From: Jake Teton-Landis Date: Mon, 25 Aug 2025 01:46:46 -0400 Subject: [PATCH 6/8] trie: now has same semantics as filtering an array by partialMatchKey --- packages/query-core/src/queryCache.ts | 433 +++++++++++++++++--------- 1 file changed, 289 insertions(+), 144 deletions(-) diff --git a/packages/query-core/src/queryCache.ts b/packages/query-core/src/queryCache.ts index e4eee512a9..34febf87ca 100644 --- a/packages/query-core/src/queryCache.ts +++ b/packages/query-core/src/queryCache.ts @@ -1,4 +1,9 @@ -import { hashKey, hashQueryKeyByOptions, matchQuery } from './utils' +import { + hashKey, + hashQueryKeyByOptions, + matchQuery, + partialMatchKey, +} from './utils' import { Query } from './query' import { notifyManager } from './notifyManager' import { Subscribable } from './subscribable' @@ -141,120 +146,296 @@ function isPrimitive(value: unknown): value is Primitive { } } -type MapTrieNode = { - key: Primitive +/** + * Like Map, but object keys have value semantics equality based + * on partialMatchKey, instead of reference equality. + * + * Lookups by object are O(NumberOfKeys) instead of O(1). + * + * ```ts + * const queryKeyMap = new QueryKeyElementMap<{ orderBy: 'likes', order: 'desc', limit: 30 }, string>() + * queryKeyMap.set({ orderBy: 'likes', order: 'desc', limit: 30 }, 'value') + * queryKeyMap.get({ orderBy: 'likes', order: 'desc', limit: 30 }) // 'value' + * + * const vanillaMap = new Map<{ orderBy: 'likes', order: 'desc', limit: 30 }, string>() + * vanillaMap.set({ orderBy: 'likes', order: 'desc', limit: 30 }, 'value') + * vanillaMap.get({ orderBy: 'likes', order: 'desc', limit: 30 }) // undefined + * ``` + */ +class QueryKeyElementMap { + #primitiveMap = new Map() + #objectMap = new Map() + + get size() { + return this.#primitiveMap.size + this.#objectMap.size + } + + get( + key: (TKeyElement & Primitive) | (TKeyElement & object), + ): TValue | undefined { + if (isPrimitive(key)) { + return this.#primitiveMap.get(key) + } + + const matchingKey = this.findMatchingObjectKey(key) + if (matchingKey) { + return this.#objectMap.get(matchingKey) + } + + return undefined + } + + set( + key: (TKeyElement & Primitive) | (TKeyElement & object), + value: TValue, + ): void { + if (isPrimitive(key)) { + this.#primitiveMap.set(key, value) + return + } + + const matchingKey = this.findMatchingObjectKey(key) + this.#objectMap.set(matchingKey ?? key, value) + } + + delete(key: (TKeyElement & Primitive) | (TKeyElement & object)): boolean { + if (isPrimitive(key)) { + return this.#primitiveMap.delete(key) + } + + const matchingKey = this.findMatchingObjectKey(key) + if (matchingKey) { + this.#objectMap.delete(matchingKey) + return true + } + return false + } + + values(): Iterable | undefined { + if (!this.#primitiveMap.size && !this.#objectMap.size) { + return undefined + } + + if (this.#primitiveMap.size && !this.#objectMap.size) { + return this.#primitiveMap.values() + } + + if (!this.#primitiveMap.size && this.#objectMap.size) { + return this.#objectMap.values() + } + + const primitiveValues = this.#primitiveMap.values() + const objectValues = this.#objectMap.values() + return (function* () { + yield* primitiveValues + yield* objectValues + })() + } + + private findMatchingObjectKey( + key: TKeyElement & object, + ): (TKeyElement & object) | undefined { + // Reference equality + if (this.#objectMap.has(key)) { + return key + } + + // Linear search for the matching key. + // This makes lookups in the trie O(NumberOfObjectKeys) + // but it also gives lookups in the trie like + // `map.get(['a', { obj: true }, 'c'])` the same semantics + // as `partialMatchKey` itself. + const keyArray = [key] + for (const candidateKey of this.#objectMap.keys()) { + if (partialMatchKey([candidateKey], keyArray)) { + return candidateKey + } + } + + return undefined + } +} + +type QueryKeyTrieNode = { + /** Element in the query key, QueryKey[number]. */ + key: (TKeyElement & Primitive) | (TKeyElement & object) /** - * - Entries who's TKey path leads exactly to this node, - * therefor their path is all primitives. - * ourPath = [p1, p2, p3, ...pN] - * theirPath = [p1, p2, p3, ...pN] + * Value stored at the end of the path leading to this node. + * This holds: `key[key.length - 1] === node.key` + * ```ts + * map.set(['a', 'b', 'c'], '123') + * // -> + * const root = { + * children: { + * 'a': { + * key: 'a', + * children: { + * 'b': { + * key: 'b', + * children: { + * 'c': { + * key: 'c', + * value: '123', + * insertionOrder: 0, + * } + * } + * } + * } + * } + * } + * } + * ``` */ - exact?: RefCountSet + value?: TValue /** - * - Entries who's path after this point contains non-primitive keys. - * Such entries cannot be looked up by value deeper in the trie. - * Implies their TKey path != this node's keyPath. - * ourPath = [p1, p2, p3, ...pN] - * theirPath = [p1, p2, p3, ...pN, nonPrimitive, ...] + * Insertion order of the *value* being stored in the trie. + * Unfortunately the natural iteration order of values in the trie does not + * match the insertion order, as expected from a map-like data structure. + * + * We need to track it explicitly. */ - nonPrimitiveSuffix?: RefCountSet - /** Child nodes storing entries who's TKey path is prefixed with this node's path. */ - children?: Map> + insertionOrder?: number + /** + * Children map to the next index element in the queryKey. + */ + children?: QueryKeyElementMap< + TKeyElement, + QueryKeyTrieNode + > +} + +type QueryKeyTrieNodeWithValue = { + key: (TKeyElement & Primitive) | (TKeyElement & object) + value: TValue + insertionOrder: number + children?: QueryKeyElementMap< + TKeyElement, + QueryKeyTrieNode + > +} + +/** + * We only consider a value to be stored in a node when insertionOrder is defined. + * This allows storing `undefined` as a value. + */ +function nodeHasValue( + node: QueryKeyTrieNode, +): node is QueryKeyTrieNodeWithValue { + return node.insertionOrder !== undefined } -/** Path length is always 1 greater than the key length, as it includes the root node. */ -function traverse( - root: MapTrieNode, - key: QueryKey, +/** + * Path length is always 1 greater than the key length, as it includes the root + * node. + * ```ts + * map.set(['a', 'b', 'c'], '123') + * const path = [root, n1, n2, n3] + * const key = ['a', 'b', 'c'] + */ +function traverse< + TKey extends QueryKey, + TValue, + TLookup extends QueryKeyTrieNode | undefined, +>( + root: QueryKeyTrieNode, + key: TKey, // May create a child node if needed lookup: ( - parent: MapTrieNode, - key: Primitive, - ) => MapTrieNode | undefined, -): Array> { - const path: Array> = [root] - let node: MapTrieNode | undefined = root - for (let i = 0; i < key.length && node; i++) { + parent: QueryKeyTrieNode, + key: (TKey[number] & Primitive) | (TKey[number] & object), + ) => TLookup, +): TLookup extends undefined ? undefined : Array { + const path: Array> = [root] + let node: QueryKeyTrieNode | undefined = root + // In hot code like this data structures, it is best to avoid creating + // Iterators with for-of loops. + // eslint-disable-next-line @typescript-eslint/prefer-for-of + for (let i = 0; i < key.length; i++) { const keyPart = key[i] - if (isPrimitive(keyPart)) { - node = lookup(node, keyPart) - } else { - node = undefined - } + node = lookup( + node, + keyPart as (TKey[number] & Primitive) | (TKey[number] & object), + ) if (node) { path.push(node) + } else { + return undefined as never } } - return path + return path as never } -function gcPath(path: Array>): void { - if (path.length === 0) { - return +function* iterateSubtreeValueNodes( + node: QueryKeyTrieNode, +): Generator, void, undefined> { + if (nodeHasValue(node)) { + yield node } - for (let i = path.length - 1; i >= 0; i--) { - const node = path[i] - if (!node) { - throw new Error('Should never occur (bug in MapTrie)') - } - if ( - node.exact?.size || - node.nonPrimitiveSuffix?.size || - node.children?.size - ) { - // Has data. Do not GC. - return - } + const children = node.children?.values() + if (!children) { + return + } - const parent = path[i - 1] - parent?.children?.delete(node.key) + for (const child of children) { + yield* iterateSubtreeValueNodes(child) } } -class MapTrieSet { - #root: MapTrieNode = { +class QueryKeyTrie { + #root: QueryKeyTrieNode = { key: undefined, } + // Provides relative insertion ordering between values in the trie. + #insertionOrder = 0 - add(key: TKey, value: TValue): void { + set(key: TKey, value: TValue): void { const path = traverse(this.#root, key, (parent, keyPart) => { - parent.children ??= new Map() + parent.children ??= new QueryKeyElementMap() let child = parent.children.get(keyPart) if (!child) { + // Note: insertionOrder is for values, not when nodes enter the trie. child = { key: keyPart } parent.children.set(keyPart, child) } return child }) - const lastPathNode = path[path.length - 1] - if (!lastPathNode) { - throw new Error('Should never occur (bug in MapTrie)') - } - if (key.length === path.length - 1) { - lastPathNode.exact ??= new RefCountSet() - lastPathNode.exact.add(value) - } else { - lastPathNode.nonPrimitiveSuffix ??= new RefCountSet() - lastPathNode.nonPrimitiveSuffix.add(value) + const lastNode = path[path.length - 1]! + if (!nodeHasValue(lastNode)) { + lastNode.insertionOrder = this.#insertionOrder++ } + lastNode.value = value } - remove(key: TKey, value: TValue): void { + delete(key: TKey): void { const path = traverse(this.#root, key, (parent, keyPart) => parent.children?.get(keyPart), ) - const lastPathNode = path[path.length - 1] - if (!lastPathNode) { - throw new Error('Should never occur (bug in MapTrie)') + if (!path) { + return + } + + const lastNode = path[path.length - 1]! + if (lastNode.insertionOrder === undefined) { + // No value stored at key. + return } - if (key.length === path.length - 1) { - lastPathNode.exact?.remove(value) - gcPath(path) - } else if (!isPrimitive(key[path.length - 1])) { - lastPathNode.nonPrimitiveSuffix?.remove(value) - gcPath(path) + + // Drop. + lastNode.value = undefined + lastNode.insertionOrder = undefined + + // GC nodes in path that are no longer needed. + for (let i = path.length - 1; i > 0; i--) { + const node = path[i]! + if (nodeHasValue(node) || node.children?.size) { + // Has data. Do not GC. + return + } + + const parent = path[i - 1] + parent?.children?.delete(node.key) } } @@ -263,64 +444,29 @@ class MapTrieSet { * Either the value has the same key and is all primitives, * Or the value's key is a suffix of the given key and contains a non-primitive key. */ - getByPrefix(key: TKey): Iterable | undefined { - let miss = false - const path = traverse(this.#root, key, (parent, keyPart) => { - const child = parent.children?.get(keyPart) - if (!child) { - miss = true - return undefined - } - return child - }) - // Failed to look up one of the primitive keys in the path. - // This means there's no match at all. - // Appears to be incorrectly reported by @typescript-eslint as always false :\ - // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition - if (miss) { + iteratePrefix(key: TKey): Array | undefined { + const path = traverse(this.#root, key, (parent, keyPart) => + parent.children?.get(keyPart), + ) + if (!path) { return undefined } - const lastNode = path[path.length - 1] - if (!lastNode) { - throw new Error('Should never occur (bug in MapTrie)') - } - - // If the `key` is all primitives then we need to recurse to find all values - // that match the prefix, as these values will be stored deeper in the trie. - // - // If the `key` contains a non-primitive part after the returned path, - // then all possible values that have the suffix are stored in this node. - const isPrimitivePath = path.length - 1 === key.length - if (!isPrimitivePath) { - return lastNode.nonPrimitiveSuffix?.[Symbol.iterator]() - } - - // See if we can avoid instantiating a generator - if ( - !lastNode.children && - (lastNode.exact || lastNode.nonPrimitiveSuffix) && - !(lastNode.exact && lastNode.nonPrimitiveSuffix) - ) { - return lastNode.exact ?? lastNode.nonPrimitiveSuffix + const lastNode = path[path.length - 1]! + if (!lastNode.children?.size) { + // No children - either return value if we have one, or nothing. + if (nodeHasValue(lastNode)) { + return [lastNode.value] + } + return undefined } - return (function* depthFirstPrefixIterator() { - const queue = [lastNode] - while (queue.length > 0) { - const node = queue.pop()! - if (node.exact) { - yield* node.exact - } - if (node.nonPrimitiveSuffix) { - yield* node.nonPrimitiveSuffix - } - if (node.children) { - const children = Array.from(node.children.values()).reverse() - children.forEach((child) => queue.push(child)) - } - } - })() + const subtreeInDepthFirstOrder = Array.from( + iterateSubtreeValueNodes(lastNode), + ) + return subtreeInDepthFirstOrder + .sort((a, b) => a.insertionOrder - b.insertionOrder) + .map((node) => node.value) } } @@ -328,7 +474,7 @@ class MapTrieSet { export class QueryCache extends Subscribable { #queries: QueryStore - #keyIndex = new MapTrieSet() + #keyIndex = new QueryKeyTrie() #knownHashFns = new RefCountSet>() constructor(public config: QueryCacheConfig = {}) { @@ -372,7 +518,7 @@ export class QueryCache extends Subscribable { add(query: Query): void { if (!this.#queries.has(query.queryHash)) { this.#queries.set(query.queryHash, query) - this.#keyIndex.add(query.queryKey, query) + this.#keyIndex.set(query.queryKey, query) const hashFn = query.options.queryKeyHashFn if (hashFn) { this.#knownHashFns.add(hashFn) @@ -393,7 +539,7 @@ export class QueryCache extends Subscribable { if (queryInMap === query) { this.#queries.delete(query.queryHash) - this.#keyIndex.remove(query.queryKey, query) + this.#keyIndex.delete(query.queryKey) const hashFn = query.options.queryKeyHashFn if (hashFn) { this.#knownHashFns.remove(hashFn) @@ -437,20 +583,14 @@ export class QueryCache extends Subscribable { return this.findExact(defaultedFilters) } - const candidates = this.#keyIndex.getByPrefix(defaultedFilters.queryKey) + const candidates = this.#keyIndex.iteratePrefix(defaultedFilters.queryKey) if (!candidates) { return undefined } - for (const query of candidates) { - if (matchQuery(defaultedFilters, query)) { - return query as unknown as - | Query - | undefined - } - } - - return undefined + return candidates.find((query) => matchQuery(defaultedFilters, query)) as + | Query + | undefined } findAll(filters: QueryFilters = {}): Array { @@ -460,9 +600,14 @@ export class QueryCache extends Subscribable { } if (filters.queryKey) { - const withPrefix = this.#keyIndex.getByPrefix(filters.queryKey) - const candidates = withPrefix ? Array.from(withPrefix) : [] - return candidates.filter((query) => matchQuery(filters, query)) + const candidates = this.#keyIndex.iteratePrefix(filters.queryKey) + if (!candidates) { + return [] + } + + return Object.keys(filters).length > 1 + ? candidates.filter((query) => matchQuery(filters, query)) + : candidates } const queries = this.getAll() From a247729e7a0498b694a5458810509beb5f960037 Mon Sep 17 00:00:00 2001 From: Jake Teton-Landis Date: Mon, 25 Aug 2025 18:06:11 -0400 Subject: [PATCH 7/8] new trie produces results in insert order, remove sort cop-out from test --- .../src/__tests__/queryCache.test.tsx | 25 +++++-------------- 1 file changed, 6 insertions(+), 19 deletions(-) diff --git a/packages/query-core/src/__tests__/queryCache.test.tsx b/packages/query-core/src/__tests__/queryCache.test.tsx index 42422a6510..5808e2ca98 100644 --- a/packages/query-core/src/__tests__/queryCache.test.tsx +++ b/packages/query-core/src/__tests__/queryCache.test.tsx @@ -348,25 +348,12 @@ describe('queryCache', () => { }) expect(predicateCallCount).toBe(results.length) - - const sortByHash = (a: Query | undefined, b: Query | undefined) => { - if (!a && !b) return 0 - if (!a) return -1 - if (!b) return 1 - if (a.queryHash < b.queryHash) return -1 - if (a.queryHash > b.queryHash) return 1 - return 0 - } - results.sort(sortByHash) - - expect(results).toEqual( - [ - exactMatchQuery, - primitiveSuffixQuery, - primitiveSuffixLength2Query, - matchingObjectSuffixQuery, - ].sort(sortByHash), - ) + expect(results).toEqual([ + exactMatchQuery, + primitiveSuffixQuery, + primitiveSuffixLength2Query, + matchingObjectSuffixQuery, + ]) }) test('should return all the queries when no filters are defined', async () => { From abda003dca25f5e24263ccb165a5ace42f715f61 Mon Sep 17 00:00:00 2001 From: Jake Teton-Landis Date: Mon, 25 Aug 2025 18:07:22 -0400 Subject: [PATCH 8/8] remove unused test import --- packages/query-core/src/__tests__/queryCache.test.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/query-core/src/__tests__/queryCache.test.tsx b/packages/query-core/src/__tests__/queryCache.test.tsx index 5808e2ca98..746959dbf0 100644 --- a/packages/query-core/src/__tests__/queryCache.test.tsx +++ b/packages/query-core/src/__tests__/queryCache.test.tsx @@ -1,6 +1,6 @@ import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest' import { queryKey, sleep } from '@tanstack/query-test-utils' -import { Query, QueryCache, QueryClient, QueryObserver } from '..' +import { QueryCache, QueryClient, QueryObserver } from '..' describe('queryCache', () => { let queryClient: QueryClient