Implement React v18 from Scratch Using WASM and Rust - [20] Implement 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：v20

Context is also a very important feature in React, so we need to include it in our WASM version as well. Let's recall how we usually use Context:

import {createContext, useContext} from 'react'

const ctx = createContext('A')

export default function App() {
  return (
    <ctxA.Provider value={'B'}>
      <Cpn />
    </ctxA.Provider>
  )
}

function Cpn() {
  const value = useContext(ctx)
  return <div>{value}</div>
}

So, we need to first export two methods from the React library:

#[wasm_bindgen(js_name = useContext)]
pub unsafe fn use_context(context: &JsValue) -> Result<JsValue, JsValue> {
  ...
}

#[wasm_bindgen(js_name = createContext)]
pub unsafe fn create_context(default_value: &JsValue) -> JsValue {
    let context = Object::new();
    Reflect::set(
        &context,
        &"$$typeof".into(),
        &JsValue::from_str(REACT_CONTEXT_TYPE),
    );
    Reflect::set(&context, &"_currentValue".into(), default_value);
    let provider = Object::new();
    Reflect::set(
        &provider,
        &"$$typeof".into(),
        &JsValue::from_str(REACT_PROVIDER_TYPE),
    );
    Reflect::set(&provider, &"_context".into(), &context);
    Reflect::set(&context, &"Provider".into(), &provider);
    context.into()
}

The code inside create_context translates to the following in JavaScript:

const context {
  $$typeof: REACT_CONTEXT_TYPE,
  Provider: null,
  _currentValue: defaultValue,
}
context.Provider = {
  $$typeof: REACT_PROVIDER_TYPE,
  _context: context,
}
return context

As you can see, ctxA.Provider is a new type of FiberNode, and we need to add a new branch to handle it. Following the flow sequence, the first step is begin_work:

fn update_context_provider(
    work_in_progress: Rc<RefCell<FiberNode>>,
) -> Option<Rc<RefCell<FiberNode>>> {
    let provider_type = { work_in_progress.borrow()._type.clone() };
    let context = derive_from_js_value(&provider_type, "_context");
    let new_props = { work_in_progress.borrow().pending_props.clone() };
    push_provider(&context, derive_from_js_value(&new_props, "value"));
    let next_children = derive_from_js_value(&new_props, "children");
    reconcile_children(work_in_progress.clone(), Some(next_children));
    work_in_progress.clone().borrow().child.clone()
}

The part that is difficult to understand here is push_provider.

static mut PREV_CONTEXT_VALUE: JsValue = JsValue::null();
static mut PREV_CONTEXT_VALUE_STACK: Vec<JsValue> = vec![];

pub fn push_provider(context: &JsValue, new_value: JsValue) {
    unsafe {
        PREV_CONTEXT_VALUE_STACK.push(PREV_CONTEXT_VALUE.clone());
        PREV_CONTEXT_VALUE = Reflect::get(context, &"_currentValue".into()).unwrap();
        Reflect::set(context, &"_currentValue".into(), &new_value);
    }
}

Correspondingly, there is also a pop_provider that goes along with it.

pub fn pop_provider(context: &JsValue) {
    unsafe {
        Reflect::set(context, &"_currentValue".into(), &PREV_CONTEXT_VALUE);
        let top = PREV_CONTEXT_VALUE_STACK.pop();
        if top.is_none() {
            PREV_CONTEXT_VALUE = JsValue::null();
        } else {
            PREV_CONTEXT_VALUE = top.unwrap();
        }
    }
}

It will be called within complete_work.

WorkTag::ContextProvider => {
  let _type = { work_in_progress.borrow()._type.clone() };
  let context = derive_from_js_value(&_type, "_context");
  pop_provider(&context);
  self.bubble_properties(work_in_progress.clone());
  None
}

We will clarify this portion of the code through the following example:

const ctxA = createContext('A0')
const ctxB = createContext('B0')

export default function App() {
  return (
    <ctxA.Provider value='A1'>
      <ctxB.Provider value='B1'>
        <ctxA.Provider value='A2'>
          <ctxB.Provider value='B2'>
            <Child />
          </ctxB.Provider>
          <Child />
        </ctxA.Provider>
        <Child />
      </ctxB.Provider>
      <Child />
    </ctxA.Provider>
  )
}

function Child() {
  const a = useContext(ctxA)
  const b = useContext(ctxB)
  return (
    <div>
      A: {a} B: {b}
    </div>
  )
}

The expected result of the above example should be:

A: A2 B: B2
A: A2 B: B1
A: A1 B: B1
A: A1 B: B0

Let's analyze it. According to the flow, when begin_work reaches the bottommost Child, it goes through four push_provider operations, and the state of the FiberNode is as follows:

When it reaches the third level Child, it performs one pop_provider operation, and the state becomes:

When it reaches the second level Child, it performs another pop_provider operation, and the state becomes:

When it reaches the first level Child, it performs the final pop_provider operation, and the state becomes:

The reason it may be difficult to understand is that it stores the values of multiple contexts in a single stack. You can go through this example a few more times to better understand it.

Once you understand this, the basic flow of Context is mostly covered. However, there is also useContext, which is quite simple. You can add the relevant code following the flow of other Hooks we discussed earlier, and the core is the read_context method in fiber_hooks.

fn read_context(context: JsValue) -> JsValue {
  let consumer = unsafe { CURRENTLY_RENDERING_FIBER.clone() };
  if consumer.is_none() {
      panic!("Can only call useContext in Function Component");
  }
  let value = derive_from_js_value(&context, "_currentValue");
  value
}

With these changes, the example mentioned above can be executed. For more details about this update, please refer to here.

However, the current implementation of Context is not yet complete. It can cause issues when combined with performance optimization-related features. For example, consider the following 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 does not re-render, and the page does not update. The reason is that the Child component hits the bailout strategy. However, the Child component actually uses context, and the value of the context has changed. The Child component should be re-rendered. We will address this issue in the next article.

Please kindly give me a star!