From f83d4acc6d952d84f01db31c4118c6d848a0ad2c Mon Sep 17 00:00:00 2001
From: daishi <daishi@axlight.com>
Date: Tue, 27 Aug 2024 10:07:47 +0900
Subject: [PATCH] refactor useShallow

---
 src/react/shallow.ts | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/src/react/shallow.ts b/src/react/shallow.ts
index a0c4b3f0e9..0816f370f7 100644
--- a/src/react/shallow.ts
+++ b/src/react/shallow.ts
@@ -6,15 +6,17 @@ import { shallow } from '../vanilla/shallow.ts'
 
 const { useRef } = ReactExports
 
-export function useShallow<S, U>(selector: (state: S) => U): (state: S) => U {
-  const prev = useRef<U>()
+const sliceCache = new WeakMap<object, unknown>()
 
+export function useShallow<S, U>(selector: (state: S) => U): (state: S) => U {
+  const key = useRef({}).current
   return (state) => {
+    const prev = sliceCache.get(key) as U | undefined
     const next = selector(state)
-    return shallow(prev.current, next)
-      ? (prev.current as U)
-      : // It might not work with React Compiler
-        // eslint-disable-next-line react-compiler/react-compiler
-        (prev.current = next)
+    if (shallow(prev, next)) {
+      return prev as U
+    }
+    sliceCache.set(key, next)
+    return next
   }
 }