refactor: useContextSelector of all the context usage (#862)

* test: test driven

* chore: all repalce with useContextSelector

* test

* chore: fix prettier remove import

.fatherrc.ts

@@ -4,6 +4,6 @@ export default {
   runtimeHelpers: true,
   preCommit: {
     eslint: true,
-    prettier: true,
+    prettier: false,
   },
 };

docs/examples/simple.tsx

@@ -1,6 +1,6 @@
 /* eslint-disable no-console,func-names,react/no-multi-comp */
-import React from 'react';
 import Table from 'rc-table';
+import React from 'react';
 import '../../assets/index.less';
 
 interface RecordType {
@@ -9,42 +9,50 @@ interface RecordType {
   c?: string;
 }
 
-const columns = [
-  { title: 'title1', dataIndex: 'a', key: 'a', width: 100 },
-  { id: '123', title: 'title2', dataIndex: 'b', key: 'b', width: 100, align: 'right' },
-  { title: 'title3', dataIndex: 'c', key: 'c', width: 200 },
-  {
-    title: 'Operations',
-    dataIndex: '',
-    key: 'd',
-    render(_: any, record: RecordType) {
-      return (
-        <a
-          onClick={e => {
-            e.preventDefault();
-            console.log('Operate on:', record);
-          }}
-          href="#"
-        >
-          Operations
-        </a>
-      );
-    },
-  },
-];
-
 const data = [
   { a: '123', key: '1' },
   { a: 'cdd', b: 'edd', key: '2' },
   { a: '1333', c: 'eee', d: 2, key: '3' },
 ];
 
-const Demo = () => (
-  <div>
-    <h2>simple table</h2>
-    <Table<RecordType> columns={columns} data={data} />
-  </div>
-);
+class Demo extends React.Component {
+  constructor(props) {
+    super(props);
+
+    this.state = {
+      count: 0,
+    };
+
+    this.columns = [
+      {
+        title: 'title1',
+        dataIndex: 'a',
+        render: this.renderColumn,
+      },
+    ];
+  }
+
+  renderColumn = () => {
+    return this.state.count;
+  };
+
+  render() {
+    return (
+      <>
+        <button
+          onClick={() => {
+            this.setState({
+              count: this.state.count + 1,
+            });
+          }}
+        >
+          Click {this.state.count} times
+        </button>
+        <Table<RecordType> columns={this.columns} data={data} />
+      </>
+    );
+  }
+}
 
 export default Demo;
 /* eslint-enable */

src/Body/BodyRow.tsx

@@ -1,16 +1,17 @@
-import * as React from 'react';
 import classNames from 'classnames';
+import * as React from 'react';
 import Cell from '../Cell';
-import TableContext from '../context/TableContext';
 import BodyContext from '../context/BodyContext';
-import { getColumnsKey } from '../utils/valueUtil';
+import TableContext from '../context/TableContext';
+import { useContextSelector } from '../ContextSelector';
 import type {
   ColumnType,
   CustomizeComponent,
   GetComponentProps,
-  Key,
   GetRowKey,
+  Key,
 } from '../interface';
+import { getColumnsKey } from '../utils/valueUtil';
 import ExpandedRow from './ExpandedRow';
 
 export interface BodyRowProps<RecordType> {
@@ -49,7 +50,10 @@ function BodyRow<RecordType extends { children?: readonly RecordType[] }>(
     cellComponent,
     childrenColumnName,
   } = props;
-  const { prefixCls, fixedInfoList } = React.useContext(TableContext);
+  const { prefixCls, fixedInfoList } = useContextSelector(TableContext, [
+    'prefixCls',
+    'fixedInfoList',
+  ]);
   const {
     flattenColumns,
     expandableType,
@@ -61,7 +65,18 @@ function BodyRow<RecordType extends { children?: readonly RecordType[] }>(
     expandIcon,
     expandedRowRender,
     expandIconColumnIndex,
-  } = React.useContext(BodyContext);
+  } = useContextSelector(BodyContext, [
+    'flattenColumns',
+    'expandableType',
+    'expandRowByClick',
+    'onTriggerExpand',
+    'rowClassName',
+    'expandedRowClassName',
+    'indentSize',
+    'expandIcon',
+    'expandedRowRender',
+    'expandIconColumnIndex',
+  ]);
   const [expandRended, setExpandRended] = React.useState(false);
 
   const expanded = expandedKeys && expandedKeys.has(props.recordKey);

src/Body/ExpandedRow.tsx

@@ -1,8 +1,9 @@
 import * as React from 'react';
-import type { CustomizeComponent } from '../interface';
 import Cell from '../Cell';
-import TableContext from '../context/TableContext';
 import ExpandedRowContext from '../context/ExpandedRowContext';
+import TableContext from '../context/TableContext';
+import { useContextSelector } from '../ContextSelector';
+import type { CustomizeComponent } from '../interface';
 
 export interface ExpandedRowProps {
   prefixCls: string;
@@ -25,9 +26,11 @@ function ExpandedRow({
   colSpan,
   isEmpty,
 }: ExpandedRowProps) {
-  const { scrollbarSize } = React.useContext(TableContext);
-  const { fixHeader, fixColumn, componentWidth, horizonScroll } =
-    React.useContext(ExpandedRowContext);
+  const scrollbarSize = useContextSelector(TableContext, 'scrollbarSize');
+  const { fixHeader, fixColumn, componentWidth, horizonScroll } = useContextSelector(
+    ExpandedRowContext,
+    ['fixHeader', 'fixColumn', 'componentWidth', 'horizonScroll'],
+  );
 
   // Cache render node
   return React.useMemo(() => {

src/Body/index.tsx

@@ -1,15 +1,16 @@
 import * as React from 'react';
-import TableContext from '../context/TableContext';
-import type { GetRowKey, Key, GetComponentProps } from '../interface';
-import ExpandedRow from './ExpandedRow';
 import BodyContext from '../context/BodyContext';
-import { getColumnsKey } from '../utils/valueUtil';
-import ResizeContext from '../context/ResizeContext';
-import BodyRow from './BodyRow';
-import useFlattenRecords from '../hooks/useFlattenRecords';
 import HoverContext from '../context/HoverContext';
 import type { PerfRecord } from '../context/PerfContext';
 import PerfContext from '../context/PerfContext';
+import ResizeContext from '../context/ResizeContext';
+import TableContext from '../context/TableContext';
+import { useContextSelector } from '../ContextSelector';
+import useFlattenRecords from '../hooks/useFlattenRecords';
+import type { GetComponentProps, GetRowKey, Key } from '../interface';
+import { getColumnsKey } from '../utils/valueUtil';
+import BodyRow from './BodyRow';
+import ExpandedRow from './ExpandedRow';
 import MeasureRow from './MeasureRow';
 
 export interface BodyProps<RecordType> {
@@ -33,9 +34,12 @@ function Body<RecordType>({
   emptyNode,
   childrenColumnName,
 }: BodyProps<RecordType>) {
-  const { onColumnResize } = React.useContext(ResizeContext);
-  const { prefixCls, getComponent } = React.useContext(TableContext);
-  const { flattenColumns } = React.useContext(BodyContext);
+  const onColumnResize = useContextSelector(ResizeContext, 'onColumnResize');
+  const { prefixCls, getComponent } = useContextSelector(TableContext, [
+    'prefixCls',
+    'getComponent',
+  ]);
+  const flattenColumns = useContextSelector(BodyContext, 'flattenColumns');
 
   const flattenData: { record: RecordType; indent: number; index: number }[] =
     useFlattenRecords<RecordType>(data, childrenColumnName, expandedKeys, getRowKey);

src/Cell/index.tsx

@@ -1,25 +1,25 @@
-import * as React from 'react';
 import classNames from 'classnames';
-import shallowEqual from 'shallowequal';
 import { supportRef } from 'rc-util/lib/ref';
+import warning from 'rc-util/lib/warning';
+import * as React from 'react';
+import shallowEqual from 'shallowequal';
+import BodyContext from '../context/BodyContext';
+import type { HoverContextProps } from '../context/HoverContext';
+import HoverContext from '../context/HoverContext';
+import PerfContext from '../context/PerfContext';
+import StickyContext from '../context/StickyContext';
+import { useContextSelector } from '../ContextSelector';
 import type {
-  DataIndex,
+  AlignType,
+  CellEllipsisType,
+  CellType,
   ColumnType,
-  RenderedCell,
   CustomizeComponent,
-  CellType,
+  DataIndex,
   DefaultRecordType,
-  AlignType,
-  CellEllipsisType,
+  RenderedCell,
 } from '../interface';
 import { getPathValue, validateValue } from '../utils/valueUtil';
-import StickyContext from '../context/StickyContext';
-import HoverContext from '../context/HoverContext';
-import BodyContext from '../context/BodyContext';
-import type { HoverContextProps } from '../context/HoverContext';
-import warning from 'rc-util/lib/warning';
-import PerfContext from '../context/PerfContext';
-import { useContextSelector } from '../ContextSelector';
 
 /** Check if cell is in hover range */
 function inHoverRange(cellStartRow: number, cellRowSpan: number, startRow: number, endRow: number) {
@@ -143,7 +143,7 @@ function Cell<RecordType extends DefaultRecordType>(
 
   const perfRecord = React.useContext(PerfContext);
   const supportSticky = React.useContext(StickyContext);
-  const { allColumnsFixedLeft } = React.useContext(BodyContext);
+  const allColumnsFixedLeft = useContextSelector(BodyContext, 'allColumnsFixedLeft');
 
   // ==================== Child Node ====================
   const [childNode, legacyCellProps] = React.useMemo<

src/ContextSelector/index.tsx

@@ -1,6 +1,7 @@
-import * as React from 'react';
-import useLayoutEffect from 'rc-util/lib/hooks/useLayoutEffect';
 import useEvent from 'rc-util/lib/hooks/useEvent';
+import useLayoutEffect from 'rc-util/lib/hooks/useLayoutEffect';
+import * as React from 'react';
+import { unstable_batchedUpdates } from 'react-dom';
 import shallowEqual from 'shallowequal';
 
 export type Selector<T, O = T> = (value: T) => O;
@@ -24,8 +25,8 @@ export interface ReturnCreateContext<T> {
   Provider: React.ComponentType<ContextSelectorProviderProps<T>>;
 }
 
-export function createContext<T>(): ReturnCreateContext<T> {
-  const Context = React.createContext<Context<T>>(null as any);
+export function createContext<T>(defaultContext?: T): ReturnCreateContext<T> {
+  const Context = React.createContext<Context<T>>(defaultContext as any);
 
   const Provider = ({ value, children }: ContextSelectorProviderProps<T>) => {
     const valueRef = React.useRef(value);
@@ -37,8 +38,10 @@ export function createContext<T>(): ReturnCreateContext<T> {
     }));
 
     useLayoutEffect(() => {
-      context.listeners.forEach(listener => {
-        listener(value);
+      unstable_batchedUpdates(() => {
+        context.listeners.forEach(listener => {
+          listener(value);
+        });
       });
     }, [value]);
 
@@ -48,23 +51,55 @@ export function createContext<T>(): ReturnCreateContext<T> {
   return { Context, Provider };
 }
 
-export function useContextSelector<T, O>(holder: ReturnCreateContext<T>, selector: Selector<T, O>) {
-  const eventSelector = useEvent(selector);
+export function useContextSelector<T, O>(
+  holder: ReturnCreateContext<T>,
+  selector: Selector<T, O>,
+): O;
+export function useContextSelector<T, O extends Partial<T>>(
+  holder: ReturnCreateContext<T>,
+  selector: (keyof T)[],
+): O;
+export function useContextSelector<T, S extends keyof T>(
+  holder: ReturnCreateContext<T>,
+  selector: S,
+): T[S];
+
+export function useContextSelector<T, O>(
+  holder: ReturnCreateContext<T>,
+  selector: Selector<T, any> | (keyof T)[] | keyof T,
+) {
+  const eventSelector = useEvent<Selector<T, O>>(
+    typeof selector === 'function'
+      ? selector
+      : ctx => {
+          if (!Array.isArray(selector)) {
+            return ctx[selector];
+          }
+
+          const obj = {} as O;
+          selector.forEach(key => {
+            (obj as any)[key] = ctx[key];
+          });
+          return obj;
+        },
+  );
   const context = React.useContext(holder?.Context);
   const { listeners, getValue } = context || {};
 
-  const [value, setValue] = React.useState(() => eventSelector(context ? getValue() : null));
+  const valueRef = React.useRef<O>();
+  valueRef.current = eventSelector(context ? getValue() : null);
+  const [, forceUpdate] = React.useState({});
 
   useLayoutEffect(() => {
     if (!context) {
       return;
     }
 
     function trigger(nextValue: T) {
-      setValue(prev => {
-        const selectedValue = eventSelector(nextValue);
-        return shallowEqual(prev, selectedValue) ? prev : selectedValue;
-      });
+      const nextSelectorValue = eventSelector(nextValue);
+      if (!shallowEqual(valueRef.current, nextSelectorValue)) {
+        forceUpdate({});
+      }
     }
 
     listeners.add(trigger);
@@ -74,5 +109,5 @@ export function useContextSelector<T, O>(holder: ReturnCreateContext<T>, selecto
     };
   }, [context]);
 
-  return value;
+  return valueRef.current;
 }