[v4-alpha] use mutable source

.size-snapshot.json

@@ -1,20 +1,20 @@
 {
   "index.js": {
-    "bundled": 4684,
-    "minified": 1888,
-    "gzipped": 846
+    "bundled": 3513,
+    "minified": 1263,
+    "gzipped": 597
   },
   "index.mjs": {
-    "bundled": 3993,
-    "minified": 1614,
-    "gzipped": 786,
+    "bundled": 2840,
+    "minified": 1019,
+    "gzipped": 525,
     "treeshaked": {
       "rollup": {
-        "code": 124,
+        "code": 14,
         "import_statements": 14
       },
       "webpack": {
-        "code": 1144
+        "code": 998
       }
     }
   },

package.json

@@ -1,7 +1,10 @@
 {
   "name": "zustand",
   "private": true,
-  "version": "3.5.9",
+  "version": "4.0.0-alpha.1",
+  "publishConfig": {
+    "tag": "alpha"
+  },
   "description": "🐻 Bear necessities for state management in React",
   "main": "./index.js",
   "types": "./index.d.ts",
@@ -151,16 +154,16 @@
     "json": "^11.0.0",
     "lint-staged": "^11.1.2",
     "prettier": "^2.3.2",
-    "react": "^17.0.2",
-    "react-dom": "^17.0.2",
+    "react": "alpha",
+    "react-dom": "alpha",
     "rollup": "^2.56.2",
     "rollup-plugin-esbuild": "^4.5.0",
     "rollup-plugin-size-snapshot": "^0.12.0",
     "shx": "^0.3.3",
     "typescript": "^4.3.5"
   },
   "peerDependencies": {
-    "react": ">=16.8"
+    "react": "alpha"
   },
   "peerDependenciesMeta": {
     "react": {

src/index.ts

@@ -1,4 +1,10 @@
-import { useEffect, useLayoutEffect, useReducer, useRef } from 'react'
+/// <reference types="react/next" />
+
+import {
+  unstable_createMutableSource as createMutableSource,
+  useMemo,
+  unstable_useMutableSource as useMutableSource,
+} from 'react'
 import createImpl, {
   Destroy,
   EqualityChecker,
@@ -12,15 +18,6 @@ import createImpl, {
 } from './vanilla'
 export * from './vanilla'
 
-// For server-side rendering: https://github.com/pmndrs/zustand/pull/34
-// Deno support: https://github.com/pmndrs/zustand/issues/347
-const isSSR =
-  typeof window === 'undefined' ||
-  !window.navigator ||
-  /ServerSideRendering|^Deno\//.test(window.navigator.userAgent)
-
-const useIsomorphicLayoutEffect = isSSR ? useEffect : useLayoutEffect
-
 export interface UseStore<T extends State> {
   (): T
   <U>(selector: StateSelector<T, U>, equalityFn?: EqualityChecker<U>): U
@@ -36,102 +33,55 @@ export default function create<TState extends State>(
   const api: StoreApi<TState> =
     typeof createState === 'function' ? createImpl(createState) : createState
 
+  const source = createMutableSource(api, () => api.getState())
+  const subscribe = (api: StoreApi<TState>, callback: () => void) =>
+    api.subscribe(callback)
+
+  const FUNCTION_SYMBOL = Symbol()
+  const functionMap = new WeakMap<Function, { [FUNCTION_SYMBOL]: Function }>()
+
   const useStore: any = <StateSlice>(
     selector: StateSelector<TState, StateSlice> = api.getState as any,
     equalityFn: EqualityChecker<StateSlice> = Object.is
   ) => {
-    const [, forceUpdate] = useReducer((c) => c + 1, 0) as [never, () => void]
-
-    const state = api.getState()
-    const stateRef = useRef(state)
-    const selectorRef = useRef(selector)
-    const equalityFnRef = useRef(equalityFn)
-    const erroredRef = useRef(false)
-
-    const currentSliceRef = useRef<StateSlice>()
-    if (currentSliceRef.current === undefined) {
-      currentSliceRef.current = selector(state)
-    }
-
-    let newStateSlice: StateSlice | undefined
-    let hasNewStateSlice = false
-
-    // The selector or equalityFn need to be called during the render phase if
-    // they change. We also want legitimate errors to be visible so we re-run
-    // them if they errored in the subscriber.
-    if (
-      stateRef.current !== state ||
-      selectorRef.current !== selector ||
-      equalityFnRef.current !== equalityFn ||
-      erroredRef.current
-    ) {
-      // Using local variables to avoid mutations in the render phase.
-      newStateSlice = selector(state)
-      hasNewStateSlice = !equalityFn(
-        currentSliceRef.current as StateSlice,
-        newStateSlice
-      )
-    }
-
-    // Syncing changes in useEffect.
-    useIsomorphicLayoutEffect(() => {
-      if (hasNewStateSlice) {
-        currentSliceRef.current = newStateSlice as StateSlice
-      }
-      stateRef.current = state
-      selectorRef.current = selector
-      equalityFnRef.current = equalityFn
-      erroredRef.current = false
-    })
-
-    const stateBeforeSubscriptionRef = useRef(state)
-    useIsomorphicLayoutEffect(() => {
-      const listener = () => {
+    const getSnapshot = useMemo(() => {
+      let lastSlice: StateSlice | undefined
+      return (api: StoreApi<TState>) => {
+        let slice = lastSlice
         try {
-          const nextState = api.getState()
-          const nextStateSlice = selectorRef.current(nextState)
-          if (
-            !equalityFnRef.current(
-              currentSliceRef.current as StateSlice,
-              nextStateSlice
-            )
-          ) {
-            stateRef.current = nextState
-            currentSliceRef.current = nextStateSlice
-            forceUpdate()
+          slice = selector(api.getState())
+          if (lastSlice !== undefined && equalityFn(lastSlice, slice)) {
+            slice = lastSlice
+          } else {
+            lastSlice = slice
           }
         } catch (error) {
-          erroredRef.current = true
-          forceUpdate()
+          // ignore and let react reschedule update
         }
+        // Unfortunately, returning a function is not supported
+        // https://github.com/facebook/react/issues/18823
+        if (typeof slice === 'function') {
+          if (functionMap.has(slice)) {
+            return functionMap.get(slice)
+          }
+          const wrappedFunction = { [FUNCTION_SYMBOL]: slice }
+          functionMap.set(slice, wrappedFunction)
+          return wrappedFunction
+        }
+        return slice
       }
-      const unsubscribe = api.subscribe(listener)
-      if (api.getState() !== stateBeforeSubscriptionRef.current) {
-        listener() // state has changed before subscription
-      }
-      return unsubscribe
-    }, [])
-
-    return hasNewStateSlice
-      ? (newStateSlice as StateSlice)
-      : currentSliceRef.current
+    }, [selector, equalityFn])
+    const snapshot = useMutableSource(source, getSnapshot, subscribe)
+    if (
+      snapshot &&
+      (snapshot as { [FUNCTION_SYMBOL]: unknown })[FUNCTION_SYMBOL]
+    ) {
+      return (snapshot as { [FUNCTION_SYMBOL]: unknown })[FUNCTION_SYMBOL]
+    }
+    return snapshot
   }
 
   Object.assign(useStore, api)
 
-  // For backward compatibility (No TS types for this)
-  useStore[Symbol.iterator] = function () {
-    console.warn(
-      '[useStore, api] = create() is deprecated and will be removed in v4'
-    )
-    const items = [useStore, api]
-    return {
-      next() {
-        const done = items.length <= 0
-        return { value: items.shift(), done }
-      },
-    }
-  }
-
   return useStore
 }

tests/basic.test.tsx

@@ -1,5 +1,6 @@
 import {
   Component as ClassComponent,
+  useCallback,
   useEffect,
   useLayoutEffect,
   useState,
@@ -150,6 +151,7 @@ it('re-renders with useLayoutEffect', async () => {
 
   const container = document.createElement('div')
   ReactDOM.render(<Component />, container)
+  await new Promise((resolve) => setTimeout(resolve, 10))
   expect(container.innerHTML).toBe('true')
   ReactDOM.unmountComponentAtNode(container)
 })
@@ -219,14 +221,14 @@ it('can update the equality checker', async () => {
 
   // This will cause a re-render due to the equality checker.
   act(() => setState({ value: 0 }))
-  await findByText('renderCount: 2, value: 0')
+  await findByText('renderCount: 1, value: 0')
 
   // Set an equality checker that always returns true to never re-render.
   rerender(<Component equalityFn={() => true} />)
 
   // This will NOT cause a re-render due to the equality checker.
   act(() => setState({ value: 1 }))
-  await findByText('renderCount: 3, value: 0')
+  await findByText('renderCount: 2, value: 0')
 })
 
 it('can call useStore with progressively more arguments', async () => {
@@ -313,7 +315,8 @@ it('can throw an error in selector', async () => {
   act(() => {
     setState({})
   })
-  await findByText('errored')
+  // in v4 with uMS, we don't use equalityFn in render
+  // await findByText('errored')
 })
 
 it('can throw an error in equality checker', async () => {
@@ -357,7 +360,8 @@ it('can throw an error in equality checker', async () => {
   act(() => {
     setState({})
   })
-  await findByText('errored')
+  // in v4 with uMS, we don't use equalityFn in render
+  // await findByText('errored')
 })
 
 it('can get the store', () => {
@@ -428,6 +432,7 @@ it('only calls selectors when necessary', async () => {
   const useStore = create<State>(() => ({ a: 0, b: 0 }))
   const { setState } = useStore
   let inlineSelectorCallCount = 0
+  let callbackSelectorCallCount = 0
   let staticSelectorCallCount = 0
 
   function staticSelector(s: State) {
@@ -437,25 +442,30 @@ it('only calls selectors when necessary', async () => {
 
   function Component() {
     useStore((s) => (inlineSelectorCallCount++, s.b))
+    useStore(useCallback((s) => (callbackSelectorCallCount++, s.b), []))
     useStore(staticSelector)
     return (
       <>
         <div>inline: {inlineSelectorCallCount}</div>
+        <div>callback: {callbackSelectorCallCount}</div>
         <div>static: {staticSelectorCallCount}</div>
       </>
     )
   }
 
   const { rerender, findByText } = render(<Component />)
   await findByText('inline: 1')
+  await findByText('callback: 1')
   await findByText('static: 1')
 
   rerender(<Component />)
   await findByText('inline: 2')
+  await findByText('callback: 1')
   await findByText('static: 1')
 
   act(() => setState({ a: 1, b: 1 }))
   await findByText('inline: 4')
+  await findByText('callback: 2')
   await findByText('static: 2')
 })