feat: allow array of keys

src/data-fetching-store.ts

@@ -3,13 +3,15 @@ import {
   type Ref,
   shallowReactive,
   getCurrentScope,
-  ShallowRef,
+  type ShallowRef,
   ref,
-  ComputedRef,
+  type ComputedRef,
   computed,
   triggerRef,
 } from 'vue'
 import type { UseQueryOptionsWithDefaults, UseQueryKey } from './use-query'
+import { type _MaybeArray, stringifyFlatObject } from './utils'
+import { TreeMapNode } from './tree-map'
 
 export type UseQueryStatus = 'pending' | 'error' | 'success'
 
@@ -70,26 +72,24 @@ export interface UseQueryEntry<TResult = unknown, TError = Error>
 
 export const useDataFetchingStore = defineStore('PiniaColada', () => {
   const entryStateRegistry = shallowReactive(
-    new Map<UseQueryKey, UseQueryStateEntry>()
+    new TreeMapNode<UseQueryStateEntry>()
   )
   // these are not reactive as they are mostly functions
-  const entryPropertiesRegistry = new Map<
-    UseQueryKey,
-    UseQueryPropertiesEntry
-  >()
+  const entryPropertiesRegistry = new TreeMapNode<UseQueryPropertiesEntry>()
 
   // FIXME: start from here: replace properties entry with a QueryEntry that is created when needed and contains all the needed part, included functions
 
   // this allows use to attach reactive effects to the scope later on
   const scope = getCurrentScope()!
 
   function ensureEntry<TResult = unknown, TError = Error>(
-    key: UseQueryKey,
+    keyRaw: UseQueryKey[],
     { fetcher, initialData, staleTime }: UseQueryOptionsWithDefaults<TResult>
   ): UseQueryEntry<TResult, TError> {
+    const key = keyRaw.map(stringifyFlatObject)
     // ensure the state
     console.log('⚙️ Ensuring entry', key)
-    if (!entryStateRegistry.has(key)) {
+    if (!entryStateRegistry.get(key)) {
       entryStateRegistry.set(
         key,
         scope.run(() => {
@@ -108,7 +108,7 @@ export const useDataFetchingStore = defineStore('PiniaColada', () => {
 
     // TODO: these needs to be created client side. Should probably be a class for better memory
 
-    if (!entryPropertiesRegistry.has(key)) {
+    if (!entryPropertiesRegistry.get(key)) {
       const propertiesEntry: UseQueryPropertiesEntry<TResult, TError> = {
         pending: null,
         previous: null,
@@ -199,8 +199,8 @@ export const useDataFetchingStore = defineStore('PiniaColada', () => {
    * @param key - the key of the query to invalidate
    * @param refetch - whether to force a refresh of the data
    */
-  function invalidateEntry(key: UseQueryKey, refetch = false) {
-    const entry = entryPropertiesRegistry.get(key)
+  function invalidateEntry(key: UseQueryKey[], refetch = false) {
+    const entry = entryPropertiesRegistry.get(key.map(stringifyFlatObject))
 
     // nothing to invalidate
     if (!entry) {
@@ -221,10 +221,10 @@ export const useDataFetchingStore = defineStore('PiniaColada', () => {
   }
 
   function setEntryData<TResult = unknown>(
-    key: UseQueryKey,
+    key: UseQueryKey[],
     data: TResult | ((data: Ref<TResult | undefined>) => void)
   ) {
-    const entry = entryStateRegistry.get(key) as
+    const entry = entryStateRegistry.get(key.map(stringifyFlatObject)) as
       | UseQueryStateEntry<TResult>
       | undefined
     if (!entry) {
@@ -242,8 +242,8 @@ export const useDataFetchingStore = defineStore('PiniaColada', () => {
     entry.error.value = null
   }
 
-  function prefetch(key: UseQueryKey) {
-    const entry = entryPropertiesRegistry.get(key)
+  function prefetch(key: UseQueryKey[]) {
+    const entry = entryPropertiesRegistry.get(key.map(stringifyFlatObject))
     if (!entry) {
       console.warn(
         `⚠️ trying to prefetch "${String(key)}" but it's not in the registry`

src/tree-map.spec.ts

@@ -1,9 +1,9 @@
 import { describe, expect, it } from 'vitest'
-import { EntryNode } from './tree-map'
+import { TreeMapNode, entryNodeSize, logTree } from './tree-map'
 
 describe('tree-map', () => {
   it('.set and .get', () => {
-    const tree = new EntryNode<string>()
+    const tree = new TreeMapNode<string>()
     tree.set(['a'], 'a')
     expect(tree.get(['a'])).toBe('a')
     tree.set(['a', 'b'], 'ab')
@@ -24,7 +24,7 @@ describe('tree-map', () => {
   })
 
   it('.delete', () => {
-    const tree = new EntryNode<string>()
+    const tree = new TreeMapNode<string>()
     tree.set(['a', 'b', 'c'], 'abc')
     tree.set(['a', 'b', 'd'], 'abd')
     tree.delete(['a', 'b', 'c'])
@@ -52,4 +52,31 @@ describe('tree-map', () => {
     expect(tree.get(['f'])).toBe(undefined)
     expect(tree.get(['a', 'k'])).toBe('ak')
   })
+
+  it('entryNodeSize', () => {
+    const tree = new TreeMapNode<string>()
+    tree.set(['a', 'b', 'c'], 'abc')
+    expect(entryNodeSize(tree)).toBe(3)
+    tree.set(['a', 'b', 'd'], 'abd')
+    expect(entryNodeSize(tree)).toBe(4)
+    tree.set(['a', 'b', 'h'], 'abh')
+    expect(entryNodeSize(tree)).toBe(5)
+    tree.set(['a', 'b', 'h'], 'abh2')
+    expect(entryNodeSize(tree)).toBe(5)
+    tree.set(['a', 'e'], 'ae')
+    expect(entryNodeSize(tree)).toBe(6)
+    tree.set(['a', 'k'], 'ak')
+    expect(entryNodeSize(tree)).toBe(7)
+    tree.set(['f'], 'f')
+    expect(entryNodeSize(tree)).toBe(8)
+    tree.set(['f'], 'f2')
+    expect(entryNodeSize(tree)).toBe(8)
+    tree.set(['g', 'h'], 'gh')
+    expect(entryNodeSize(tree)).toBe(10)
+
+    const t2 = new TreeMapNode<unknown>()
+    t2.set(['todos', 2], {})
+    t2.set(['todos', 4], {})
+    expect(entryNodeSize(t2)).toBe(3)
+  })
 })

src/tree-map.ts

@@ -1,8 +1,14 @@
-export type EntryNodeKey = string | number
+import { type _JSONPrimitive } from './utils'
 
-export class EntryNode<T> {
+export type EntryNodeKey = _JSONPrimitive
+
+/**
+ * Internal data structure used to store the data of `useQuery()`.
+ * @internal
+ */
+export class TreeMapNode<T = unknown> {
   value: T | undefined
-  children = new Map<string | number, EntryNode<T>>()
+  children = new Map<EntryNodeKey, TreeMapNode<T>>()
 
   constructor()
   constructor(keys: EntryNodeKey[], value: T)
@@ -24,11 +30,11 @@ export class EntryNode<T> {
     } else {
       // this.children ??= new Map<EntryNodeKey,
       const [top, ...otherKeys] = keys
-      const node: EntryNode<T> | undefined = this.children.get(top)
+      const node: TreeMapNode<T> | undefined = this.children.get(top)
       if (node) {
         node.set(otherKeys, value)
       } else {
-        this.children.set(top, new EntryNode(otherKeys, value))
+        this.children.set(top, new TreeMapNode(otherKeys, value))
       }
     }
   }
@@ -61,3 +67,69 @@ export class EntryNode<T> {
     }
   }
 }
+
+// Below are debugging internals
+
+/**
+ * Calculates the size of the node and all its children. Used in tests.
+ *
+ * @internal
+ * @param node - The node to calculate the size of
+ * @returns The size of the node and all its children
+ */
+export function entryNodeSize(node: TreeMapNode): number {
+  return (
+    node.children.size +
+    [...node.children.values()].reduce(
+      (acc, child) => acc + entryNodeSize(child),
+      0
+    )
+  )
+}
+
+export function logTree(
+  tree: TreeMapNode,
+  log: (str: string) => any = console.log
+) {
+  log(printTreeMap(tree))
+}
+
+const MAX_LEVEL = 1000
+function printTreeMap(
+  tree: TreeMapNode | TreeMapNode['children'],
+  level = 0,
+  parentPre = '',
+  treeStr = ''
+): string {
+  // end of recursion
+  if (typeof tree !== 'object' || level >= MAX_LEVEL) return ''
+
+  if (tree instanceof Map) {
+    const total = tree.size
+    let index = 0
+    for (const [key, child] of tree) {
+      const hasNext = index++ < total - 1
+      const { children } = child
+
+      treeStr += `${`${parentPre}${hasNext ? '├' : '└'}${
+        '─' + (children.size > 0 ? '┬' : '')
+      } `}${key}${child.value != null ? ' · ' + String(child.value) : ''}\n`
+
+      if (children) {
+        treeStr += printTreeMap(
+          children,
+          level + 1,
+          `${parentPre}${hasNext ? '│' : ' '} `
+        )
+      }
+    }
+  } else {
+    const children = tree.children
+    treeStr = `${String(tree.value ?? '<root>')}\n`
+    if (children) {
+      treeStr += printTreeMap(children, level + 1)
+    }
+  }
+
+  return treeStr
+}

src/use-mutation.test-d.ts

@@ -8,7 +8,7 @@ it('types the parameters for the key', () => {
       expectTypeOf(one).toBeString()
       expectTypeOf(two).toBeNumber()
       expectTypeOf(result).toEqualTypeOf<{ name: string }>()
-      return 'foo'
+      return ['foo']
     },
   })
 })

src/use-mutation.ts

@@ -4,8 +4,8 @@ import { type UseQueryKey } from './use-query'
 import { type _MaybeArray, toArray } from './utils'
 
 type _MutatorKeys<TParams extends readonly any[], TResult> =
-  | UseQueryKey[]
-  | ((result: TResult, ...args: TParams) => _MaybeArray<UseQueryKey>)
+  | _MaybeArray<UseQueryKey>[]
+  | ((result: TResult, ...args: TParams) => _MaybeArray<UseQueryKey>[])
 
 export interface UseMutationOptions<
   TResult = unknown,
@@ -16,8 +16,9 @@ export interface UseMutationOptions<
    */
   mutator: (...args: TParams) => Promise<TResult>
 
+  // TODO: move this to a plugin that calls invalidateEntry()
   /**
-   * Keys to invalidate for useQuery to trigger their refetch.
+   * Keys to invalidate if the mutation succeeds so that `useQuery()` refetch if used.
    */
   keys?: _MutatorKeys<TParams, TResult>
 }
@@ -62,11 +63,11 @@ export function useMutation<
         if (pendingPromise === promise) {
           data.value = _data
           if (options.keys) {
-            const keys = toArray(
+            const keys = (
               typeof options.keys === 'function'
                 ? options.keys(_data, ...args)
                 : options.keys
-            )
+            ).map(toArray)
             for (const key of keys) {
               store.invalidateEntry(key, true)
             }

src/use-query.spec.ts

@@ -5,6 +5,8 @@ import { createPinia } from 'pinia'
 import { defineComponent } from 'vue'
 import { GlobalMountOptions } from 'node_modules/@vue/test-utils/dist/types'
 import { delay, runTimers } from '../test/utils'
+import { useDataFetchingStore } from './data-fetching-store'
+import { entryNodeSize } from './tree-map'
 
 describe('useQuery', () => {
   beforeEach(() => {
@@ -223,4 +225,51 @@ describe('useQuery', () => {
       expect(wrapper.vm.data).toBe(42)
     })
   })
+
+  describe('shared state', () => {
+    it('reuses the same state for the same key', async () => {
+      const pinia = createPinia()
+      mountSimple({ key: 'todos' }, { plugins: [pinia] })
+      mountSimple({ key: ['todos'] }, { plugins: [pinia] })
+      await runTimers()
+
+      const cacheClient = useDataFetchingStore()
+      expect(entryNodeSize(cacheClient.entryStateRegistry)).toBe(1)
+
+      mountSimple({ key: ['todos', 2] }, { plugins: [pinia] })
+      await runTimers()
+
+      expect(entryNodeSize(cacheClient.entryStateRegistry)).toBe(2)
+    })
+
+    it('populates the entry registry', async () => {
+      const pinia = createPinia()
+      mountSimple({ key: ['todos', 5] }, { plugins: [pinia] })
+      mountSimple({ key: ['todos', 2] }, { plugins: [pinia] })
+      await runTimers()
+
+      const cacheClient = useDataFetchingStore()
+      expect(entryNodeSize(cacheClient.entryStateRegistry)).toBe(3)
+    })
+
+    it('order in object keys does not matter', async () => {
+      const pinia = createPinia()
+      mountSimple(
+        { key: ['todos', { id: 5, a: true, b: 'hello' }] },
+        { plugins: [pinia] }
+      )
+      mountSimple(
+        { key: ['todos', { a: true, id: 5, b: 'hello' }] },
+        { plugins: [pinia] }
+      )
+      mountSimple(
+        { key: ['todos', { id: 5, a: true, b: 'hello' }] },
+        { plugins: [pinia] }
+      )
+      await runTimers()
+
+      const cacheClient = useDataFetchingStore()
+      expect(entryNodeSize(cacheClient.entryStateRegistry)).toBe(2)
+    })
+  })
 })

src/use-query.test-d.ts

@@ -55,3 +55,27 @@ it('expects an async fetcher', () => {
     }).data
   )
 })
+
+it('can use a function as a key', () => {
+  useQuery({
+    fetcher: async () => 42,
+    key: ['todos'],
+  })
+
+  useQuery({
+    fetcher: async () => 42,
+    key: ['todos', '2', 2],
+  })
+})
+
+it('can use objects in keys', () => {
+  useQuery({
+    fetcher: async () => 42,
+    key: { id: 1 },
+  })
+
+  useQuery({
+    fetcher: async () => 42,
+    key: ['todos', { id: 1, a: true, b: 'hello' }, 5, true],
+  })
+})