feat: add useShallow

src/shallow.ts

@@ -1,3 +1,5 @@
+import { useMemo, useRef } from 'react'
+
 export function shallow<T>(objA: T, objB: T) {
   if (Object.is(objA, objB)) {
     return true
@@ -59,3 +61,14 @@
   }
   return shallow(objA, objB)
 }) as typeof shallow
+
+export function useShallow<S, U>(selector: (state: S) => U): (state: S) => U {
+  const prev = useRef<U>()
+
+  return (state) => {
+    const next = selector(state)
+    return shallow(prev.current, next)
+      ? (prev.current as U)
+      : (prev.current = next)
+  }
+}

tests/shallow.test.tsx

@@ -1,6 +1,9 @@
+import { useState } from 'react'
+import { act, fireEvent, render, renderHook } from '@testing-library/react'
 import { describe, expect, it } from 'vitest'
 import { create } from 'zustand'
-import { shallow } from 'zustand/shallow'
+import { useStore } from 'zustand/react'
+import { shallow, useShallow } from 'zustand/shallow'
 
 describe('shallow', () => {
   it('compares primitive values', () => {
@@ -131,3 +134,126 @@ describe('unsupported cases', () => {
     ).not.toBe(false)
   })
 })
+
+describe('useShallow', () => {
+  const useTestShallowSimple = <S, U>(props: {
+    selector: (state: S) => U
+    state: S
+  }): { selectorOutput: U; useShallowOutput: U } => {
+    const useShallowOutput = useShallow(props.selector)(props.state)
+    return {
+      selectorOutput: props.selector(props.state),
+      useShallowOutput,
+    }
+  }
+
+  it('input and output selectors always return shallow equal values', () => {
+    const initialProps = {
+      selector: (state: Record<string, number>) => Object.keys(state),
+      state: { a: 1, b: 2 } as Record<string, number>,
+    }
+
+    const res = renderHook((props) => useTestShallowSimple(props), {
+      initialProps,
+    })
+
+    expect(res.result.current.selectorOutput).toEqual(
+      res.result.current.useShallowOutput
+    )
+
+    res.rerender({
+      state: { a: 1, b: 2, c: 3 },
+      selector: initialProps.selector,
+    })
+
+    expect(res.result.current.selectorOutput).toEqual(
+      res.result.current.useShallowOutput
+    )
+  })
+
+  it('returns the previously computed instance when possible', () => {
+    const initialProps = {
+      selector: Object.keys,
+      state: { a: 1, b: 2 } as Record<string, number>,
+    }
+
+    const res = renderHook((props) => useTestShallowSimple(props), {
+      initialProps,
+    })
+
+    const output1 = res.result.current.useShallowOutput
+
+    res.rerender({
+      state: initialProps.state,
+      selector: (state: Record<string, number>) => Object.keys(state), // New selector instance
+    })
+    const output2 = res.result.current.useShallowOutput
+
+    expect(output1).toBe(output2)
+    res.rerender(initialProps)
+
+    expect(res.result.current.useShallowOutput).toBe(output1)
+  })
+
+  it('only re-renders if selector output has changed according to shallow', () => {
+    let countRenders = 0
+    const store = create((): Record<string, unknown> => ({ a: 1, b: 2, c: 3 }))
+    const TestShallow = ({
+      selector = (state) => Object.keys(state).sort(),
+    }: {
+      selector?: (state: Record<string, unknown>) => string[]
+    }) => {
+      const output = useStore(store, useShallow(selector))
+
+      ++countRenders
+
+      return <div data-testid="test-shallow">{output.join(',')}</div>
+    }
+
+    expect(countRenders).toBe(0)
+    const res = render(<TestShallow />)
+
+    expect(countRenders).toBe(1)
+    expect(res.getByTestId('test-shallow').textContent).toBe('a,b,c')
+
+    act(() => {
+      store.setState({ a: 4 }) // This will not cause a re-render.
+    })
+
+    expect(countRenders).toBe(1)
+
+    act(() => {
+      store.setState({ d: 10 }) // This will cause a re-render.
+    })
+
+    expect(countRenders).toBe(2)
+    expect(res.getByTestId('test-shallow').textContent).toBe('a,b,c,d')
+  })
+
+  it('does not cause stale closure issues', () => {
+    const store = create((): Record<string, unknown> => ({ a: 1, b: 2, c: 3 }))
+    const TestShallowWithState = () => {
+      const [count, setCount] = useState(0)
+      const output = useStore(
+        store,
+        useShallow((state) => Object.keys(state).concat([count.toString()]))
+      )
+
+      return (
+        <div
+          data-testid="test-shallow"
+          onClick={() => setCount((prev) => ++prev)}>
+          {output.join(',')}
+        </div>
+      )
+    }
+
+    const res = render(<TestShallowWithState />)
+
+    expect(res.getByTestId('test-shallow').textContent).toBe('a,b,c,0')
+
+    fireEvent.click(res.getByTestId('test-shallow'))
+
+    expect(res.getByTestId('test-shallow').textContent).toBe('a,b,c,1')
+  })
+})