Address spurious consistency check re-renders by using useLayoutEffec…

…t inside useSnapshot instead of useEffect

src/react.ts

@@ -2,6 +2,7 @@ import {
   useCallback,
   useDebugValue,
   useEffect,
+  useLayoutEffect,
   useMemo,
   useRef,
   useSyncExternalStore,
@@ -148,7 +149,7 @@ export function useSnapshot<T extends object>(
     () => snapshot(proxyObject),
   )
   inRender = false
-  useEffect(() => {
+  useLayoutEffect(() => {
     lastSnapshot.current = currSnapshot
   })
   if (import.meta.env?.MODE !== 'production') {

tests/basic.test.tsx

@@ -173,6 +173,118 @@ it('no extra re-renders (render func calls in non strict mode)', async () => {
   expect(renderFn2).lastCalledWith(2)
 })
 
+it('no extra re-renders with nested useSnapshot (render func calls in non strict mode)', async () => {
+  const obj = proxy({ childCount: 0, parentCount: 0 })
+
+  const childRenderFn = vi.fn()
+  const Child = () => {
+    const snap = useSnapshot(obj)
+    childRenderFn(snap.childCount)
+    return (
+      <>
+        <div>childCount: {snap.childCount}</div>
+        <button onClick={() => ++obj.childCount}>childButton</button>
+      </>
+    )
+  }
+
+  const parentRenderFn = vi.fn()
+  const Parent = () => {
+    const snap = useSnapshot(obj)
+    parentRenderFn(snap.parentCount)
+    return (
+      <>
+        <div>parentCount: {snap.parentCount}</div>
+        <button onClick={() => ++obj.parentCount}>parentButton</button>
+        <Child />
+      </>
+    )
+  }
+
+  const { getByText } = render(<Parent />)
+
+  await waitFor(() => {
+    getByText('childCount: 0')
+    getByText('parentCount: 0')
+  })
+
+  expect(childRenderFn).toBeCalledTimes(1)
+  expect(childRenderFn).lastCalledWith(0)
+  expect(parentRenderFn).toBeCalledTimes(1)
+  expect(parentRenderFn).lastCalledWith(0)
+
+  obj.parentCount += 1
+
+  await waitFor(() => {
+    getByText('childCount: 0')
+    getByText('parentCount: 1')
+  })
+
+  expect(childRenderFn).toBeCalledTimes(2)
+  expect(childRenderFn).lastCalledWith(0)
+  expect(parentRenderFn).toBeCalledTimes(2)
+  expect(parentRenderFn).lastCalledWith(1)
+})
+
+it('no extra re-renders with child useSnapshot component (render func calls in non strict mode)', async () => {
+  const obj = proxy({ childCount: 0, anotherValue: 0 })
+
+  const childRenderFn = vi.fn()
+  const Child = () => {
+    const snap = useSnapshot(obj)
+    childRenderFn(snap.childCount)
+    return (
+      <>
+        <div>childCount: {snap.childCount}</div>
+        <button onClick={() => ++obj.childCount}>childButton</button>
+      </>
+    )
+  }
+
+  const parentRenderFn = vi.fn()
+  const Parent = () => {
+    const [parentCount, setParentCount] = useState(0)
+
+    parentRenderFn(parentCount)
+
+    return (
+      <>
+        <div>parentCount: {parentCount}</div>
+        <button onClick={() => setParentCount((v) => v + 1)}>
+          parentButton
+        </button>
+        <Child />
+      </>
+    )
+  }
+
+  const { getByText } = render(<Parent />)
+
+  await waitFor(() => {
+    getByText('childCount: 0')
+    getByText('parentCount: 0')
+  })
+
+  expect(childRenderFn).toBeCalledTimes(1)
+  expect(childRenderFn).lastCalledWith(0)
+  expect(parentRenderFn).toBeCalledTimes(1)
+  expect(parentRenderFn).lastCalledWith(0)
+
+  obj.anotherValue += 1
+
+  fireEvent.click(getByText('parentButton'))
+
+  await waitFor(() => {
+    getByText('childCount: 0')
+    getByText('parentCount: 1')
+  })
+
+  expect(childRenderFn).toBeCalledTimes(2)
+  expect(childRenderFn).lastCalledWith(0)
+  expect(parentRenderFn).toBeCalledTimes(2)
+  expect(parentRenderFn).lastCalledWith(1)
+})
+
 it('object in object', async () => {
   const obj = proxy({ object: { count: 0 } })