From 5ee5de39f4d0c6b159878174c4db2d99aadf4e15 Mon Sep 17 00:00:00 2001
From: Josep M Sobrepere <josepot@gmail.com>
Date: Thu, 2 May 2019 17:40:07 +0200
Subject: [PATCH] Avoid unnecessary selector evaluations

---
 src/hooks/useSelector.js       | 44 +++++++++++++---------------
 test/hooks/useSelector.spec.js | 53 ++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 24 deletions(-)

diff --git a/src/hooks/useSelector.js b/src/hooks/useSelector.js
index 8e40cfa10..e3277ca20 100644
--- a/src/hooks/useSelector.js
+++ b/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,41 +42,30 @@ 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
-
+  let state, setState
   try {
-    selectedState = selector(store.getState())
+    ;[state, setState] = useState(selector(store.getState()))
   } catch (err) {
-    let errorMessage = `An error occured while selecting the store state: ${
+    const 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)
   }
+  const [error, setError] = useState(null)
 
-  const latestSelectedState = useRef(selectedState)
+  const latestSelector = useRef(selector)
+  const latestSelectedState = useRef(state)
 
   useIsomorphicLayoutEffect(() => {
     latestSelector.current = selector
-    latestSelectedState.current = selectedState
-    latestSubscriptionCallbackError.current = undefined
-  })
+    setState(selector(store.getState()))
+  }, [selector])
 
   useIsomorphicLayoutEffect(() => {
     function checkForUpdates() {
@@ -88,15 +77,14 @@ export function useSelector(selector) {
         }
 
         latestSelectedState.current = newSelectedState
+        setState(latestSelectedState.current)
       } catch (err) {
         // we ignore all errors here, since when the component
         // is re-rendered, the selectors are called again, and
         // will throw again, if neither props nor store state
         // changed
-        latestSubscriptionCallbackError.current = err
+        setError(err)
       }
-
-      forceRender({})
     }
 
     subscription.onStateChange = checkForUpdates
@@ -107,5 +95,13 @@ export function useSelector(selector) {
     return () => subscription.tryUnsubscribe()
   }, [store, subscription])
 
-  return selectedState
+  if (error) {
+    const errorMessage = `An error occured while selecting the store state: ${
+      error.message
+    }.\n\nOriginal stack trace:${error.stack}`
+
+    throw new Error(errorMessage)
+  }
+
+  return state
 }
diff --git a/test/hooks/useSelector.spec.js b/test/hooks/useSelector.spec.js
index 74043f1c1..6897ec308 100644
--- a/test/hooks/useSelector.spec.js
+++ b/test/hooks/useSelector.spec.js
@@ -47,6 +47,31 @@ describe('React', () => {
       })
 
       describe('lifeycle interactions', () => {
+        it.only('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('subscribes to the store synchronously', () => {
           let rootSubscription
 
@@ -183,6 +208,34 @@ describe('React', () => {
 
           expect(renderedItems).toEqual([0, 0])
         })
+
+        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>
+          )
+
+          rtl.act(forceRender)
+
+          // this line verifies the susbcription callback uses the same memoized selector and therefore
+          // does not cause a re-render
+          store.dispatch({ type: '' })
+
+          expect(renderedItems).toEqual([0, 1])
+        })
       })
 
       describe('edge cases', () => {