From 9b552e46a83efdeeba50392585a9852d96df456a Mon Sep 17 00:00:00 2001
From: Francesco Giordano <francescogior@users.noreply.github.com>
Date: Sun, 18 Jun 2023 22:39:36 +0200
Subject: [PATCH] Pass displayed data as second parameter of functional
 optimistic data (#2668)

---
 _internal/src/types.ts               |  4 +-
 _internal/src/utils/mutate.ts        |  2 +-
 test/use-swr-local-mutation.test.tsx | 84 ++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+), 2 deletions(-)

diff --git a/_internal/src/types.ts b/_internal/src/types.ts
index 3b399b392..5db1b6679 100644
--- a/_internal/src/types.ts
+++ b/_internal/src/types.ts
@@ -317,7 +317,9 @@ export type MutatorOptions<Data = any> = {
   populateCache?:
     | boolean
     | ((result: any, currentData: Data | undefined) => Data)
-  optimisticData?: Data | ((currentData?: Data) => Data)
+  optimisticData?:
+    | Data
+    | ((currentData: Data | undefined, displayedData: Data | undefined) => Data)
   rollbackOnError?: boolean | ((error: unknown) => boolean)
   throwOnError?: boolean
 }
diff --git a/_internal/src/utils/mutate.ts b/_internal/src/utils/mutate.ts
index 36331933d..2e86ff298 100644
--- a/_internal/src/utils/mutate.ts
+++ b/_internal/src/utils/mutate.ts
@@ -139,7 +139,7 @@ export async function internalMutate<Data>(
     // Do optimistic data update.
     if (hasOptimisticData) {
       optimisticData = isFunction(optimisticData)
-        ? optimisticData(committedData)
+        ? optimisticData(committedData, displayedData)
         : optimisticData
 
       // When we set optimistic data, backup the current committedData data in `_c`.
diff --git a/test/use-swr-local-mutation.test.tsx b/test/use-swr-local-mutation.test.tsx
index d1a26becf..7e4ced9e5 100644
--- a/test/use-swr-local-mutation.test.tsx
+++ b/test/use-swr-local-mutation.test.tsx
@@ -1123,6 +1123,90 @@ describe('useSWR - local mutation', () => {
     expect(renderedData).toEqual([undefined, 'loading', 'final'])
   })
 
+  it('should be able to use functional optimistic data config and use second param `displayedData` to keep UI consistent in slow networks', async () => {
+    const key1 = createKey()
+    const key2 = createKey()
+    let data1 = 0
+    let data2 = 0
+
+    function useOptimisticData1Mutate() {
+      const { mutate } = useSWRConfig()
+      return () => {
+        return mutate(key1, () => createResponse(data1++, { delay: 1000 }), {
+          optimisticData(currentData) {
+            return currentData + 1 // optimistic update current data
+          }
+        })
+      }
+    }
+
+    function useOptimisticData2Mutate() {
+      const { mutate } = useSWRConfig()
+      return () => {
+        return mutate(key2, () => createResponse(data2++, { delay: 1000 }), {
+          optimisticData(_, displayedData) {
+            return displayedData + 1 // optimistic update displayed data
+          }
+        })
+      }
+    }
+
+    function Page() {
+      const mutateWithOptimisticallyUpdatedCurrentData =
+        useOptimisticData1Mutate()
+      const mutateWithOptimisticallyUpdatedDisplayedData =
+        useOptimisticData2Mutate()
+      const { data: renderedData1 } = useSWR<number>(key1, () =>
+        createResponse(data1, { delay: 1000 })
+      )
+      const { data: renderedData2 } = useSWR<number>(key2, () =>
+        createResponse(data2, { delay: 1000 })
+      )
+
+      return (
+        <div>
+          <button onClick={mutateWithOptimisticallyUpdatedCurrentData}>
+            incrementCurrent
+          </button>
+          <button onClick={mutateWithOptimisticallyUpdatedDisplayedData}>
+            incrementDisplayed
+          </button>
+          <div>
+            data: <span data-testid="data1">{renderedData1}</span>
+          </div>
+          <div>
+            data: <span data-testid="data2">{renderedData2}</span>
+          </div>
+        </div>
+      )
+    }
+
+    renderWithConfig(<Page />)
+    await act(() => sleep(1000)) // Wait for initial data to load
+    fireEvent.click(screen.getByText('incrementCurrent'))
+    fireEvent.click(screen.getByText('incrementDisplayed'))
+    fireEvent.click(screen.getByText('incrementCurrent'))
+    fireEvent.click(screen.getByText('incrementDisplayed'))
+    const renderedData1 = parseInt(
+      (await screen.findByTestId('data1')).innerHTML,
+      10
+    )
+    const renderedData2 = Number((await screen.findByTestId('data2')).innerHTML)
+    await act(() => sleep(2000)) // Wait for revalidation roundtrip
+    const renderedRevalidatedData1 = Number(
+      (await screen.findByTestId('data1')).innerHTML
+    )
+    const renderedRevalidatedData2 = Number(
+      (await screen.findByTestId('data2')).innerHTML
+    )
+    expect(data1).toEqual(2)
+    expect(renderedData1).toEqual(1)
+    expect(renderedRevalidatedData1).toEqual(2)
+    expect(data2).toEqual(2)
+    expect(renderedData2).toEqual(2)
+    expect(renderedRevalidatedData2).toEqual(2)
+  })
+
   it('should prevent race conditions with optimistic UI', async () => {
     const key = createKey()
     const renderedData = []