useSelector是如何触发更新的以及手写一个简单的useSelector

概述

useSelector是react-redux@7中加入的hook，可以在不使用connect()的情况下将函数组件连接到redux，这样代码写起来会更加清晰，更加方便。

使用起来也很简单，我们写一个简单的加减数组件来看一下

// index.tsx
import React from 'react';
import ReactDOM from 'react-dom';
import { createStore, Reducer } from 'redux';
import { Provider } from 'react-redux';
import Sub from './sub';

export interface StoreState {
  count: number
}
export interface StoreAction {
  type: 'change'
  payload: StoreState
}

const reducer: Reducer<StoreState, StoreAction> = (state, action) => ({ ...state, ...action.payload });
const store = createStore(reducer, { count: 0 });

function App() {
  return (
    <Provider store={store}>
      <Sub />
    </Provider>
  );
}

ReactDOM.render(<App />, document.getElementById('root'));

// Sub.tsx
import React from 'react';
import { StoreState, StoreAction } from './index';
import { useDispatch, useSelector } from 'react-redux';

export default function Sub() {
  const count = useSelector<StoreState, number>((state) => state.count);
  const dispatch = useDispatch<StoreAction>();
  const customEqalityCount = useSelector<StoreState, number>((state) => state.count, (a, b) => a > b);
  
  return (
    <div>
      <div>{count}</div>
      <div onClick={() => dispatch(
        {
          type: 'change',
          payload: { count: count + 1 },
        },
      )}>
        点击增加
      </div>
      <div onClick={() => dispatch(
        {
          type: 'change',
          payload: { count: count - 1 },
        },
      )}>
        点击减少
      </div>
      <div>{customEqalityCount}</div>
    </div>
  );
}

在index.tsx中创建了一个store，丢到Provider中，在子组件中使用useSelector获取store中最新的state, useDispatch更新store中的state，这样一个简单的加减数功能就完成了，注意在Sub中第二个useSelector使用了两个参数，向它传递了一个新旧count比较函数，只有该函数返回false的时候才会触发更新。在这里我传了一个(a, b) => a > b，意味着该值只能减少不能增加。

可以在sandbox里玩下试试

原理浅析

其实原理也不复杂，使用了useContext的特性，但看过源码后，发现直接想的一些细节很妙。
我们可以先想一下，实现一个useSelector有哪些问题需要解决：

useSelector如何获取store
如何知道store中的state已经变了
如何触发组件re-render
如何记录变化前的state
如何返回用户希望拿到的state

第一个问题最简单，直接使用useContext就可以拿到。怎么知道state已经变了呢，这里我一开始有个误区，以为直接把store或者把store.getState()获取的state放到useEffect的依赖里就可以知道了。可问题是store会变吗，答案是不会，store是一个对象，只要store通过createStore()创建，这个对象的引用就不会变。state确实会变，但这个变化react可以知道吗，state只是一个值，是一个闭包，而不是react通过useState创建的，react是不知道他是否变化的，换句话说state改变时不会通知react。

那么如何解决呢，答案就在谜面上，在store.subscribe()里订阅就可以了，我们可以在回调函数中比较变化前后的state，去触发更新。

OK，如何触发组件re-render呢，这个也比较简单，用 useState 或 useReducer 记录一个无意义的状态，在需要重新渲染的时候，改变它就可以了。

如何记录变化前后的state呢，可以用useState useReducer吗，当然不可以，我们记录之前的state是为了与现在的state进行比较，从而决定是否触发组件更新，使用这两个api可能引起额外的非必要更新，那能记录状态且不会触发re-render的api只有useRef了。

如何返回用户想要的state呢，哈哈哈，自定义hook是个函数呀，直接返回就完事了。

手写一个简单的useSelector

原理差不多搞清楚了之后，我们就可以来试着模拟一个useSelector，实践一下。注：以下实现简化了源码中的很多细节

首先，我们需要有一个context

// context.ts
import { createContext } from "react";
import { AnyAction, Store } from "redux";

const StoreContext = createContext<Store<any, AnyAction>>(null as any);
export default StoreContext;

有了context就可以写provider跟useDispatch了

// Provider.tsx
import React from 'react';
import { AnyAction, Store, Action } from 'redux';
import StoreContext from './context';

interface ProviderParams<T extends Action = AnyAction, S = any> {
  store: Store<S, T>,
  children: JSX.Element
}

export default function Provider
<T extends Action = AnyAction>({ store, children }: ProviderParams<T>) {
  // @ts-ignore
  return <StoreContext.Provider value={store}>{children}</StoreContext.Provider>;
}

// useDispatch.ts
import { useContext } from 'react';
import { Action, Dispatch } from 'redux';
import StoreContext from './context';

export default function useDispatch<T extends Action>(): Dispatch<T> {
  const store = useContext(StoreContext);
  return store.dispatch;
}

然后就可以将最开始index.tsx的代码改一下，引用我们自己文件。

// index.tsx
import React from 'react';
import ReactDOM from 'react-dom';
import { createStore, Reducer } from 'redux';
import Provider from './Provider';
import Sub from './sub';
// ...

我们在这里创建store，通过context传下去。
接下来就可写useSelector了

import {
  useContext, useEffect, useReducer, useRef,
} from 'react';
import StoreContext from './context';

type EqualityFn<T> = (a: T, b: T) => boolean;

export default function useSelector<T, Selected extends unknown>(
  selector: (state: T) => Selected,
  equalityFn?: EqualityFn<Selected>,
): Selected {
  const store = useContext(StoreContext);
  const [, forceRender] = useReducer((s) => s + 1, 0);

  const latestStoreState = useRef<T>(store.getState());
  const latestSelectedState = useRef<Selected>(selector(latestStoreState.current));

  useEffect(() => {
    function checkUpdate() {
      const newState = store.getState();
      if (newState === latestStoreState) return;

      const newSelectedState = selector(newState);
      if (!equalityFn) equalityFn = (a, b) => a === b;

      if (!equalityFn(newSelectedState, latestSelectedState.current)) {
        latestSelectedState.current = newSelectedState;
        latestStoreState.current = newState;
        forceRender();
      }
    }
    const unsubscribe = store.subscribe(checkUpdate);
    return () => unsubscribe();
  }, [store]);

  return latestSelectedState.current;
}

最后改下sub.tsx中的代码，引用我们自己的文件

import React from 'react';
import { StoreState, StoreAction } from './index';
import useDispatch from './useDispatch';
import useSelector from './useSelector';
//...