React 实现 useReducer

本文将带大家实现 useReducer。

先看下如何使用。

js 复制代码

function FunctionComponent() {
  const [count1, setCount1] = useReducer((x) => x + 1, 0);

  return (
    <div className="border">
      <button
        onClick={() => {
          setCount1();
        }}
      >
        {count1}
      </button>
    </div>
  );
}

ReactDOM.createRoot(document.getElementById("root") as HTMLElement).render(
  (<FunctionComponent />) as any
);

Render 阶段

BeginWork 阶段

reconcileSingleElement 增加了判断节点是否可复用的逻辑。

js 复制代码

export function createWorkInProgress(current: Fiber, pendingProps: any): Fiber {
  let workInProgress = current.alternate;
  if (workInProgress === null) {
    workInProgress = createFiber(current.tag, pendingProps, current.key);
    workInProgress.elementType = current.elementType;
    workInProgress.type = current.type;
    workInProgress.stateNode = current.stateNode;

    workInProgress.alternate = current;
    current.alternate = workInProgress;
  } else {
    workInProgress.pendingProps = pendingProps;
    workInProgress.type = current.type;
    workInProgress.flags = NoFlags;
  }

  workInProgress.flags = current.flags;
  workInProgress.childLanes = current.childLanes;
  workInProgress.lanes = current.lanes;

  workInProgress.child = current.child;
  workInProgress.memoizedProps = current.memoizedProps;
  workInProgress.memoizedState = current.memoizedState;
  workInProgress.updateQueue = current.updateQueue;

  workInProgress.sibling = current.sibling;
  workInProgress.index = current.index;

  return workInProgress;
}

function useFiber(fiber: Fiber, pendingProps: any) {
    const clone = createWorkInProgress(fiber, pendingProps);
    clone.index = 0;
    clone.sibling = null;
    return clone;
}

// 协调单个节点，对于页面初次渲染，创建 fiber，不涉及对比复用老节点
// new (1)
// old 2 [1] 3 4
function reconcileSingleElement(
    returnFiber: Fiber,
    currentFirstChild: Fiber | null,
    element: ReactElement
) {
    // 节点复用的条件
    // 1. 同一层级下 2. key 相同 3. 类型相同
    // element 和 currentFirstChild 对应同一个父级，第一个条件满足
    const key = element.key;
    let child = currentFirstChild;

    while (child !== null) {
      if (child.key === key) {
        const elementType = element.type;
        // 可以复用
        if (child.elementType === elementType) {
          // todo 后面其它 fiber 可以删除了
          
          const existing = useFiber(child, element.props);
          existing.return = returnFiber;
          return existing;
        } else {
          // 前提：React 不认为同一层级下有两个相同的 key 值
          break;
        }
      } else {
        // 删除单个节点
        // deleteChild(returnFiber, child);
      }
      // 老 fiber 节点是单链表
      child = child.sibling;
    }

    let createdFiber = createFiberFromElement(element);
    createdFiber.return = returnFiber;
    return createdFiber;
}

CompleteWork 阶段

因为现在还没实现合成事件，可以先通过 addEventListener 模拟事件。

finalizeInitialChildren 遍历新旧 props，更新属性，

js 复制代码

// 初始化、更新属性
// old {className='red', onClick:f, data-red='red'}
// new {className='red green', onClick:f}
function finalizeInitialChildren(
  domElement: Element,
  prevProps: any,
  nextProps: any
) {
  // 遍历老的 props
  for (const propKey in prevProps) {
    const prevProp = prevProps[propKey];
    if (propKey === "children") {
      if (isStr(prevProp) || isNum(prevProp)) {
        // 属性
        domElement.textContent = "";
      }
    } else {
      // 设置属性
      if (propKey === 'onClick') {
          // 模拟事件
          document.removeEventListener("click", prevProp)
      } else {
          // 如果新的 props 中没有，直接删除
          if(!(prevProp in nextProps)){
              (domElement as any)[propKey] = "";
          }
      }
    }
  }
  
  // 遍历新的 props
  for (const propKey in nextProps) {
    const nextProp = nextProps[propKey];
    if (propKey === "children") {
      if (isStr(nextProp) || isNum(nextProp)) {
        // 属性
        domElement.textContent = nextProp + "";
      }
    } else {
      // 设置属性
      if (propKey === 'onClick') {
          // 模拟事件
          document.addEventListener("click", nextProp)
      } else {
          (domElement as any)[propKey] = nextProp;
      }
    }
  }
}

updateFunctionComponent 增加 renderWithHooks。

js 复制代码

// 当前正在工作的函数组件的 fiber
let currentlyRenderingFiber: Fiber | null = null;
// 当前正在工作的 hook
let workInProgressHook: Hook | null = null;
let currentHook: Hook | null = null;

function finishRenderingHooks() {
  currentlyRenderingFiber = null;
  currentHook = null;
  workInProgressHook = null;
}

export function renderWithHooks<Props>(
  current: Fiber | null,
  workInProgress: Fiber,
  Component: any,
  props: Props
): any {
  currentlyRenderingFiber = workInProgress;
  workInProgress.memoizedState = null;

  let children = Component(props);

  finishRenderingHooks();

  return children;
}

function updateFunctionComponent(current: Fiber | null, workInProgress: Fiber) {
  const { type, pendingProps } = workInProgress;
  // 函数执行结果就是 children
  const children = renderWithHooks(current, workInProgress, type, pendingProps);
  reconcileChildren(current, workInProgress, children);
  return workInProgress.child;
}

completeWork 区分挂载和更新，更新时调用 updateHostComponent。

js 复制代码

function updateHostComponent(
  current: Fiber | null,
  workInProgress: Fiber,
  type: string,
  newProps: any
) {
  if (current?.memoizedProps === newProps) {
    return;
  }

  finalizeInitialChildren(
    workInProgress.stateNode as Element,
    current?.memoizedProps,
    newProps
  );
}

function completeWork(
  current: Fiber | null,
  workInProgress: Fiber
): Fiber | null {
  const newProps = workInProgress.pendingProps;
  switch (workInProgress.tag) {
    case HostRoot: {
      return null;
    }
    case HostComponent: {
      // 原生标签,type 是标签名
      const { type } = workInProgress;
      if(current !== null && workInProgress.stateNode !== null) {
          updateHostComponent(current, workInProgress, type, newProps)
      } else {
        // 1. 创建真实 DOM
        const instance = document.createElement(type);
        // 2. 初始化 DOM 属性
        finalizeInitialChildren(instance, null, newProps);
        // 3. 把子 dom 挂载到父 dom 上
        appendAllChildren(instance, workInProgress);
        workInProgress.stateNode = instance;
      }
      return null;
    }
  }

  throw new Error(
    `Unknown unit of work tag (${workInProgress.tag}). This error is likely caused by a bug in ` +
      "React. Please file an issue."
  );
}

实现 useReducer

定义 Hook 类型。

js 复制代码

type Hook = {
  memoizedState: any; // 不同情况下，取值也不同，useState/useReducer 存的是 state，useEffect/useLayoutEffect 存的是 effect 单向循环链表
  next: null | Hook;
};

构建 hook 链表

js 复制代码

// 1. 返回当前 useX 函数对应的 hook
// 2. 构建 hook 链表
// 源码中把 mount 阶段和 update 阶段实现在了两个函数中，这里简化一下，放在一个函数中实现
function updateWorkInProgressHook(): Hook {
  let hook: Hook;

  const current = currentlyRenderingFiber?.alternate

  if (current) {
    // update 阶段，复用老的
    currentlyRenderingFiber!.memoizedState = current.memoizedState;

    if (workInProgressHook != null) {
      workInProgressHook = hook = workInProgressHook.next!
      currentHook = currentHook?.next as Hook;
    } else {
      // hook 单链表的头节点
      hook = workInProgressHook = currentlyRenderingFiber?.memoizedState
      currentHook = current.memoizedState;
    }
  } else {
    // mount 阶段
    currentHook = null;
    hook = {
      memoizedState: null,
      next: null
    }

    // 构建 hook 链表
    if (workInProgressHook) {
      workInProgressHook = workInProgressHook.next = hook
    } else {
      // hook 单链表的头节点
      workInProgressHook = currentlyRenderingFiber?.memoizedState = hook
    }
  }

  return hook;
}

实现调度更新

js 复制代码

// 根据 sourceFiber 找根节点
function getRootForUpdatedFiber(sourceFiber: Fiber): FiberRoot {
  let node = sourceFiber;
  let parent = node.return;

  while (parent !== null) {
    node = parent;
    parent = node.return;
  }

  return node.tag === HostRoot ? node.stateNode : null;
}

function dispatchReducerAction<S, I, A>(
  fiber: Fiber,
  hook: Hook,
  reducer: ((state: S, action: A) => S) | null,
  action: any
) {
  hook.memoizedState = reducer ? reducer(hook.memoizedState, action) : action;

  const root = getRootForUpdatedFiber(fiber);

  fiber.alternate = { ...fiber };
  if (fiber.sibling) {
    fiber.sibling.alternate = fiber.sibling;
  }

  scheduleUpdateOnFiber(root, fiber, true);
}

整合前两步，实现 useReducer 函数。

js 复制代码

export function useReducer<S, I, A>(
  reducer: ((state: S, action: A) => S) | null,
  initialArg: I,
  init?: (initialArg: I) => S
) {
  // 1.  构建 hook 链表(mount、update)
  const hook: Hook = updateWorkInProgressHook();

  let initialState: S;
  if (init !== undefined) {
    initialState = init(initialArg);
  } else {
    initialState = initialArg as any;
  }

  // 2. 区分函数组件是初次挂载还是更新
  if(!currentlyRenderingFiber?.alternate) {
    hook.memoizedState = initialState;
  }

  // 3. dispatch
  const dispatch = dispatchReducerAction.bind(
    null,
    currentlyRenderingFiber!,
    hook,
    reducer as any
  )

  return [hook.memoizedState, dispatch];
}