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!