Implement React v18 from Scratch Using WASM and Rust - [24] Suspense(1) - Render Fallback

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

Suspense is undoubtedly one of the most appealing features in the new version of React, so let's implement it ourselves. This article is the first part, focusing on implementing the Fallback rendering of Suspense.

Consider the following code as an example:

import { Suspense } from 'react';

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

function Child() {
  throw new Promise((resolve) => setTimeout(resolve, 1000));
}

For the Suspense node, it has two child branches corresponding to Primary and Fallback. The root node of the Primary branch is of type Offscreen, and the root node of the Fallback branch is of type Fragment:

Specifically, in the example above:

During the initial render, the code enters the Primary branch. When it reaches the Child component, an exception is thrown because the component throws a Promise object. This triggers the "unwind" process, which searches for the nearest Suspense node and adds the DidCapture flag to it, and then continues the render process from that node.

Since the Suspense node has the DidCapture flag, the code enters the Fallback branch. The subsequent steps involve the normal render and commit processes, eventually rendering the content within the Fallback.

That's the functionality we want to implement. Now let's briefly go through the code.

First, let's take a look at begin_work.rs where we need to add handling for Suspense:

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

    let mut show_fallback = false;
    let did_suspend =
        (work_in_progress.borrow().flags.clone() & Flags::DidCapture) != Flags::NoFlags;

    if did_suspend {
        show_fallback = true;
        work_in_progress.borrow_mut().flags -= Flags::DidCapture;
    }

    let next_primary_children = derive_from_js_value(&next_props, "children");
    let next_fallback_children = derive_from_js_value(&next_props, "fallback");
    push_suspense_handler(work_in_progress.clone());

    if current.is_none() {
        if show_fallback {
            return Some(mount_suspense_fallback_children(
                work_in_progress.clone(),
                next_primary_children.clone(),
                next_fallback_children.clone(),
            ));
        } else {
            return Some(mount_suspense_primary_children(
                work_in_progress.clone(),
                next_primary_children.clone(),
            ));
        }
    } else {
        if show_fallback {
            return Some(update_suspense_fallback_children(
                work_in_progress.clone(),
                next_primary_children.clone(),
                next_fallback_children.clone(),
            ));
        } else {
            return Some(update_suspense_primary_children(
                work_in_progress.clone(),
                next_primary_children.clone(),
            ));
        }
    }
}

Here, we handle four branches based on whether the Fallback is shown and whether it's the first update.

Next, let's look at work_loop.rs:

loop {
  unsafe {
      if WORK_IN_PROGRESS_SUSPENDED_REASON != NOT_SUSPENDED && WORK_IN_PROGRESS.is_some() {
          // unwind process
          ...
      }
  }
  match if should_time_slice {
      work_loop_concurrent()
  } else {
      work_loop_sync()
  } {
      Ok(_) => {
          break;
      }
      Err(e) => handle_throw(root.clone(), e),
  };
}

When an exception is thrown in the component, it enters the Err branch, where we mainly add the handle_throw process, which is currently simple:

fn handle_throw(root: Rc<RefCell<FiberRootNode>>, thrown_value: JsValue) {
    unsafe {
        WORK_IN_PROGRESS_SUSPENDED_REASON = SUSPENDED_ON_DATA;
        WORK_IN_PROGRESS_THROWN_VALUE = Some(thrown_value);
    }
}

The loop continues, entering the unwind process:

loop {
  unsafe {
      if WORK_IN_PROGRESS_SUSPENDED_REASON != NOT_SUSPENDED && WORK_IN_PROGRESS.is_some() {
          let thrown_value = WORK_IN_PROGRESS_THROWN_VALUE.clone().unwrap();

          WORK_IN_PROGRESS_SUSPENDED_REASON = NOT_SUSPENDED;
          WORK_IN_PROGRESS_THROWN_VALUE = None;

          throw_and_unwind_work_loop(
              root.clone(),
              WORK_IN_PROGRESS.clone().unwrap(),
              thrown_value,
              lane.clone(),
          );
      }
  }
  ...
}

fn throw_and_unwind_work_loop(
    root: Rc<RefCell<FiberRootNode>>,
    unit_of_work: Rc<RefCell<FiberNode>>,
    thrown_value: JsValue,
    lane: Lane,
) {
    unwind_unit_of_work(unit_of_work);
}

The task here is to find the nearest Suspense node and mark it with DidCapture.

With this, our task is complete. However, to pave the way for the next article, let's implement a bit more functionality.

Using the same code example, when the initial render reaches the Child component, we should capture the Promise object it throws, call its then method, and trigger the re-rendering logic in the provided function.

This way, when the Promise object's state becomes fulfilled, it will re-enter the render process. At this point, if the Child component still throws an exception, the same process repeats. However, for now, we won't handle this because we haven't implemented the use hook yet. So, this is just a temporary approach for testing.

Let's see how we can capture the Promise object and trigger the re-rendering when it becomes fulfilled:

First, let's add the throw_exception method in throw_and_unwind_work_loop:

fn throw_and_unwind_work_loop(
    root: Rc<RefCell<FiberRootNode>>,
    unit_of_work: Rc<RefCell<FiberNode>>,
    thrown_value: JsValue,
    lane: Lane,
) {
    throw_exception(root.clone(), thrown_value, lane.clone());
}

fn attach_ping_listener(root: Rc<RefCell<FiberRootNode>>, wakeable: JsValue, lane: Lane) {
    let then_value = derive_from_js_value(&wakeable, "then");
    let then = then_value.dyn_ref::<Function>().unwrap();
    let closure = Closure::wrap(Box::new(move || {
        root.clone().borrow_mut().mark_root_updated(lane.clone());
        ensure_root_is_scheduled(root.clone());
    }) as Box<dyn Fn()>);
    let ping = closure.as_ref().unchecked_ref::<Function>().clone();
    closure.forget();
    then.call2(&wakeable, &ping, &ping)
        .expect("failed to call then function");
}

pub fn throw_exception(root: Rc<RefCell<FiberRootNode>>, value: JsValue, lane: Lane) {
    if !value.is_null()
        && type_of(&value, "object")
        && derive_from_js_value(&value, "then").is_function()
    {
        let suspense_boundary = get_suspense_handler();
        if suspense_boundary.is_some() {
            let suspense_boundary = suspense_boundary.unwrap();
            suspense_boundary.borrow_mut().flags |= Flags::ShouldCapture;
        }

        attach_ping_listener(root, value, lane)
    }
}

The ping function is the function passed to then. The core logic is to set the current lane as the priority for the next update and call ensure_root_is_scheduled to start a new update. However, it was found during testing that this is not enough because the performance optimization in begin_work.rs will bail out the update starting from the root node. This issue also exists in big-react (switch to the master branch and run the "suspense-use" example to reproduce it, see the issue for details).

To solve this problem, a temporary solution is to bubble up the update priority one more time before the unwind process. This way, when the update starts again from the root node, it won't enter the bailout process due to the flags on subtree_flags.

loop {
    unsafe {
        if WORK_IN_PROGRESS_SUSPENDED_REASON != NOT_SUSPENDED && WORK_IN_PROGRESS.is_some() {
            let thrown_value = WORK_IN_PROGRESS_THROWN_VALUE.clone().unwrap();

            WORK_IN_PROGRESS_SUSPENDED_REASON = NOT_SUSPENDED;
            WORK_IN_PROGRESS_THROWN_VALUE = None;

            mark_update_lane_from_fiber_to_root(
                WORK_IN_PROGRESS.clone().unwrap(),
                lane.clone(),
            );

            ...
        }
    }
    ...
}

You can find the details of this update here.

Please kindly give me a star!

Blog

Implement React v18 from Scratch Using WASM and Rust - [24] Suspense(1) - Render Fallback

ayou

Join Our Newsletter. No Spam, Only the good stuff.

Related