Avoid unnecessary selector evaluations

Author: josepot

src/hooks/useSelector.js

@@ -1,4 +1,4 @@
-import { useReducer, useRef, useEffect, useMemo, useLayoutEffect } from 'react'
+import { useState, useRef, useEffect, useMemo, useLayoutEffect } from 'react'
 import invariant from 'invariant'
 import { useReduxContext } from './useReduxContext'
 import shallowEqual from '../utils/shallowEqual'
@@ -42,52 +42,27 @@ export function useSelector(selector) {
   invariant(selector, `You must pass a selector to useSelectors`)
 
   const { store, subscription: contextSub } = useReduxContext()
-  const [, forceRender] = useReducer(s => s + 1, 0)
-
   const subscription = useMemo(() => new Subscription(store, contextSub), [
     store,
     contextSub
   ])
 
-  const latestSubscriptionCallbackError = useRef()
-  const latestSelector = useRef(selector)
-
-  let selectedState = undefined
-
-  try {
-    selectedState = selector(store.getState())
-  } catch (err) {
-    let errorMessage = `An error occured while selecting the store state: ${
-      err.message
-    }.`
+  const [subscriptionResult, setSubscriptionResult] = useState()
 
-    if (latestSubscriptionCallbackError.current) {
-      errorMessage += `\nThe error may be correlated with this previous error:\n${
-        latestSubscriptionCallbackError.current.stack
-      }\n\nOriginal stack trace:`
-    }
-
-    throw new Error(errorMessage)
-  }
-
-  const latestSelectedState = useRef(selectedState)
-
-  useIsomorphicLayoutEffect(() => {
-    latestSelector.current = selector
-    latestSelectedState.current = selectedState
-    latestSubscriptionCallbackError.current = undefined
-  })
+  const latestSelector = useRef()
+  const latestResult = useRef()
+  const latestSubscriptionCallbackError = useRef()
+  const latestSubscriptionResult = useRef()
 
   useIsomorphicLayoutEffect(() => {
     function checkForUpdates() {
+      let newSelectedState
       try {
-        const newSelectedState = latestSelector.current(store.getState())
+        newSelectedState = latestSelector.current(store.getState())
 
-        if (shallowEqual(newSelectedState, latestSelectedState.current)) {
+        if (shallowEqual(newSelectedState, latestResult.current)) {
           return
         }
-
-        latestSelectedState.current = newSelectedState
       } catch (err) {
         // we ignore all errors here, since when the component
         // is re-rendered, the selectors are called again, and
@@ -96,16 +71,52 @@ export function useSelector(selector) {
         latestSubscriptionCallbackError.current = err
       }
 
-      forceRender({})
+      const newSubscriptionResult = new Array(1)
+      newSubscriptionResult[0] = newSelectedState
+      setSubscriptionResult(newSubscriptionResult)
     }
 
     subscription.onStateChange = checkForUpdates
     subscription.trySubscribe()
 
-    checkForUpdates()
-
     return () => subscription.tryUnsubscribe()
   }, [store, subscription])
 
-  return selectedState
+  let result = latestResult.current
+
+  useIsomorphicLayoutEffect(() => {
+    latestResult.current = result
+    latestSelector.current = selector
+    latestSubscriptionCallbackError.current = undefined
+    latestSubscriptionResult.current = subscriptionResult
+  })
+
+  try {
+    return (result =
+      // If the selector has changed or if there has been an error while
+      // it was being evaluated inside "checkForUpdates", then the selector
+      // has to be re-evaluated
+      selector !== latestSelector.current ||
+      latestSubscriptionCallbackError.current
+        ? selector(store.getState())
+        : // If this update was triggered by "checkForUpdates", then we
+        // should return the result that was already calculated in there.
+        // Otherwise, this is an update triggered by an ancestor and
+        // there is no need to recalculate the result
+        subscriptionResult !== latestSubscriptionResult.current
+        ? subscriptionResult[0]
+        : result)
+  } catch (err) {
+    let errorMessage = `An error occured while selecting the store state: ${
+      err.message
+    }.`
+
+    if (latestSubscriptionCallbackError.current) {
+      errorMessage += `\nThe error may be correlated with this previous error:\n${
+        latestSubscriptionCallbackError.current.stack
+      }\n\nOriginal stack trace:`
+    }
+
+    throw new Error(errorMessage)
+  }
 }

test/hooks/useSelector.spec.js

@@ -1,6 +1,6 @@
 /*eslint-disable react/prop-types*/
 
-import React from 'react'
+import React, { useCallback, useReducer } from 'react'
 import { createStore } from 'redux'
 import { renderHook, act } from 'react-hooks-testing-library'
 import * as rtl from 'react-testing-library'
@@ -31,6 +31,31 @@ describe('React', () => {
           expect(result.current).toEqual(0)
         })
 
+        it('always uses the latest state', () => {
+          store = createStore(c => c + 1, -1)
+
+          const Comp = () => {
+            const selector = useCallback(c => c + 1, [])
+            const value = useSelector(selector)
+            renderedItems.push(value)
+            return <div />
+          }
+
+          rtl.render(
+            <ProviderMock store={store}>
+              <Comp />
+            </ProviderMock>
+          )
+
+          expect(renderedItems).toEqual([1])
+
+          act(() => {
+            store.dispatch({ type: '' })
+          })
+
+          expect(renderedItems).toEqual([1, 2])
+        })
+
         it('selects the state and renders the component when the store updates', () => {
           const { result } = renderHook(() => useSelector(s => s.count), {
             wrapper: props => <ProviderMock {...props} store={store} />
@@ -156,6 +181,39 @@ describe('React', () => {
         })
       })
 
+      it('uses the latest selector', () => {
+        let selectorId = 0
+        let forceRender
+
+        const Comp = () => {
+          const [, f] = useReducer(c => c + 1, 0)
+          forceRender = f
+          const renderedSelectorId = selectorId++
+          const value = useSelector(() => renderedSelectorId)
+          renderedItems.push(value)
+          return <div />
+        }
+
+        rtl.render(
+          <ProviderMock store={store}>
+            <Comp />
+          </ProviderMock>
+        )
+
+        expect(renderedItems).toEqual([0])
+
+        rtl.act(forceRender)
+        expect(renderedItems).toEqual([0, 1])
+
+        rtl.act(() => {
+          store.dispatch({ type: '' })
+        })
+        expect(renderedItems).toEqual([0, 1])
+
+        rtl.act(forceRender)
+        expect(renderedItems).toEqual([0, 1, 2])
+      })
+
       describe('edge cases', () => {
         it('ignores transient errors in selector (e.g. due to stale props)', () => {
           const spy = jest.spyOn(console, 'error').mockImplementation(() => {})