从零实现 React v18，但 WASM 版 - [22] 实现 memo

Posted on | In

模仿 big-react，使用 Rust 和 WebAssembly，从零实现 React v18 的核心功能。深入理解 React 源码的同时，还锻炼了 Rust 的技能，简直赢麻了！

代码地址：https://github.com/ParadeTo/big-react-wasm

本文对应 tag：v22

Based on big-react，I am going to implement React v18 core features from scratch using WASM and Rust.

Code Repository：https://github.com/ParadeTo/big-react-wasm

The tag related to this article：v22

前面几篇文章都是围绕 React 性能优化相关的特性展开的，不过还差一个 memo，今天就来实现一下。以下面代码为例：

The previous articles were focused on exploring performance optimization features related to React. However, there is still one missing feature: memo. Today, let’s implement it. Take the following code as an example:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import {useState, memo} from 'react'

export default function App() {
  const [num, update] = useState(0)
  console.log('App render', num)
  return (
    <div onClick={() => update(num + 1)}>
      <Cpn num={num} name={'cpn1'} />
      <Cpn num={0} name={'cpn2'} />
    </div>
  )
}

const Cpn = memo(function ({num, name}) {
  console.log('render', name)
  return (
    <div>
      {name}: {num}
      <Child />
    </div>
  )
})

function Child() {
  console.log('Child render')
  return <p>i am child</p>
}

首次渲染时，会打印：

When initially rendered, the following will be printed:

1
2
3
4
5
App render 0
render cpn1
Child render
render cpn2
Child render

点击后，应该只有第一个 Cpn 组件会重新渲染，控制台打印：

After clicking, only the first Cpn component should be re-rendered, and the console will print:

1
2
3
App render 1
render cpn1
Child render

下面我们来看看要怎么实现。

Now let’s see how to implement this.

首先，需要从 react 这个库中导出 memo 方法，如下所示：

First, we need to import the memo method from the React library, as shown below:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
#[wasm_bindgen]
pub unsafe fn memo(_type: &JsValue, compare: &JsValue) -> JsValue {
    let fiber_type = Object::new();

    Reflect::set(
        &fiber_type,
        &"$$typeof".into(),
        &JsValue::from_str(REACT_MEMO_TYPE),
    );
    Reflect::set(&fiber_type, &"type".into(), _type);

    let null = JsValue::null();
    Reflect::set(
        &fiber_type,
        &"compare".into(),
        if compare.is_undefined() {
            &null
        } else {
            compare
        },
    );
    fiber_type.into()
}

翻译成 JS 的话，是这样：

In JavaScript, the translation would be as follows:

1
2
3
4
5
6
7
8
9
10
11
export function memo(
	type: FiberNode['type'],
	compare?: (oldProps: Props, newProps: Props) => boolean
) {
	const fiberType = {
		$$typeof: REACT_MEMO_TYPE,
		type,
		compare: compare === undefined ? null : compare
	};
	return fiberType;
}

跟之前的 context Provider 类似，这里也是返回了一个对象，并且把传入的组件保存在了 type 字段中，同时把第二个参数存在了 compare 字段中，该字段的作用应该都清楚，就不赘述了。很明显，这里又是一个新的 FiberNode 类型，我们需要在 begin work 中增加对该类型的处理：

Similar to the previous context Provider, here we also return an object. The passed-in component is saved in the type field, and the second argument is stored in the compare field. The purpose of the compare field should be clear, so I won’t elaborate on it. Clearly, this is a new FiberNode type, and we need to add handling for this type in the begin work phase.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

fn update_memo_component(
    work_in_progress: Rc<RefCell<FiberNode>>,
    render_lane: Lane,
) -> Result<Option<Rc<RefCell<FiberNode>>>, JsValue> {
    let current = { work_in_progress.borrow().alternate.clone() };
    let next_props = { work_in_progress.borrow().pending_props.clone() };

    if current.is_some() {
        let current = current.unwrap();
        let prev_props = current.borrow().memoized_props.clone();
        if !check_scheduled_update_or_context(current.clone(), render_lane.clone()) {
            let mut props_equal = false;
            let compare = derive_from_js_value(&work_in_progress.borrow()._type, "compare");
            if compare.is_function() {
                let f = compare.dyn_ref::<Function>().unwrap();
                props_equal = f
                    .call2(&JsValue::null(), &prev_props, &next_props)
                    .unwrap()
                    .as_bool()
                    .unwrap();
            } else {
                props_equal = shallow_equal(&prev_props, &next_props);
            }

            if props_equal && Object::is(&current.borrow()._ref, &work_in_progress.borrow()._ref) {
                unsafe { DID_RECEIVE_UPDATE = false };
                work_in_progress.borrow_mut().pending_props = prev_props;
                work_in_progress.borrow_mut().lanes = current.borrow().lanes.clone();
                return Ok(bailout_on_already_finished_work(
                    work_in_progress.clone(),
                    render_lane,
                ));
            }
        }
    }
    let Component = { derive_from_js_value(&work_in_progress.borrow()._type, "type") };
    update_function_component(work_in_progress.clone(), Component, render_lane)
}

这里的代码很好懂，如果有 current，说明不是首次渲染，可以看是否可以进行性能优化。

首先还是通过 check_scheduled_update_or_context 判断子孙组件中是否有满足这次更新优先级的节点，如果没有则进行 memo 相关的性能优化，具体来说为：

  • 获取 compare 函数，如果没有则用默认的 shallow_equal（该函数用于对比两个对象是否相等，key 相同且其值相同时两个对象相等，值的比较为浅比较）
  • 将新旧 props 传入上面得到的函数
  • 如果 compare 返回 true，则进入 bailout 逻辑

否则，进入 update_function_component 逻辑，因为 memo 只是在 FunctionComponent 外面多套了一层而已。注意到这里的 update_function_component 的参数跟之前不一样了，之前只有 work_in_progress 和 render_lane 是因为只考虑 FunctionComponent 的情况下，可以从 work_in_progress 的 _type 中获取 Component，现在加入了 MemoComponent，则需要从 work_in_progress 的 _type 中的 type 来获取 Component。

其他比较细微的改动就不介绍了，详情请见这里。

The code here is easy to understand. If there is a current value, it means it’s not the initial render, so we can check if there are any nodes in the descendant components that meet the priority of this update, and if not, we can perform performance optimizations related to memoization. Specifically:

  • We retrieve the compare function, and if it doesn’t exist, we use the default shallow_equal function (which compares two objects for equality by comparing their keys and values, performing a shallow comparison).
  • We pass the new and old props to the function obtained above.
  • If the compare function returns true, we enter the bailout logic.

Otherwise, we enter the update_function_component logic because memo is just an additional layer outside the FunctionComponent. Note that the parameters for update_function_component are different now. Previously, we only had work_in_progress and render_lane because we only considered the case of a FunctionComponent, where we could retrieve the Component from work_in_progress‘s _type. Now, with the addition of MemoComponent, we need to retrieve the Component from work_in_progress‘s _type‘s type.

I won’t go into other minor changes, but you can find more details here.