feat: support functional optimisticData (#1861)

* feat: support functional optimisticData

* Update src/types.ts

Co-authored-by: Shu Ding <[email protected]>

Author: huozhi

src/types.ts

@@ -145,7 +145,7 @@ export type MutatorCallback<Data = any> = (
 export type MutatorOptions<Data = any> = {
   revalidate?: boolean
   populateCache?: boolean | ((result: any, currentData: Data) => Data)
-  optimisticData?: Data
+  optimisticData?: Data | ((currentData?: Data) => Data)
   rollbackOnError?: boolean
 }
 

src/utils/mutate.ts

@@ -27,7 +27,7 @@ export const internalMutate = async <Data>(
     : options.populateCache
   const revalidate = options.revalidate !== false
   const rollbackOnError = options.rollbackOnError !== false
-  const optimisticData = options.optimisticData
+  const customOptimisticData = options.optimisticData
 
   // Serilaize key
   const [key, , keyInfo] = serialize(_key)
@@ -55,11 +55,14 @@ export const internalMutate = async <Data>(
   // Update global timestamps.
   const beforeMutationTs = getTimestamp()
   MUTATION[key] = [beforeMutationTs, 0]
-  const hasOptimisticData = !isUndefined(optimisticData)
+  const hasCustomOptimisticData = !isUndefined(customOptimisticData)
   const rollbackData = cache.get(key)
 
   // Do optimistic data update.
-  if (hasOptimisticData) {
+  if (hasCustomOptimisticData) {
+    const optimisticData = isFunction(customOptimisticData)
+      ? customOptimisticData(rollbackData)
+      : customOptimisticData
     cache.set(key, optimisticData)
     broadcastState(cache, key, optimisticData)
   }
@@ -88,7 +91,7 @@ export const internalMutate = async <Data>(
     if (beforeMutationTs !== MUTATION[key][0]) {
       if (error) throw error
       return data
-    } else if (error && hasOptimisticData && rollbackOnError) {
+    } else if (error && hasCustomOptimisticData && rollbackOnError) {
       // Rollback. Always populate the cache in this case but without
       // transforming the data.
       populateCache = true

test/use-swr-local-mutation.test.tsx

@@ -1038,6 +1038,72 @@ describe('useSWR - local mutation', () => {
     expect(renderedData).toEqual([undefined, 'foo', 'bar', 'baz', 'foo'])
   })
 
+  it('should support optimistic updates via functional `optimisticData`', async () => {
+    const key = createKey()
+    const renderedData = []
+    let mutate
+
+    function Page() {
+      const { data, mutate: boundMutate } = useSWR(key, () =>
+        createResponse('foo', { delay: 20 })
+      )
+      mutate = boundMutate
+      renderedData.push(data)
+      return <div>data: {String(data)}</div>
+    }
+
+    renderWithConfig(<Page />)
+    await screen.findByText('data: foo')
+
+    await act(() =>
+      mutate(createResponse('baz', { delay: 20 }), {
+        optimisticData: data => 'functional_' + data
+      })
+    )
+    await sleep(30)
+    expect(renderedData).toEqual([
+      undefined,
+      'foo',
+      'functional_foo',
+      'baz',
+      'foo'
+    ])
+  })
+
+  it('should be able use mutate to manipulate data via functional `optimisticData`', async () => {
+    const key = createKey()
+    const renderedData = []
+
+    function useOptimisticDataMutate(_key, data, fallback) {
+      const { mutate } = useSWRConfig()
+      return () => {
+        return mutate(_key, createResponse(data, { delay: 20 }), {
+          optimisticData() {
+            return fallback
+          }
+        })
+      }
+    }
+
+    function Page() {
+      const mutateWithOptData = useOptimisticDataMutate(key, 'final', 'loading')
+      const { data } = useSWR(key)
+      renderedData.push(data)
+      return (
+        <div>
+          <button onClick={() => mutateWithOptData()}>mutate</button>
+          <div>data: {String(data)}</div>
+        </div>
+      )
+    }
+
+    renderWithConfig(<Page />)
+    fireEvent.click(screen.getByText('mutate'))
+    await sleep(30)
+
+    expect(renderedData).toEqual([undefined, 'loading', 'final'])
+  })
+
   it('should rollback optimistic updates when mutation fails', async () => {
     const key = createKey()
     const renderedData = []
@@ -1152,7 +1218,7 @@ describe('useSWR - local mutation', () => {
       const { data, mutate } = useSWR(key, () => 'foo')
       mutatePage = () =>
         mutate(new Promise(res => setTimeout(() => res('baz'), 20)), {
-          optimisticData: 'bar',
+          optimisticData: () => 'bar',
           revalidate: false,
           populateCache: v => '!' + v
         })