useSelector refactoring

useSelector.js

function useSelectorWithStoreAndSubscription(
  selector,
  equalityFn,
  store,
  contextSub
) {
  const subscription = useMemo(() => new Subscription(store, contextSub), [
    store,
    contextSub
  ])

  const [lastError, setLastError] = useState()
  let selectedState, setSelectedState
  try {
    const initialSelectedState = useMemo(() => selector(store.getState()), [selector, store, lastError])
    // use { value } as wrapper to support equalityFn that are laxer than Object.is
    void ([{ value: selectedState }, setSelectedState] = useState({ value: initialSelectedState }))

    if (!Object.is(initialSelectedState, selectedState)) {
      // can happen from useMemo result changes
      // ignore equalityFn for this to avoid infinite rerenders on e.g. () => false
      setSelectedState({ value: initialSelectedState })
      selectedState = initialSelectedState
    }
  } catch (err) {
    let errorMessage = `An error occurred while selecting the store state: ${err.message}.`

    if (lastError) {
      errorMessage += `\nThe error may be correlated with this previous error:\n${lastError.stack}\n\nOriginal stack trace:`
    }

    throw new Error(errorMessage)
  }
  
  const checkForUpdates = useRef()

  useIsomorphicLayoutEffect(() => {
    // should this be in the LayoutEffect (i.e. remember only the `selectedState` from the previous commit),
    // or should this be in the body of the component (i.e. always up to date with the most recent render pass)
    checkForUpdates.current = () => {
      try {
        const newSelectedState = selector(store.getState())
        
        // attempting to do this instead causes missed updates. Why? React bug?
        // setSelectedState(selectedState => equalityFn(newSelectedState, selectedState) ? selectedState : { value: newSelectedState })
  
        if (equalityFn(newSelectedState, selectedState)) {
          return
        }
        
        setSelectedState({ value: newSelectedState })
      } 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
        setLastError(err)
      }
    }
    if (lastError) {
      setLastError(undefined)
    }
  })

  useIsomorphicLayoutEffect(() => {
    subscription.onStateChange = () => checkForUpdates.current()
    subscription.trySubscribe()

    checkForUpdates.current()

    return () => subscription.tryUnsubscribe()
  }, [store, subscription])

  return selectedState
}