From 71858776c189b3911752a24f932f5da3f640fd97 Mon Sep 17 00:00:00 2001
From: Shu Ding <g@shud.in>
Date: Mon, 5 Apr 2021 00:10:49 +0800
Subject: [PATCH] Bug fixes (#1096)

* fix type: allow getKey to be null

* fix leaking internal state when setSize on null key
---
 src/types.ts                   |  7 +++----
 src/use-swr-infinite.ts        | 10 ++++++++--
 test/use-swr-infinite.test.tsx | 28 ++++++++++++++++++++++++++++
 3 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/src/types.ts b/src/types.ts
index 9cda764f6..031749948 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -129,10 +129,9 @@ export interface SWRResponse<Data, Error> {
   isValidating: boolean
 }
 
-export type KeyLoader<Data = any> = (
-  index: number,
-  previousPageData: Data | null
-) => ValueKey
+export type KeyLoader<Data = any> =
+  | ((index: number, previousPageData: Data | null) => ValueKey)
+  | null
 
 /**
  * @deprecated `SWRInfiniteConfigInterface` will be renamed to `SWRInfiniteConfiguration`.
diff --git a/src/use-swr-infinite.ts b/src/use-swr-infinite.ts
index b1c9b2ed5..971125321 100644
--- a/src/use-swr-infinite.ts
+++ b/src/use-swr-infinite.ts
@@ -44,7 +44,7 @@ function useSWRInfinite<Data = any, Error = any>(
   // get the serialized key of the first page
   let firstPageKey: string | null = null
   try {
-    ;[firstPageKey] = cache.serializeKey(getKey(0, null))
+    ;[firstPageKey] = cache.serializeKey(getKey ? getKey(0, null) : null)
   } catch (err) {
     // not ready
   }
@@ -104,7 +104,7 @@ function useSWRInfinite<Data = any, Error = any>(
       let previousPageData = null
       for (let i = 0; i < pageSize; ++i) {
         const [pageKey, pageArgs] = cache.serializeKey(
-          getKey(i, previousPageData)
+          getKey ? getKey(i, previousPageData) : null
         )
 
         if (!pageKey) {
@@ -159,6 +159,9 @@ function useSWRInfinite<Data = any, Error = any>(
 
   const mutate = useCallback(
     (data: MutatorCallback, shouldRevalidate = true) => {
+      // It is possible that the key is still falsy.
+      if (!contextCacheKey) return undefined
+
       if (shouldRevalidate && typeof data !== 'undefined') {
         // we only revalidate the pages that are changed
         const originalData = dataRef.current
@@ -178,6 +181,9 @@ function useSWRInfinite<Data = any, Error = any>(
   // extend the SWR API
   const setSize = useCallback(
     (arg: number | ((size: number) => number)) => {
+      // It is possible that the key is still falsy.
+      if (!pageSizeCacheKey) return undefined
+
       let size
       if (typeof arg === 'function') {
         size = arg(resolvePageSize())
diff --git a/test/use-swr-infinite.test.tsx b/test/use-swr-infinite.test.tsx
index 76a173abf..abffd974c 100644
--- a/test/use-swr-infinite.test.tsx
+++ b/test/use-swr-infinite.test.tsx
@@ -543,4 +543,32 @@ describe('useSWRInfinite', () => {
     await screen.findByText('A:page-2-2')
     await screen.findByText('B:page-2-2')
   })
+
+  it('should support null as getKey', async () => {
+    function Page() {
+      const { data, setSize } = useSWRInfinite<string, string>(
+        null,
+        () => 'data'
+      )
+
+      return (
+        <div
+          onClick={() => {
+            // load next page
+            setSize(size => size + 1)
+          }}
+        >
+          data:{data || ''}
+        </div>
+      )
+    }
+
+    render(<Page />)
+    screen.getByText('data:')
+    await screen.findByText('data:')
+
+    // load next page
+    fireEvent.click(screen.getByText('data:'))
+    await screen.findByText('data:')
+  })
 })