Implement React v18 from Scratch Using WASM and Rust - [27] Implement useTransition

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

useTransition is a new hook introduced in React that allows you to update state without blocking the UI. The official React website provides an example that demonstrates the difference before and after using useTransition. There is also an article that analyzes the underlying principles. Now let's implement it. You can find the details of the changes in this link.

First, let's follow the process of adding a new hook and add the relevant code. Eventually, we will come to fiber_hooks.rs:

fn mount_transition() -> Vec<JsValue> {
    let result = mount_state(&JsValue::from(false)).unwrap();
    let is_pending = result[0].as_bool().unwrap();
    let set_pending = result[1].clone().dyn_into::<Function>().unwrap();
    let hook = mount_work_in_progress_hook();
    let set_pending_cloned = set_pending.clone();
    let closure = Closure::wrap(Box::new(move |callback: Function| {
        start_transition(set_pending_cloned.clone(), callback);
    }) as Box<dyn Fn(Function)>);
    let start: Function = closure.as_ref().unchecked_ref::<Function>().clone();
    closure.forget();
    hook.as_ref().unwrap().clone().borrow_mut().memoized_state =
        Some(MemoizedState::MemoizedJsValue(start.clone().into()));
    vec![JsValue::from_bool(is_pending), start.into()]
}

During the mount_transition process, the following data structure is formed:

So when update_transition is called, we can retrieve the values from the hooks:

fn update_transition() -> Vec<JsValue> {
    let result = update_state(&JsValue::undefined()).unwrap();
    let is_pending = result[0].as_bool().unwrap();
    let hook = update_work_in_progress_hook();
    if let MemoizedState::MemoizedJsValue(start) = hook
        .as_ref()
        .unwrap()
        .clone()
        .borrow()
        .memoized_state
        .as_ref()
        .unwrap()
    {
        return vec![JsValue::from_bool(is_pending), start.into()];
    }
    panic!("update_transition")
}

The key lies in the implementation of start_transition:

fn start_transition(set_pending: Function, callback: Function) {
    set_pending.call1(&JsValue::null(), &JsValue::from_bool(true));
    let prev_transition = unsafe { REACT_CURRENT_BATCH_CONFIG.transition };

    // low priority
    unsafe { REACT_CURRENT_BATCH_CONFIG.transition = Lane::TransitionLane.bits() };
    callback.call0(&JsValue::null());
    set_pending.call1(&JsValue::null(), &JsValue::from_bool(false));

    unsafe { REACT_CURRENT_BATCH_CONFIG.transition = prev_transition };
}

According to the analysis in this article, the implementation first updates isPending to true with the current priority. Then it lowers the priority, executes the callback, and updates isPending to false. Finally, it restores the previous priority.

The update process with lowered priority uses Concurrent Mode, which is why it doesn't block the UI:

if cur_priority == Lane::SyncLane {
  ...
} else {
    if is_dev() {
        log!("Schedule in macrotask, priority {:?}", update_lanes);
    }
    let scheduler_priority = lanes_to_scheduler_priority(cur_priority.clone());
    let closure = Closure::wrap(Box::new(move |did_timeout_js_value: JsValue| {
        let did_timeout = did_timeout_js_value.as_bool().unwrap();
        perform_concurrent_work_on_root(root_cloned.clone(), did_timeout)
    }) as Box<dyn Fn(JsValue) -> JsValue>);
    let function = closure.as_ref().unchecked_ref::<Function>().clone();
    closure.forget();
    new_callback_node = Some(unstable_schedule_callback_no_delay(
        scheduler_priority,
        function,
    ))
}

With this, the implementation of useTransition is mostly complete. However, there were a few bugs encountered during the process:

The first bug is in begin_work.rs:

work_in_progress.borrow_mut().lanes = Lane::NoLane;

When a FiberNode has multiple Lanes, this approach causes issues. It should be changed to:

work_in_progress.borrow_mut().lanes -= render_lane;

So that only the currently rendered Lane is removed each time.

The second bug is in work_loop.rs:

log!("render over {:?}", *root.clone().borrow());
WORK_IN_PROGRESS_ROOT_RENDER_LANE = Lane::NoLane;

Originally, this line was in the render_root function, resetting the variable after the Render phase is complete. But in Concurrent Mode, when the Render process is interrupted, this variable should not be reset. Therefore, this line is moved to perform_concurrent_work_on_root:

if exit_status == ROOT_COMPLETED {
    ...
    unsafe { WORK_IN_PROGRESS_ROOT_RENDER_LANE = Lane::NoLane };
}

So that the variable is only reset when the Render process is completed.

The third bug is in update_queue.rs, as shown in the following image:

Additionally, the Scheduler has been refactored. Previously, the min-heap was defined as follows:

static mut TASK_QUEUE: Vec<Task> = vec![];
static mut TIMER_QUEUE: Vec<Task> = vec![];

This required implementing a separate peek_mut function when modifying properties of the Task in the heap:

let mut task = peek_mut(&mut TASK_QUEUE);
task.callback = JsValue::null();

Now it has been changed to:

static mut TASK_QUEUE: Vec<Rc<RefCell<Task>>> = vec![];
static mut TIMER_QUEUE: Vec<Rc<RefCell<Task>>> = vec![];

And the peek function can be used uniformly:

let task = peek(&TASK_QUEUE);
task.borrow_mut().callback = JsValue::null();

Please kindly give me a star!