Implement React v18 from Scratch Using WASM and Rust - [22] Implement memo

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

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:

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:

App render 0
render cpn1
Child render
render cpn2
Child render

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

App render 1
render cpn1
Child render

Now let's see how to implement this.

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

#[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()
}

In JavaScript, the translation would be as follows:

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;
}

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.

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)
}

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.