Implement React v18 from Scratch Using WASM and Rust - [21] Performance Optimization for Context

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：v21

The previous article implemented the Context feature but couldn't combine it with performance optimization-related features. In this article, we will address this issue. Let's use the previous example:

const ctx = createContext(0)

export default function App() {
  const [num, update] = useState(0)
  const memoChild = useMemo(() => {
    return <Child />
  }, [])
  console.log('App render ', num)
  return (
    <ctx.Provider value={num}>
      <div
        onClick={() => {
          update(1)
        }}>
        {memoChild}
      </div>
    </ctx.Provider>
  )
}

function Child() {
  console.log('Child render')
  const val = useContext(ctx)

  return <div>ctx: {val}</div>
}

After clicking, the Child component doesn't re-render, and the page doesn't update. The reason is that Child triggers the bailout strategy, but it actually uses context, so we can say that Child depends on the ctx Context. Therefore, we need to add a field in the FiberNode to store the dependent Context:

#[derive(Clone, Debug)]
pub struct ContextItem {
    context: JsValue,
    memoized_state: JsValue,
    next: Option<Rc<RefCell<ContextItem>>>,
}

pub struct FiberDependencies {
    pub first_context: Option<Rc<RefCell<ContextItem>>>,
    pub lanes: Lane,
}
...
pub struct FiberNode {
  ...
  pub dependencies: Option<Rc<RefCell<FiberDependencies>>>,
  ...
}

When should we update this field? Of course, it should be done when calling useContext, which is read_context in fiber_hooks:

// fiber_hooks.rs
use crate::fiber_context::read_context as read_context_origin;
...
fn read_context(context: JsValue) -> JsValue {
    let consumer = unsafe { CURRENTLY_RENDERING_FIBER.clone() };
    read_context_origin(consumer, context)
}

// fiber_context.rs
pub fn read_context(consumer: Option<Rc<RefCell<FiberNode>>>, context: JsValue) -> JsValue {
    if consumer.is_none() {
        panic!("Can only call useContext in Function Component");
    }
    let consumer = consumer.unwrap();
    let value = derive_from_js_value(&context, "_currentValue");

    let context_item = Rc::new(RefCell::new(ContextItem {
        context,
        next: None,
        memoized_state: value.clone(),
    }));

    if unsafe { LAST_CONTEXT_DEP.is_none() } {
        unsafe { LAST_CONTEXT_DEP = Some(context_item.clone()) };
        consumer.borrow_mut().dependencies = Some(Rc::new(RefCell::new(FiberDependencies {
            first_context: Some(context_item),
            lanes: Lane::NoLane,
        })));
    } else {
        let next = Some(context_item.clone());
        unsafe {
            LAST_CONTEXT_DEP.clone().unwrap().borrow_mut().next = next.clone();
            LAST_CONTEXT_DEP = next;
        }
    }
    value
}

This function adds all the Contexts that the current FiberNode depends on to a linked list and attaches it to the dependencies property. For example, in the following example:

function App() {
  const value1 = useContext(ctx1)
  const value2 = useContext(ctx2)
}

The data structure would look like this:

Next, let's continue with the update process. First, we need to modify update_context_provider in the begin work phase. The most important change is the following line:

propagate_context_change(work_in_progress.clone(), context, render_lane);

From the function name, it should notify other FiberNodes that depend on this Context that its value has changed. Specifically, in propagate_context_change, it searches for FiberNodes that depend on this Context in a depth-first manner. It adds render_lane to their lanes and to the child_lanes of their ancestors (up to the Context's Provider), as shown in the diagram:

During the traversal, if it encounters a Provider of the same Context, it skips it and its subtree, as shown in the diagram:

This is because Context follows the principle of proximity. For example, in the following example, the value of ctx in Child comes from the ancestor node closest to it. Therefore, changes in the outermost Provider do not affect it:

function App() {
  return (
    <ctx.Provider value='a'>
      <ctx.Provider value='b'>
        <div>
          <Child />
        </div>
      </ctx.Provider>
    </ctx.Provider>
  )
}

Now, when the begin work phase reaches these FiberNodes, they will be skipped because their child_lanes include render_lane. However, they will still reach the FiberNode that depends on the Context because its lanes include render_lane. As a result, the corresponding component will re-render and execute useContext to obtain the new value.

With these changes, the example at the beginning of the article should work correctly. You can find the detailed modifications here.

Please kindly give me a star!