Implement React v18 from Scratch Using WASM and Rust - [26] Implement React.lazy

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

Suspense has another useful feature that combines with React.lazy for component lazy loading. Let's continue implementing it. The changes can be found in detail here.

We'll use the following example to illustrate:

import { Suspense, lazy } from 'react';

function delay(promise) {
  return new Promise((resolve) => {
    setTimeout(() => {
      resolve(promise);
    }, 2000);
  });
}

const Cpn = lazy(() =>
  import('./Cpn').then((res) => delay(res))
);

export default function App() {
  return (
    <Suspense fallback={<div>loading</div>}>
      <Cpn />
    </Suspense>
  );
}

First, we need to export this method in the react library:

#[wasm_bindgen]
pub fn lazy(ctor: &JsValue) -> JsValue {
    let payload = Object::new();
    Reflect::set(&payload, &"_status".into(), &JsValue::from(UNINITIALIZED));
    Reflect::set(&payload, &"_result".into(), ctor);

    let lazy_type = Object::new();

    Reflect::set(
        &lazy_type,
        &"$typeof".into(),
        &JsValue::from_str(REACT_LAZY_TYPE),
    );
    Reflect::set(&lazy_type, &"_payload".into(), &payload);
    let closure = Closure::wrap(
        Box::new(lazy_initializer) as Box<dyn Fn(JsValue) -> Result<JsValue, JsValue>>
    );
    let f = closure.as_ref().unchecked_ref::<Function>().clone();
    closure.forget();
    Reflect::set(&lazy_type, &"_init".into(), &f);
    lazy_type.into()
}

The translation to JavaScript is as follows:

const payload = {
  _status: UNINITIALIZED,
  _result: ctor,
};

const lazy_type = {
  $typeof: REACT_LAZY_TYPE,
  _payload: payload,
  _init: lazy_initializer,
};

Here, the focus is on the lazy_initializer function. Let's explain it using the JavaScript version:

function lazy_initializer(payload) {
  if (payload._status === Uninitialized) {
    const ctor = payload._result;
    const thenable = ctor();
    thenable.then(
      (moduleObject) => {
        payload._status = Resolved;
        payload._result = moduleObject;
      },
      (error) => {
        payload._status = Rejected;
        payload._result = error;
      }
    );

    payload._status = Pending;
    payload._result = thenable;
  }

  if (payload._status === Resolved) {
    const moduleObject = payload._result;
    return moduleObject.default;
  } else {
    throw payload._result;
  }
}

This function is similar to the "use" hook implemented in the previous article. Here, ctor corresponds to () => import('./Cpn').then((res) => delay(res)). When executed, it returns a Promise object. It only returns the result (res) when the object's status is Resolved. In this case, res is a module object, and its default property contains the content exported using export default. For other statuses, it throws _result. When the status is Pending, _result is the Promise object itself, and when the status is Rejected, _result is an error object.

Next, the main file to be modified is begin_work.rs:

....
        WorkTag::LazyComponent => update_lazy_component(work_in_progress.clone(), render_lane),
    };
}

fn update_lazy_component(
    work_in_progress: Rc<RefCell<FiberNode>>,
    render_lane: Lane,
) -> Result<Option<Rc<RefCell<FiberNode>>>, JsValue> {
    let lazy_type = { work_in_progress.borrow()._type.clone() };
    let payload = derive_from_js_value(&lazy_type, "_payload");
    let init_jsvalue = derive_from_js_value(&lazy_type, "_init");
    let init = init_jsvalue.dyn_ref::<Function>().unwrap();
    // return value OR throw
    let Component = init.call1(&JsValue::null(), &payload)?;
    work_in_progress.borrow_mut()._type = Component.clone();
    work_in_progress.borrow_mut().tag = WorkTag::FunctionComponent;
    let child = update_function_component(work_in_progress, Component.clone(), render_lane);
    child
}
....

The key line here is let Component = init.call1(&JsValue::null(), &payload)?;. When init is executed and throws an exception, according to the flow in the previous article, it will find the nearest Suspense component and restart the render process, rendering the fallback of Suspense. When the Promise object is resolved, the update process will be triggered again, and when it reaches this line, the returned value of init will be the exported component, which is Cpn.

In addition, handle_throw in work_loop.rs needs to be modified to handle non-use errors:

fn handle_throw(root: Rc<RefCell<FiberRootNode>>, mut thrown_value: JsValue) {
    /*
        throw possibilities:
            1. use thenable
            2. error (Error Boundary), lazy
    */
    if Object::is(&thrown_value, &SUSPENSE_EXCEPTION) {
        unsafe { WORK_IN_PROGRESS_SUSPENDED_REASON = SUSPENDED_ON_DATA };
        thrown_value = get_suspense_thenable();
    } else {
        let is_wakeable = !thrown_value.is_null()
            && type_of(&thrown_value, "object")
            && derive_from_js_value(&thrown_value, "then").is_function();
        unsafe {
            WORK_IN_PROGRESS_SUSPENDED_REASON = if is_wakeable {
                SUSPENDED_ON_DEPRECATED_THROW_PROMISE
            } else {
                SUSPENDED_ON_ERROR
            };
        };
    }

    unsafe {
        WORK_IN_PROGRESS_THROWN_VALUE = Some(thrown_value);
    }
}

Finally, the previous article mentioned an issue with bailout affecting the normal operation of Suspense. The solution was to move the code responsible for bubbling update priority to fiber_throw.rs:

let closure = Closure::wrap(Box::new(move || {
  ...
  mark_update_lane_from_fiber_to_root(source_fiber.clone(), lane.clone());
  ensure_root_is_scheduled(root.clone());
}) as Box<dyn Fn()>);
...

Additionally, in begin_work.rs, we exclude the Suspense component from the bailout logic:

if !has_scheduled_update_or_context
    && current.borrow().tag != WorkTag::SuspenseComponent
{
  ...
  return Ok(bailout_on_already_finished_work(
      work_in_progress,
      render_lane,
  ));
}

Please kindly give me a star!