feat: deprecate equalityFn and add createWithEqualityFn (#1945)

* feat: deprecate equalityFn and add createWithEqualityFn

* add link to deprecation message

package.json

@@ -65,6 +65,15 @@
       "module": "./esm/shallow.js",
       "default": "./shallow.js"
     },
+    "./traditional": {
+      "types": "./traditional.d.ts",
+      "import": {
+        "types": "./esm/traditional.d.mts",
+        "default": "./esm/traditional.mjs"
+      },
+      "module": "./esm/traditional.js",
+      "default": "./traditional.js"
+    },
     "./context": {
       "types": "./context.d.ts",
       "import": {
@@ -84,6 +93,7 @@
     "build:middleware": "rollup -c --config-middleware",
     "build:middleware:immer": "rollup -c --config-middleware_immer",
     "build:shallow": "rollup -c --config-shallow",
+    "build:traditional": "rollup -c --config-traditional",
     "build:context": "rollup -c --config-context",
     "postbuild": "yarn patch-d-ts && yarn copy && yarn patch-esm-ts",
     "prettier": "prettier \"*.{js,json,md}\" \"{examples,src,tests,docs}/**/*.{js,jsx,ts,tsx,md,mdx}\" --write",

src/context.ts

@@ -6,8 +6,9 @@ import {
   useRef,
 } from 'react'
 import type { ReactNode } from 'react'
-import { useStore } from 'zustand'
 import type { StoreApi } from 'zustand'
+// eslint-disable-next-line import/extensions
+import { useStoreWithEqualityFn } from 'zustand/traditional'
 
 type UseContextStore<S extends StoreApi<unknown>> = {
   (): ExtractState<S>
@@ -62,7 +63,7 @@ function createContext<S extends StoreApi<unknown>>() {
         'Seems like you have not used zustand provider as an ancestor.'
       )
     }
-    return useStore(
+    return useStoreWithEqualityFn(
       store,
       selector as (state: ExtractState<S>) => StateSlice,
       equalityFn

src/react.ts

@@ -27,10 +27,19 @@ export function useStore<S extends WithReact<StoreApi<unknown>>>(
   api: S
 ): ExtractState<S>
 
+export function useStore<S extends WithReact<StoreApi<unknown>>, U>(
+  api: S,
+  selector: (state: ExtractState<S>) => U
+): U
+
+/**
+ * @deprecated Use `useStoreWithEqualityFn` from 'zustand/traditional'
+ * https://github.com/pmndrs/zustand/discussions/1937
+ */
 export function useStore<S extends WithReact<StoreApi<unknown>>, U>(
   api: S,
   selector: (state: ExtractState<S>) => U,
-  equalityFn?: (a: U, b: U) => boolean
+  equalityFn: (a: U, b: U) => boolean
 ): U
 
 export function useStore<TState, StateSlice>(
@@ -38,6 +47,11 @@ export function useStore<TState, StateSlice>(
   selector: (state: TState) => StateSlice = api.getState as any,
   equalityFn?: (a: StateSlice, b: StateSlice) => boolean
 ) {
+  if (import.meta.env?.MODE !== 'production' && equalityFn) {
+    console.warn(
+      "[DEPRECATED] Use `createWithEqualityFn` from 'zustand/traditional'. https://github.com/pmndrs/zustand/discussions/1937"
+    )
+  }
   const slice = useSyncExternalStoreWithSelector(
     api.subscribe,
     api.getState,
@@ -51,9 +65,13 @@ export function useStore<TState, StateSlice>(
 
 export type UseBoundStore<S extends WithReact<ReadonlyStoreApi<unknown>>> = {
   (): ExtractState<S>
+  <U>(selector: (state: ExtractState<S>) => U): U
+  /**
+   * @deprecated Use `createWithEqualityFn` from 'zustand/traditional'
+   */
   <U>(
     selector: (state: ExtractState<S>) => U,
-    equals?: (a: U, b: U) => boolean
+    equalityFn: (a: U, b: U) => boolean
   ): U
 } & S
 

src/traditional.ts

@@ -0,0 +1,98 @@
+import { useDebugValue } from 'react'
+// import { useSyncExternalStoreWithSelector } from 'use-sync-external-store/shim/with-selector'
+// This doesn't work in ESM, because use-sync-external-store only exposes CJS.
+// See: https://github.com/pmndrs/valtio/issues/452
+// The following is a workaround until ESM is supported.
+// eslint-disable-next-line import/extensions
+import useSyncExternalStoreExports from 'use-sync-external-store/shim/with-selector'
+import { createStore } from './vanilla.ts'
+import type {
+  Mutate,
+  StateCreator,
+  StoreApi,
+  StoreMutatorIdentifier,
+} from './vanilla.ts'
+
+const { useSyncExternalStoreWithSelector } = useSyncExternalStoreExports
+
+type ExtractState<S> = S extends { getState: () => infer T } ? T : never
+
+type ReadonlyStoreApi<T> = Pick<StoreApi<T>, 'getState' | 'subscribe'>
+
+type WithReact<S extends ReadonlyStoreApi<unknown>> = S & {
+  getServerState?: () => ExtractState<S>
+}
+
+export function useStoreWithEqualityFn<S extends WithReact<StoreApi<unknown>>>(
+  api: S
+): ExtractState<S>
+
+export function useStoreWithEqualityFn<
+  S extends WithReact<StoreApi<unknown>>,
+  U
+>(
+  api: S,
+  selector: (state: ExtractState<S>) => U,
+  equalityFn?: (a: U, b: U) => boolean
+): U
+
+export function useStoreWithEqualityFn<TState, StateSlice>(
+  api: WithReact<StoreApi<TState>>,
+  selector: (state: TState) => StateSlice = api.getState as any,
+  equalityFn?: (a: StateSlice, b: StateSlice) => boolean
+) {
+  const slice = useSyncExternalStoreWithSelector(
+    api.subscribe,
+    api.getState,
+    api.getServerState || api.getState,
+    selector,
+    equalityFn
+  )
+  useDebugValue(slice)
+  return slice
+}
+
+export type UseBoundStoreWithEqualityFn<
+  S extends WithReact<ReadonlyStoreApi<unknown>>
+> = {
+  (): ExtractState<S>
+  <U>(
+    selector: (state: ExtractState<S>) => U,
+    equalityFn?: (a: U, b: U) => boolean
+  ): U
+} & S
+
+type CreateWithEqualityFn = {
+  <T, Mos extends [StoreMutatorIdentifier, unknown][] = []>(
+    initializer: StateCreator<T, [], Mos>,
+    defaultEqualityFn: <U>(a: U, b: U) => boolean
+  ): UseBoundStoreWithEqualityFn<Mutate<StoreApi<T>, Mos>>
+  <T>(): <Mos extends [StoreMutatorIdentifier, unknown][] = []>(
+    initializer: StateCreator<T, [], Mos>,
+    defaultEqualityFn: <U>(a: U, b: U) => boolean
+  ) => UseBoundStoreWithEqualityFn<Mutate<StoreApi<T>, Mos>>
+}
+
+const createWithEqualityFnImpl = <T>(
+  createState: StateCreator<T, [], []>,
+  defaultEqualityFn?: <U>(a: U, b: U) => boolean
+) => {
+  const api = createStore(createState)
+
+  const useBoundStoreWithEqualityFn: any = (
+    selector?: any,
+    equalityFn = defaultEqualityFn
+  ) => useStoreWithEqualityFn(api, selector, equalityFn)
+
+  Object.assign(useBoundStoreWithEqualityFn, api)
+
+  return useBoundStoreWithEqualityFn
+}
+
+export const createWithEqualityFn = (<T>(
+  createState: StateCreator<T, [], []> | undefined,
+  defaultEqualityFn?: <U>(a: U, b: U) => boolean
+) =>
+  createState
+    ? createWithEqualityFnImpl(createState, defaultEqualityFn)
+    : createWithEqualityFnImpl) as CreateWithEqualityFn

tests/basic.test.tsx

@@ -11,6 +11,7 @@ import ReactDOM from 'react-dom'
 import { afterEach, expect, it, vi } from 'vitest'
 import { create } from 'zustand'
 import type { StoreApi } from 'zustand'
+import { createWithEqualityFn } from 'zustand/traditional'
 
 const consoleError = console.error
 afterEach(() => {
@@ -89,7 +90,10 @@ it('uses the store with selectors', async () => {
 })
 
 it('uses the store with a selector and equality checker', async () => {
-  const useBoundStore = create(() => ({ item: { value: 0 } }))
+  const useBoundStore = createWithEqualityFn(
+    () => ({ item: { value: 0 } }),
+    Object.is
+  )
   const { setState } = useBoundStore
   let renderCount = 0
 
@@ -214,7 +218,10 @@ it('can update the selector', async () => {
 it('can update the equality checker', async () => {
   type State = { value: number }
   type Props = { equalityFn: (a: State, b: State) => boolean }
-  const useBoundStore = create<State>(() => ({ value: 0 }))
+  const useBoundStore = createWithEqualityFn<State>(
+    () => ({ value: 0 }),
+    Object.is
+  )
   const { setState } = useBoundStore
   const selector = (s: State) => s
 
@@ -258,7 +265,10 @@ it('can call useBoundStore with progressively more arguments', async () => {
     equalityFn?: (a: number, b: number) => boolean
   }
 
-  const useBoundStore = create<State>(() => ({ value: 0 }))
+  const useBoundStore = createWithEqualityFn<State>(
+    () => ({ value: 0 }),
+    Object.is
+  )
   const { setState } = useBoundStore
 
   let renderCount = 0
@@ -357,7 +367,7 @@ it('can throw an error in equality checker', async () => {
   type State = { value: string | number }
 
   const initialState: State = { value: 'foo' }
-  const useBoundStore = create(() => initialState)
+  const useBoundStore = createWithEqualityFn(() => initialState, Object.is)
   const { setState } = useBoundStore
   const selector = (s: State) => s
   const equalityFn = (a: State, b: State) =>