Repository: pmndrs/zustand (ts)

Issue:
feature request: option for useStore.subscribe fire once on initialisation
Currently `useStore.subscribe` only fires when the subscribed part of the store changes. It would be useful if there was an option to force the function to fire once upon initialisation. Maybe I'm missing something, here's an example:

```javascript
useEffect(() => {
    // I don't want to have to repeat myself up here just for initialisation
    const { position, rotation } = getSceneObject(props.id)(useStore.getState())
    target.current.position.set(position[0], position[1] - 3, position[2])
    plane.current.position.set(...position)
    plane.current.rotation.set(rotation[0], Math.PI, 0)

    return useStore.subscribe(
      ({ position, rotation }: SceneObject) => {
        target.current.position.set(position[0], position[1] - 3, position[2])
        plane.current.position.set(...position)
        plane.current.rotation.set(rotation[0], Math.PI, 0)
      },
      (state) => getSceneObject(props.id)(state)
    )
  }, [props.id])
````
Missing type inference when calling store.subscribe with a selector
Type inference breaks when using `store.subscribe` with a selector.

```js
import createStore from "zustand/vanilla";

const store = createStore(() => ({
  hello: "world"
}));

store.subscribe((state) => {
  // 👍 state type is inferred here
  console.log(state.hello);
});

store.subscribe(
  (state) => {
    // 💥 state type "unknown" here
    console.log(state.hello);
  },
  ({ hello }) => ({ hello })
);
```

Here's a [Codesandbox](https://codesandbox.io/s/thirsty-kirch-yy5mh).

Is this a known limitation?

Pull request description:
feat(middleware): subscribeWithSelector middleware
deprecating the core subscribeWithSelector feature.
close #475 
close #555

Patch:
diff --git a/readme.md b/readme.md
index a813c70..9113088 100644
--- a/readme.md
+++ b/readme.md
@@ -177,19 +177,10 @@ const useStore = create(() => ({ paw: true, snout: true, fur: true }))
 const paw = useStore.getState().paw
 // Listening to all changes, fires synchronously on every change
 const unsub1 = useStore.subscribe(console.log)
-// Listening to selected changes, in this case when "paw" changes
-const unsub2 = useStore.subscribe(console.log, state => state.paw)
-// Subscribe also supports an optional equality function
-const unsub3 = useStore.subscribe(console.log, state => [state.paw, state.fur], shallow)
-// Subscribe also exposes the previous value
-const unsub4 = useStore.subscribe((paw, previousPaw) => console.log(paw, previousPaw), state => state.paw)
 // Updating state, will trigger listeners
 useStore.setState({ paw: false })
 // Unsubscribe listeners
 unsub1()
-unsub2()
-unsub3()
-unsub4()
 // Destroying the store (removing all listeners)
 useStore.destroy()
 
@@ -198,6 +189,30 @@ function Component() {
   const paw = useStore(state => state.paw)
 ```
 
+### Using subscribe with selector
+
+If you need to subscribe with selector,
+`subscribeWithSelector` middleware will help.
+
+With this middleware `subscribe` accepts an additional signature:
+```ts
+subscribe(selector, callback, options?: { equalityFn, fireImmediately }): Unsubscribe
+```
+
+```js
+import { subscribeWithSelector } from 'zustand/middleware'
+const useStore = create(subscribeWithSelector(() => ({ paw: true, snout: true, fur: true })))
+
+// Listening to selected changes, in this case when "paw" changes
+const unsub2 = useStore.subscribe(state => state.paw, console.log)
+// Subscribe also exposes the previous value
+const unsub3 = useStore.subscribe(state => state.paw, (paw, previousPaw) => console.log(paw, previousPaw))
+// Subscribe also supports an optional equality function
+const unsub4 = useStore.subscribe(state => [state.paw, state.fur], console.log, { equalityFn: shallow })
+// Subscribe and fire immediately
+const unsub5 = useStore.subscribe(state => state.paw, console.log, { fireImmediately: true })
+```
+
 ## Using zustand without React
 
 Zustands core can be imported and used without the React dependency. The only difference is that the create function does not return a hook, but the api utilities.
diff --git a/src/middleware.ts b/src/middleware.ts
index 56397d0..0a257c0 100644
--- a/src/middleware.ts
+++ b/src/middleware.ts
@@ -1,10 +1,15 @@
 import {
+  EqualityChecker,
   GetState,
   PartialState,
   SetState,
   State,
   StateCreator,
+  StateListener,
+  StateSelector,
+  StateSliceListener,
   StoreApi,
+  Subscribe,
 } from './vanilla'
 
 const DEVTOOLS = Symbol()
@@ -192,6 +197,104 @@ export const devtools =
     return initialState
   }
 
+type SubscribeWithSelector<T extends State> = {
+  (listener: StateListener<T>): () => void
+  <StateSlice>(
+    selector: StateSelector<T, StateSlice>,
+    listener: StateSliceListener<StateSlice>,
+    options?: {
+      equalityFn?: EqualityChecker<StateSlice>
+      fireImmediately?: boolean
+    }
+  ): () => void
+}
+
+export function subscribeWithSelector<S extends State>(
+  fn: (
+    set: SetState<S>,
+    get: GetState<S>,
+    api: Omit<StoreApi<S>, 'subscribe'> & {
+      subscribe: SubscribeWithSelector<S>
+      subscribeWithSelectorEnabled: true
+    }
+  ) => S
+): (
+  set: SetState<S>,
+  get: GetState<S>,
+  api: Omit<StoreApi<S>, 'subscribe'> & {
+    subscribe: SubscribeWithSelector<S>
+    subscribeWithSelectorEnabled: true
+  }
+) => S
+
+export function subscribeWithSelector<
+  S extends State,
+  CustomSetState extends SetState<S>
+>(
+  fn: (
+    set: CustomSetState,
+    get: GetState<S>,
+    api: Omit<StoreApi<S>, 'subscribe'> & {
+      subscribe: SubscribeWithSelector<S>
+      subscribeWithSelectorEnabled: true
+    }
+  ) => S
+): (
+  set: CustomSetState,
+  get: GetState<S>,
+  api: Omit<StoreApi<S>, 'subscribe'> & {
+    subscribe: SubscribeWithSelector<S>
+    subscribeWithSelectorEnabled: true
+  }
+) => S
+
+export function subscribeWithSelector<
+  S extends State,
+  CustomSetState extends SetState<S>,
+  CustomGetState extends GetState<S>,
+  CustomStoreApi extends Omit<StoreApi<S>, 'subscribe'> & {
+    subscribe: SubscribeWithSelector<S>
+    subscribeWithSelectorEnabled: true
+  }
+>(
+  fn: (set: CustomSetState, get: CustomGetState, api: CustomStoreApi) => S
+): (set: CustomSetState, get: CustomGetState, api: CustomStoreApi) => S
+
+export function subscribeWithSelector<
+  S extends State,
+  CustomSetState extends SetState<S>,
+  CustomGetState extends GetState<S>,
+  CustomStoreApi extends Omit<StoreApi<S>, 'subscribe'> & {
+    subscribe: SubscribeWithSelector<S>
+    subscribeWithSelectorEnabled: true
+  }
+>(fn: (set: CustomSetState, get: CustomGetState, api: CustomStoreApi) => S) {
+  return (set: CustomSetState, get: CustomGetState, api: CustomStoreApi): S => {
+    const origSubscribe = api.subscribe as Subscribe<S>
+    api.subscribe = ((selector: any, optListener: any, options: any) => {
+      let listener: StateListener<S> = selector // if no selector
+      if (optListener) {
+        const equalityFn = options?.equalityFn || Object.is
+        let currentSlice = selector(api.getState())
+        listener = (state) => {
+          const nextSlice = selector(state)
+          if (!equalityFn(currentSlice, nextSlice)) {
+            const previousSlice = currentSlice
+            optListener((currentSlice = nextSlice), previousSlice)
+          }
+        }
+        if (options?.fireImmediately) {
+          optListener(currentSlice, currentSlice)
+        }
+      }
+      return origSubscribe(listener)
+    }) as SubscribeWithSelector<S>
+    api.subscribeWithSelectorEnabled = true
+    const initialState = fn(set, get, api)
+    return initialState
+  }
+}
+
 type Combine<T, U> = Omit<T, keyof U> & U
 export const combine =
   <
diff --git a/src/vanilla.ts b/src/vanilla.ts
index c8a1245..f354432 100644
--- a/src/vanilla.ts
+++ b/src/vanilla.ts
@@ -16,6 +16,9 @@ export type StateListener<T> = (state: T, previousState: T) => void
 export type StateSliceListener<T> = (slice: T, previousSlice: T) => void
 export interface Subscribe<T extends State> {
   (listener: StateListener<T>): () => void
+  /**
+   * @deprecated Please use `subscribeWithSelector` middleware
+   */
   <StateSlice>(
     listener: StateSliceListener<StateSlice>,
     selector?: StateSelector<T, StateSlice>,
@@ -88,6 +91,7 @@ export default function create<
     selector: StateSelector<TState, StateSlice> = getState as any,
     equalityFn: EqualityChecker<StateSlice> = Object.is
   ) => {
+    console.warn('[DEPRECATED] Please use `subscribeWithSelector` middleware')
     let currentSlice: StateSlice = selector(state)
     function listenerToAdd() {
       const nextSlice = selector(state)


Test patch:
diff --git a/tests/context.test.tsx b/tests/context.test.tsx
index d7873fa..11668dd 100644
--- a/tests/context.test.tsx
+++ b/tests/context.test.tsx
@@ -1,7 +1,8 @@
 import { Component as ClassComponent, useEffect, useState } from 'react'
 import { render } from '@testing-library/react'
-import create from 'zustand'
+import create, { StoreApi } from 'zustand'
 import createContext from 'zustand/context'
+import { subscribeWithSelector } from 'zustand/middleware'
 
 const consoleError = console.error
 afterEach(() => {
@@ -63,13 +64,23 @@ it('uses context store with selectors', async () => {
 })
 
 it('uses context store api', async () => {
-  const { Provider, useStoreApi } = createContext<CounterState>()
+  const { Provider, useStoreApi: useStoreApiOrig } =
+    createContext<CounterState>()
 
   const createStore = () =>
-    create<CounterState>((set) => ({
-      count: 0,
-      inc: () => set((state) => ({ count: state.count + 1 })),
-    }))
+    create(
+      subscribeWithSelector<{ count: number; inc: () => void }>((set) => ({
+        count: 0,
+        inc: () => set((state) => ({ count: state.count + 1 })),
+      }))
+    )
+
+  const useStoreApi = useStoreApiOrig as () => Omit<
+    StoreApi<CounterState>,
+    'subscribe'
+  > & {
+    subscribe: ReturnType<typeof createStore>['subscribe']
+  }
 
   function Counter() {
     const storeApi = useStoreApi()
@@ -77,8 +88,8 @@ it('uses context store api', async () => {
     useEffect(
       () =>
         storeApi.subscribe(
-          () => setCount(storeApi.getState().count),
-          (state) => state.count
+          (state) => state.count,
+          () => setCount(storeApi.getState().count)
         ),
       [storeApi]
     )
diff --git a/tests/middlewareTypes.test.tsx b/tests/middlewareTypes.test.tsx
index 2c9ced3..6cb1e00 100644
--- a/tests/middlewareTypes.test.tsx
+++ b/tests/middlewareTypes.test.tsx
@@ -1,7 +1,14 @@
 import { produce } from 'immer'
 import type { Draft } from 'immer'
 import create, { UseBoundStore } from 'zustand'
-import { NamedSet, combine, devtools, persist, redux } from 'zustand/middleware'
+import {
+  NamedSet,
+  combine,
+  devtools,
+  persist,
+  redux,
+  subscribeWithSelector,
+} from 'zustand/middleware'
 import { State, StateCreator } from 'zustand/vanilla'
 
 type TImmerConfigFn<T extends State> = (fn: (draft: Draft<T>) => void) => void
@@ -309,3 +316,87 @@ it('should combine devtools and combine', () => {
   }
   TestComponent
 })
+
+it('should combine subscribeWithSelector and combine', () => {
+  const useStore = create(
+    subscribeWithSelector(
+      combine({ count: 1 }, (set, get) => ({
+        inc: () => set({ count: get().count + 1 }, false, 'inc'),
+      }))
+    )
+  )
+
+  const TestComponent = (): JSX.Element => {
+    useStore().count
+    useStore().inc()
+    useStore.getState().count
+    useStore.getState().inc()
+    useStore.subscribe(
+      (state) => state.count,
+      (count) => console.log(count * 2)
+    )
+
+    return <></>
+  }
+  TestComponent
+})
+
+it('should combine devtools and subscribeWithSelector', () => {
+  const useStore = create(
+    devtools(
+      subscribeWithSelector<
+        {
+          count: number
+          inc: () => void
+        },
+        NamedSet<{
+          count: number
+          inc: () => void
+        }>
+      >((set, get) => ({
+        count: 1,
+        inc: () => set({ count: get().count + 1 }, false, 'inc'),
+      }))
+    )
+  )
+
+  const TestComponent = (): JSX.Element => {
+    useStore().count
+    useStore().inc()
+    useStore.getState().count
+    useStore.getState().inc()
+    useStore.subscribe(
+      (state) => state.count,
+      (count) => console.log(count * 2)
+    )
+
+    return <></>
+  }
+  TestComponent
+})
+
+it('should combine devtools, subscribeWithSelector and combine', () => {
+  const useStore = create(
+    devtools(
+      subscribeWithSelector(
+        combine({ count: 1 }, (set, get) => ({
+          inc: () => set({ count: get().count + 1 }, false, 'inc'),
+        }))
+      )
+    )
+  )
+
+  const TestComponent = (): JSX.Element => {
+    useStore().count
+    useStore().inc()
+    useStore.getState().count
+    useStore.getState().inc()
+    useStore.subscribe(
+      (state) => state.count,
+      (count) => console.log(count * 2)
+    )
+
+    return <></>
+  }
+  TestComponent
+})
diff --git a/tests/subscribe.test.tsx b/tests/subscribe.test.tsx
index 7fe45f3..154aa12 100644
--- a/tests/subscribe.test.tsx
+++ b/tests/subscribe.test.tsx
@@ -1,4 +1,5 @@
 import create from 'zustand'
+import { subscribeWithSelector } from 'zustand/middleware'
 
 describe('subscribe()', () => {
   it('should not be called if new state identity is the same', () => {
@@ -24,9 +25,11 @@ describe('subscribe()', () => {
   it('should not be called when state slice is the same', () => {
     const spy = jest.fn()
     const initialState = { value: 1, other: 'a' }
-    const { setState, subscribe } = create(() => initialState)
+    const { setState, subscribe } = create(
+      subscribeWithSelector(() => initialState)
+    )
 
-    subscribe(spy, (s) => s.value)
+    subscribe((s) => s.value, spy)
     setState({ other: 'b' })
     expect(spy).not.toHaveBeenCalled()
   })
@@ -34,9 +37,11 @@ describe('subscribe()', () => {
   it('should be called when state slice changes', () => {
     const spy = jest.fn()
     const initialState = { value: 1, other: 'a' }
-    const { setState, subscribe } = create(() => initialState)
+    const { setState, subscribe } = create(
+      subscribeWithSelector(() => initialState)
+    )
 
-    subscribe(spy, (s) => s.value)
+    subscribe((s) => s.value, spy)
     setState({ value: initialState.value + 1 })
     expect(spy).toHaveBeenCalledTimes(1)
     expect(spy).toHaveBeenCalledWith(initialState.value + 1, initialState.value)
@@ -45,9 +50,11 @@ describe('subscribe()', () => {
   it('should not be called when equality checker returns true', () => {
     const spy = jest.fn()
     const initialState = { value: 1, other: 'a' }
-    const { setState, subscribe } = create(() => initialState)
+    const { setState, subscribe } = create(
+      subscribeWithSelector(() => initialState)
+    )
 
-    subscribe(spy, undefined, () => true)
+    subscribe((s) => s, spy, { equalityFn: () => true })
     setState({ value: initialState.value + 2 })
     expect(spy).not.toHaveBeenCalled()
   })
@@ -55,13 +62,11 @@ describe('subscribe()', () => {
   it('should be called when equality checker returns false', () => {
     const spy = jest.fn()
     const initialState = { value: 1, other: 'a' }
-    const { setState, subscribe } = create(() => initialState)
-
-    subscribe(
-      spy,
-      (s) => s.value,
-      () => false
+    const { setState, subscribe } = create(
+      subscribeWithSelector(() => initialState)
     )
+
+    subscribe((s) => s.value, spy, { equalityFn: () => false })
     setState({ value: initialState.value + 2 })
     expect(spy).toHaveBeenCalledTimes(1)
     expect(spy).toHaveBeenCalledWith(initialState.value + 2, initialState.value)
@@ -70,9 +75,11 @@ describe('subscribe()', () => {
   it('should unsubscribe correctly', () => {
     const spy = jest.fn()
     const initialState = { value: 1, other: 'a' }
-    const { setState, subscribe } = create(() => initialState)
+    const { setState, subscribe } = create(
+      subscribeWithSelector(() => initialState)
+    )
 
-    const unsub = subscribe(spy, (s) => s.value)
+    const unsub = subscribe((s) => s.value, spy)
 
     setState({ value: initialState.value + 1 })
     unsub()
@@ -85,7 +92,9 @@ describe('subscribe()', () => {
   it('should keep consistent behavior with equality check', () => {
     const spy = jest.fn()
     const initialState = { value: 1, other: 'a' }
-    const { getState, setState, subscribe } = create(() => initialState)
+    const { getState, setState, subscribe } = create(
+      subscribeWithSelector(() => initialState)
+    )
 
     const isRoughEqual = (x: number, y: number) => Math.abs(x - y) < 1
     setState({ value: 0 })
@@ -100,7 +109,7 @@ describe('subscribe()', () => {
       spy(s.value, prevValue)
       prevValue = s.value
     })
-    const unsub2 = subscribe(spy2, (s) => s.value, isRoughEqual)
+    const unsub2 = subscribe((s) => s.value, spy2, { equalityFn: isRoughEqual })
     setState({ value: 0.5 })
     setState({ value: 1 })
     unsub()