填坑 | React Context原理

很久没有更新我的 从0带你实现 React 18 专栏了，还剩下两个部分，一个是Context，一个是Suspense&unwind。这次来说说Context原理

Context的使用大家肯定不陌生了，被Context包裹的子节点可以通过 Consumer 或者 useContext获取Context的共享值，那么这是如何实现的呢？

我们简单的画图说明

一段简单的DOM结构: 

<Context_A.Provider value={A1}>
    <FunctionComponet1> // A1
        <Context_A.Provider value={A2}>
            <FunctionComponet2> // A2
            <FunctionComponet3> // A2
        </Context_A.Provider>
    </FunctionComponet1>
    <Context_B.Provider value={B1}> 
      <FunctionComponet4> // B1
    </Context_B.Provider>
</Context_A.Provider>

其对应的Fiber结构如下所示: 

可以看到，最外层有Context_A包裹，内部包含函数组件1 和 Context_B, 而函数组件1再次包含Context_A, Context_A 包含函数组件2 3 。 Context_B 包含函数组件4

在这四个函数组件内，其中 组件1可以拿到 A1的值，组件2 3 可以拿到 A2的值，也就是内层Context_A的值，组件4可以拿到B1， Context是可以层层嵌套的。

你也许会发现，嵌套Context的Value类似于栈，每经过一层嵌套就把最新的值Push到栈中，而回溯的时候每回溯到一层Context就把顶部的值Pop，React中就是这么实现的

React对于不同的Context，比如Context A , B 共用了一个Stack，因为Fiber节点的遍历是深度优先的，递归的顺序是稳定的，也就是说，当我向下递的顺序为 Context A -> Context B -> Context C的时候，向上回溯（归）的顺序一定是 Context C -> Context B -> Context A 天然的保证了某个节点在 递，归的两个阶段，压入栈和弹出栈的值是同一个。

当向下递到第一层 Context_A 时，会将当前Context_A的value A1 压入栈

每个Context内部都保存一个最新的值 _currentValue 当便利到一个Context时，都会将value属性传入的最新值爆存在 _currentValue中，并且推入栈

此时，Context_A下面的子节点，都能通过Context_A._currentValue拿到A1值，在函数组件1内，输出Context Value为栈顶的值 A1

继续向下递，遍历到第二层 Context_A的时候，将新的值A2 PUSH STACK 如图

下面的函数组件 2 3 读到的currentValue值就是 A2了，旧的的A1值就无法被 2 3 读取到了，相当于覆盖。

递到叶子结点，开始回溯，经过内层 Context_A 弹出A2，此时把弹出的A1值，赋给 Context_A的_currentValue 

继续向上回归，遍历到组件1的兄弟节点 Context_B 此时把B1 推入栈 并且 B1 赋给 _currentValue

此时 函数组件4 即可拿到 B的值 B1，如果在其中拿Context_A的值取到的也是 Context_A._currentValue =  A1，你可以理解为, _currentValue的值是随着遍历变化的，而组件读取Context都是读取其._currentValue值。

最后，回溯，弹出B1 设置 Context_B.current_value = null 弹出A1 设置Context_A._currentValue = null 如下

如果你创建Context的时候，通过参数传入初始值:

                        const Context = createContext(initialValue) 

那么._currentValue的默认值就是这个初始化值，而不是null

下面我们简单看一下实现，具体实现你可以看我的git仓库 https://github.com/Gravity2333/My-React/tree/learn

createContext & Context对象

我们使用 createContext创建一个Context对象，其结构如下所示

Context对象包含一个 _currentValue属性 记录最新的Context值

一个Provider属性指向Provider对象

一个Consumer属性指向Consumer对象

Provider和Consumer对象中包含内部属性 _context 指回 Context对象

实现代码如下:

// lib/react/context
import { REACT_CONSUMER_TYPE, REACT_CONTEXT_TYPE, REACT_PROVIDER_TYPE } from "../share/ReactSymbols";

export type Context<T> = {
  $$typeof: symbol | number;
  Provider: ContextProviderType<T>;
  Consumer: ContextConsumer<T>;
  _currentValue: T;
};

/** 创建一个Context对象 */
export default function createContext<T>(defaultValue: T): Context<T> {
  const _context: Context<T> = {
    $$typeof: REACT_CONTEXT_TYPE,
    Provider: null,
    Consumer: null,
    _currentValue: defaultValue,
  };

  // add provider
  _context.Provider = {
    $$typeof: REACT_PROVIDER_TYPE,
    _context,
  };

  // add consumer
  _context.Consumer = {
    $$typeof: REACT_CONSUMER_TYPE,
    _context,
  };

  return _context;
}

updateContextProvider

这个阶段在Beginwork阶段完成，也就是递的阶段。

这里需要注意，Context对象是不直接生成Fiber对象的，我们平时使用Context的时候也是使用 

Context.Provider / Context.Consumer 

// react-reconsiler/beginwork.ts
/** 更新ContextProvider */
function updateContextProvider(wip: FiberNode, renderLane: Lane) {
  const context = wip.type._context;
  const pendingProps = wip.pendingProps;
  // 获取最新值
  const newValue = pendingProps?.value;

  // 最新值推入 推入Context Stack
  pushContext(context, newValue);

  // reconcile child
  reconcileChildren(wip, pendingProps.children);
  return wip.child;
}

逻辑很简单，拿到最新的值，推入上下文栈，并且协调子节点，和其他类型的节点没什么区别

需要注意，整个Context对象是从 wip.type._context中取得，因为我们在使用 Context.Provider时，通常

<Context.Provider value=""> ... </Context.Provider value="">
等价于
createElement(Context.Provider, {value} ,...)

可以看到，整个Provider对象是作为type保存到Element元素中，生成Fiber的时候也会将其爆存在Fiber.type中。我们通过Provider对象._context属性即可获得Context对象！

 PushContext & PopContext 

简单说一下上下文栈，实现也很简单，如下:

// fiberContext.ts

/**
 * 用来记录上一个context的值
 */
let prevContextValue: any = null;

/**
 * 用来记录prevContextValue的栈
 * 当Context嵌套的时候，会把上一个prevContextValue push到stack中
 *
 */
const prevContextValueStack: any[] = [];

/** 推入context beginwork调用 */
export function pushContext(context: Context<any>, newValue: any) {
  // 保存prevContextValue
  prevContextValueStack.push(prevContextValue);

  // 保存currentValue
  prevContextValue = context._currentValue;

  // 给context设置newValue
  context._currentValue = newValue;
}

/** 弹出context completetwork调用 */
export function popContext(context: Context<any>) {
  /** 从prevContextValue恢复context的值 */
  context._currentValue = prevContextValue;

  /** 从prevContextValueStack恢复prevContextValue的值 */
  prevContextValue = prevContextValueStack.pop();
}

可以看到，设置了一个prevContextValue作为中间变量，你可以理解为一个类似缓存的设计，每次Push的时候，都会先将 prevContextValue的值推入栈，并且将最新的值爆存在这个变量中，Pop相反。这个过程也会将最新的值爆存在 Context._currentValue中。

 CompleteWork阶段

归的阶段，就是将Context从上下文栈中弹出，如下:

// completeWork.ts 
... 
case ContextProvider:
      // pop Context
      popContext(wip.type._context);
      bubbleProperties(wip);
      return null;

获取Context的值

看了上面的Context对象的结构你就能猜出来，可以通过 Context._currentValue的方式获取Contetx的值，但是一般不希望使用者直接操作内置对象，所以React暴露了两种操作方式

• Context.Consumer
• useContext hook 

Consumer

Consumer的使用方式很简单，就是传入一个函数作为children，Consumer会调用这个函数并且传入当前上下文的值,如下

<Context.Consumer>{
    contextValue => {
        ...
    }
}</Context.Consumer>

在BeginWork阶段，我们需要对Consumer类型的Fiber单独处理，这个由UpdateContextConsumer实现，如下

/** 更新Consumer */
function updateContextConsumer(wip: FiberNode, renderLane: Lane) {
  const context = wip.type?._context;
  const pendingProps = wip.pendingProps || {};
  const consumerFn = pendingProps.children;

  if (typeof consumerFn === "function") {
    const children = consumerFn(context?._currentValue);
    reconcileChildren(wip, children);
    return wip.child;
  }

  return null;
}

原理很简单，就是把children取出来，作为一个函数运行，并且传入当前Context._currentValue

 useContext

useContext作为一个hooks函数，在mount和update阶段，都调用readContext函数。

readContext函数调用readContextImpl函数

readContextImpl用来实现读取._currentValue 实现如下, 很简单，就不赘述了

// fiberContext.ts
/** 获取context内容
 *  还需要生成新的dependencies
 */
export function readContextImpl<T>(consumer: FiberNode, context: Context<T>) {
  if (!consumer) {
    throw new Error("useContext Error: 请不要在函数组件外部调用hooks");
  }

  return context._currentValue;
}

通过上面的逻辑，我们就完成了Context的最基本功能，但是下面还要解决一个性能优化的问题

Context兼容Bailout

在不引入bailout策略的情况下，父节点的内容变动，整个fiber树都会重新渲染。

这种情况下，某个Context变动，整棵树所有的Fiber节点都会重新渲染，能保证Context的更新被所有子节点感知到。

引入bailout策略之后，react每次更新会跳过那些 props type state 不变化的Fiber，那么设想，如果父Context变动，子节点都设置了 memo包裹，那么这些子节点中，即便有读取Context内容的节点也会被bailout策略跳过，因为其 state props type都没变化。

我们需要一种机制，当上层的Context变化的时候，通知其所有订阅的子节点跳过bailout策略。

我们需要在Fiber节点上设置一个属性，记录当前的Fiber节点（函数节点）订阅了哪些Context

我们在Fiber节点上加入 dependences

/** Fiber节点类 */
export class FiberNode {
  ... 
  /** Fiber依赖的Context */
  dependencies: Dependencies<any>;

/** 记录当前Fiber对象依赖的Context
 *   firstContext 依赖的第一个Context
 *   lanes 所有依赖context对应更新的
 */
export type Dependencies<T> = {
  firstContext: ContextItem<T>;
  lanes: Lanes;
};

/** context在dependencies中以链表的方式连接 其链表项为 ContextItem */
export type ContextItem<T> = {
  context: Context<T>;
  next: ContextItem<T>;
  memorizedState: T;
};

dependencies 是一个对象，其中firstContext维护了一个Context线性链表，以及这些Context在本次更新中的lanes

ContextItem记录了当前的Context对象以及next，memorizedState记录了_currentValue

我们需要在一个地方统一的挂载dependencies，我们发现上面实现的readContextImpl可以用来记录当前fiber有哪些订阅的Context, 修改上述readContextImpl实现

/** 获取context内容
 *  还需要生成新的dependencies
 */
export function readContextImpl<T>(consumer: FiberNode, context: Context<T>) {
  if (!consumer) {
    throw new Error("useContext Error: 请不要在函数组件外部调用hooks");
  }
  // 建立新的dependencies
  const contextItem: ContextItem<T> = {
    context,
    next: null,
    memorizedState: context._currentValue,
  };

  // 绑定到consumer
  if (lastContextDepItem === null) {
    // 第一个contextItem
    consumer.dependencies = {
      firstContext: contextItem,
      lanes: NoLanes,
    };

    lastContextDepItem = contextItem;
  } else {
    // 不是第一个
    lastContextDepItem = lastContextDepItem.next = contextItem;
  }

  return context._currentValue;
}

在每次读取Context值的时候，都会在其Fiber上挂depenencies，在下次更新中，就可以得知Fiber订阅了哪些Context了 

在每次便利到Provider时，我们需要对比新的Value是否有变化，这个对比通过Object.is完成，即对比内存地址。

如果没有变化，就bailout跳出，如果变化了，就需要 “通知” 所有订阅了这个Context的组件，这个通知过程由 propagateContextChange完成

/** 更新ContextProvider */
function updateContextProvider(wip: FiberNode, renderLane: Lane) {
  const context = wip.type._context;
  const memorizedProps = wip.memorizedProps;
  const pendingProps = wip.pendingProps;
  const newValue = pendingProps?.value;
  const oldValue = memorizedProps?.value;
  // 推入Context
  pushContext(context, newValue);

  // TODO bailout逻辑
  if (
    Object.is(oldValue, newValue) &&
    memorizedProps.children === pendingProps.children
  ) {
    /** 两次value相等 并且children不能变化，children变化 哪怕value不变 也要更新下面的Fiber */
    return bailoutOnAlreadyFinishedWork(wip, renderLane);
  } else {
    /** 传播Context变化 */
    propagateContextChange(wip, context, renderLane);
  }

  // reconcile child
  reconcileChildren(wip, pendingProps.children);
  return wip.child;
}

propagateContextChange的逻辑就是，从当前的Provider开始，开启一个子深度优先遍历，遍历其子树，对于每个Fiber子节点都检查其dependcies，如果包含了当前Context就在其Fiber.lane加入当前更新的优先级lane，这样在遍历到这个节点的时候，就会因为其存在状态 state，不会bailout跳过。 同时，这个遍历的过程向下终止到 叶子结点 或者 相同Context的Provider，因为相同Context Provider下的子节点就不由当前的Provider管理了，实现如下：

/** 传播context变化 */
export function propagateContextChange(
  wip: FiberNode,
  context: Context<any>,
  renderLane: Lane
) {
  let nextFiber = wip.child;
  while (nextFiber !== null && nextFiber !== wip) {
    const deps = nextFiber.dependencies;
    if (deps) {
      let contextItem = deps.firstContext;
      while (contextItem !== null) {
        if (contextItem.context === context) {
          // 找到对应的Context 设置lane
          // 设置fiber和alternate的lane
          nextFiber.lanes = mergeLane(nextFiber.lanes, renderLane);
          if (nextFiber.alternate !== null) {
            nextFiber.alternate.lanes = mergeLane(
              nextFiber.alternate.lanes,
              renderLane
            );
          }
          // 从当前Fiber到Provide的Fiber 标记childLanes
          scheduleContextOnParentPath(nextFiber, wip, renderLane);
          // 设置deps.lane
          deps.lanes = mergeLane(deps.lanes, renderLane);
          // break
          break;
        }else{
          contextItem =contextItem.next
        }
      }
    } else if (
      ((nextFiber.tag === ContextProvider && nextFiber.type !== wip.type) ||
        nextFiber.tag !== ContextProvider) &&
      nextFiber.child !== null
    ) {
      nextFiber = nextFiber.child;
      continue;
    }
    // 回溯
    while (nextFiber.sibling === null) {
      if (nextFiber.return === null || nextFiber.return === wip) {
        return // 直接return 跳出所有循环
      }
      nextFiber = nextFiber.return;
    }

    nextFiber = nextFiber.sibling;
  }
}

/** 在parent路径上调度Context
 * 从当前找到context的节点，到provider节点 标记childLanes
 */
function scheduleContextOnParentPath(
  from: FiberNode,
  to: FiberNode,
  renderLane: Lane
) {
  if (from === to) return;
  let parentNode = from.return;
  while (parentNode !== null && from !== to) {
    parentNode.childLanes = mergeLane(parentNode.childLanes, renderLane);
    if (parentNode.alternate !== null) {
      parentNode.alternate.childLanes = mergeLane(
        parentNode.alternate.childLanes,
        renderLane
      );
    }
    parentNode = parentNode.return;
  }
}

我们当前的实现不考虑类组件，所以Context都是在函数组件中完成读取的，这就要我们在开始渲染函数组件之前，先检查这个函数是否包含Context的变化，这个逻辑由 prepareToReadContext完成。

/** 处理函数节点的比较 */
function updateFunctionComponent(
  wip: FiberNode,
  Component: Function,
  renderLane: Lane
): FiberNode {
  /** 重制dependencies信息 */
  prepareToReadContext(wip, renderLane);
  // renderWithHooks 中检查，如果状态改变 则置didReceiveUpdate = true
  const nextChildElement = renderWithHooks(wip, Component, renderLane);
  if (wip.alternate !== null && !didReceiveUpdate) {
    // bailout
    // 重置hook
    bailoutHook(wip, renderLane);
    return bailoutOnAlreadyFinishedWork(wip, renderLane);
  }
  reconcileChildren(wip, nextChildElement);
  return wip.child;
}

这个函数有两个功能，一个是在进入当前函数节点之前判断当前函数节点是否存在因为Context变动而发生的更新 一个是清空dependcies，准备开始下一轮的Context订阅收集

/** readContext之前的准备工作  每次函数调用之前执行 重新设置dependencies*/
export function prepareToReadContext(wip: FiberNode, renderLane: Lane) {
  lastContextDepItem = null; // 置空
  if (wip.dependencies !== null) {
    if (isSubsetOfLanes(wip.dependencies.lanes, renderLane)) {
      // 当前函数组件内的Context 包含当前renderLane更新优先级的更新
      markWipReceiveUpdate(); // 当前函数组件 不能bailout
    }
    wip.dependencies = null; // 置空
  }
}

可以看到，在popagate阶段，会讲所有的Context更新对应的lane合并到 ContextItem的lane上，所以在 prepareToReadContext 中，只需要检查一下当前渲染的lane，是否为ContextItem.lanes的子lane即可判断当前函数组件是否存在因Context变动触发的更新。

最后，对Consumer组件的渲染，我们也进行改造，把通过 context._currentValue读取Context值，统一改成 readContextImpl

/** 更新Consumer */
function updateContextConsumer(wip: FiberNode, renderLane: Lane) {
  const context = wip.type?._context;
  const pendingProps = wip.pendingProps || {};
  const consumerFn = pendingProps.children;

  if (typeof consumerFn === "function") {
    const children = consumerFn(readContextImpl(wip,context));
    reconcileChildren(wip, children);
    return wip.child;
  }

  return null;
}

通过这种方式，我们就可以让Context兼容bailout策略