From f8a4d58dbd704fdddbfbf279c2c848580eb7db5d Mon Sep 17 00:00:00 2001
From: Daishi Kato <dai-shi@users.noreply.github.com>
Date: Sun, 21 Jan 2024 14:41:20 +0900
Subject: [PATCH] [v5]: refactor useMemoSelector (#2302)

* [v5]: refactor useMemoSelector

* add a test

* Revert "[v5]: refactor useMemoSelector"

This reverts commit b3c8b15586a270d12c335e566975021adf86c815.

* Revert "Revert "[v5]: refactor useMemoSelector""

This reverts commit 3c47301d23e18dffb7d72df36595f83570d15d08.
---
 src/react.ts         | 10 ++++------
 src/react/shallow.ts |  2 +-
 tests/basic.test.tsx | 17 +++++++++++++++++
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/react.ts b/src/react.ts
index b2ee3963..3e0dc16b 100644
--- a/src/react.ts
+++ b/src/react.ts
@@ -25,15 +25,13 @@ const useMemoSelector = <TState, StateSlice>(
   selector: (state: TState) => StateSlice,
 ) =>
   useMemo(() => {
-    let lastSlice: StateSlice
-    let lastState: TState
+    let prev: readonly [TState, StateSlice] | undefined
     return () => {
       const state = getState()
-      if (!Object.is(lastState, state)) {
-        lastSlice = selector(state)
-        lastState = state
+      if (!prev || !Object.is(prev[0], state)) {
+        prev = [state, selector(state)]
       }
-      return lastSlice
+      return prev[1]
     }
   }, [getState, selector])
 
diff --git a/src/react/shallow.ts b/src/react/shallow.ts
index 98e62dae..585b99b8 100644
--- a/src/react/shallow.ts
+++ b/src/react/shallow.ts
@@ -1,4 +1,4 @@
-// import { useDebugValue } from 'react'
+// import { useRef } from 'react'
 // That doesnt work in ESM, because React libs are CJS only.
 // The following is a workaround until ESM is supported.
 // eslint-disable-next-line import/extensions
diff --git a/tests/basic.test.tsx b/tests/basic.test.tsx
index 90fc0510..0bfca4d6 100644
--- a/tests/basic.test.tsx
+++ b/tests/basic.test.tsx
@@ -680,3 +680,20 @@ it('works with non-object state', async () => {
   fireEvent.click(getByText('button'))
   await findByText('count: 2')
 })
+
+it('works with "undefined" state', async () => {
+  const useUndefined = create(() => undefined)
+
+  const Component = () => {
+    const str = useUndefined((v) => v || 'undefined')
+    return <div>str: {str}</div>
+  }
+
+  const { findByText } = render(
+    <StrictMode>
+      <Component />
+    </StrictMode>,
+  )
+
+  await findByText('str: undefined')
+})