From dd19510e0b7276d1bc30f20379d4087669633e02 Mon Sep 17 00:00:00 2001
From: Francois Best <github@francoisbest.com>
Date: Wed, 13 Nov 2024 21:12:16 +0100
Subject: [PATCH] fix: Handle dynamic default values in useQueryStates

---
 packages/nuqs/src/useQueryStates.ts | 54 +++++++++++++++++++++--------
 1 file changed, 40 insertions(+), 14 deletions(-)

diff --git a/packages/nuqs/src/useQueryStates.ts b/packages/nuqs/src/useQueryStates.ts
index c96f81012..6ece5c74f 100644
--- a/packages/nuqs/src/useQueryStates.ts
+++ b/packages/nuqs/src/useQueryStates.ts
@@ -40,6 +40,7 @@ export type Values<T extends UseQueryStatesKeysMap> = {
     ? NonNullable<ReturnType<T[K]['parse']>>
     : ReturnType<T[K]['parse']> | null
 }
+type NullableValues<T extends UseQueryStatesKeysMap> = Nullable<Values<T>>
 
 type UpdaterFn<T extends UseQueryStatesKeysMap> = (
   old: Values<T>
@@ -80,7 +81,7 @@ export function useQueryStates<KeyMap extends UseQueryStatesKeysMap>(
     urlKeys = defaultUrlKeys
   }: Partial<UseQueryStatesOptions<KeyMap>> = {}
 ): UseQueryStatesReturn<KeyMap> {
-  type V = Values<KeyMap>
+  type V = NullableValues<KeyMap>
   const stateKeys = Object.keys(keyMap).join(',')
   const resolvedUrlKeys = useMemo(
     () =>
@@ -99,6 +100,17 @@ export function useQueryStates<KeyMap extends UseQueryStatesKeysMap>(
   if (Object.keys(queryRef.current).length !== Object.keys(keyMap).length) {
     queryRef.current = Object.fromEntries(initialSearchParams?.entries() ?? [])
   }
+  const defaultValues = useMemo(
+    () =>
+      Object.fromEntries(
+        Object.keys(keyMap).map(key => [key, keyMap[key]!.defaultValue ?? null])
+      ) as Values<KeyMap>,
+    [
+      Object.values(keyMap)
+        .map(({ defaultValue }) => defaultValue)
+        .join(',')
+    ]
+  )
 
   const [internalState, setInternalState] = useState<V>(() => {
     const source = initialSearchParams ?? new URLSearchParams()
@@ -137,7 +149,7 @@ export function useQueryStates<KeyMap extends UseQueryStatesKeysMap>(
     }
     const handlers = Object.keys(keyMap).reduce(
       (handlers, stateKey) => {
-        handlers[stateKey as keyof V] = ({
+        handlers[stateKey as keyof KeyMap] = ({
           state,
           query
         }: CrossHookSyncPayload) => {
@@ -147,7 +159,7 @@ export function useQueryStates<KeyMap extends UseQueryStatesKeysMap>(
           // for the subsequent setState to pick it up.
           stateRef.current = {
             ...stateRef.current,
-            [stateKey as keyof V]: state ?? defaultValue ?? null
+            [stateKey as keyof KeyMap]: state ?? defaultValue ?? null
           }
           queryRef.current[urlKey] = query
           debug(
@@ -162,7 +174,7 @@ export function useQueryStates<KeyMap extends UseQueryStatesKeysMap>(
         }
         return handlers
       },
-      {} as Record<keyof V, (payload: CrossHookSyncPayload) => void>
+      {} as Record<keyof KeyMap, (payload: CrossHookSyncPayload) => void>
     )
 
     for (const stateKey of Object.keys(keyMap)) {
@@ -183,7 +195,7 @@ export function useQueryStates<KeyMap extends UseQueryStatesKeysMap>(
     (stateUpdater, callOptions = {}) => {
       const newState: Partial<Nullable<KeyMap>> =
         typeof stateUpdater === 'function'
-          ? stateUpdater(stateRef.current)
+          ? stateUpdater(applyDefaultValues(stateRef.current, defaultValues))
           : stateUpdater === null
             ? (Object.fromEntries(
                 Object.keys(keyMap).map(key => [key, null])
@@ -241,10 +253,16 @@ export function useQueryStates<KeyMap extends UseQueryStatesKeysMap>(
       startTransition,
       resolvedUrlKeys,
       updateUrl,
-      rateLimitFactor
+      rateLimitFactor,
+      defaultValues
     ]
   )
-  return [internalState, update]
+
+  const outputState = useMemo(
+    () => applyDefaultValues(internalState, defaultValues),
+    [internalState, defaultValues]
+  )
+  return [outputState, update]
 }
 
 // --
@@ -254,26 +272,34 @@ function parseMap<KeyMap extends UseQueryStatesKeysMap>(
   urlKeys: Partial<Record<keyof KeyMap, string>>,
   searchParams: URLSearchParams,
   cachedQuery?: Record<string, string | null>,
-  cachedState?: Values<KeyMap>
-) {
+  cachedState?: NullableValues<KeyMap>
+): NullableValues<KeyMap> {
   return Object.keys(keyMap).reduce((obj, stateKey) => {
     const urlKey = urlKeys?.[stateKey] ?? stateKey
-    const { defaultValue, parse } = keyMap[stateKey]!
+    const { parse } = keyMap[stateKey]!
     const queuedQuery = getQueuedValue(urlKey)
     const query =
       queuedQuery === undefined
         ? (searchParams?.get(urlKey) ?? null)
         : queuedQuery
     if (cachedQuery && cachedState && cachedQuery[urlKey] === query) {
-      obj[stateKey as keyof KeyMap] =
-        cachedState[stateKey] ?? defaultValue ?? null
+      obj[stateKey as keyof KeyMap] = cachedState[stateKey] ?? null
       return obj
     }
     const value = query === null ? null : safeParse(parse, query, stateKey)
-    obj[stateKey as keyof KeyMap] = value ?? defaultValue ?? null
+    obj[stateKey as keyof KeyMap] = value ?? null
     if (cachedQuery) {
       cachedQuery[urlKey] = query
     }
     return obj
-  }, {} as Values<KeyMap>)
+  }, {} as NullableValues<KeyMap>)
+}
+
+function applyDefaultValues<KeyMap extends UseQueryStatesKeysMap>(
+  state: NullableValues<KeyMap>,
+  defaults: Partial<Values<KeyMap>>
+) {
+  return Object.fromEntries(
+    Object.keys(state).map(key => [key, state[key] ?? defaults[key] ?? null])
+  ) as Values<KeyMap>
 }