![Cover image for Implement React v18 from Scratch Using WASM and Rust - [18] Implement useRef, useCallback, useMemo](https://media.dev.to/cdn-cgi/image/width=1000,height=420,fit=cover,gravity=auto,format=auto/https%3A%2F%2Fdev-to-uploads.s3.amazonaws.com%2Fuploads%2Farticles%2Fy4nf1pli06yvz690kyeg.png)
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：v18
We have already implemented two commonly used hooks, useState and useEffect, earlier. Today, we will continue to implement three more hooks: useRef, useCallback, and useMemo.
Since the framework has already been set up, we can simply follow the same pattern and add these three hooks to our react package.
// react/src/lib.rs
#[wasm_bindgen(js_name = useRef)]
pub unsafe fn use_ref(initial_value: &JsValue) -> JsValue {
let use_ref = &CURRENT_DISPATCHER.current.as_ref().unwrap().use_ref;
use_ref.call1(&JsValue::null(), initial_value)
}
#[wasm_bindgen(js_name = useMemo)]
pub unsafe fn use_memo(create: &JsValue, deps: &JsValue) -> Result<JsValue, JsValue> {
let use_memo = &CURRENT_DISPATCHER.current.as_ref().unwrap().use_memo;
use_memo.call2(&JsValue::null(), create, deps)
}
#[wasm_bindgen(js_name = useCallback)]
pub unsafe fn use_callback(callback: &JsValue, deps: &JsValue) -> JsValue {
let use_callback = &CURRENT_DISPATCHER.current.as_ref().unwrap().use_callback;
use_callback.call2(&JsValue::null(), callback, deps)
}
// react/src/current_dispatcher.rs
pub unsafe fn update_dispatcher(args: &JsValue) {
...
let use_ref = derive_function_from_js_value(args, "use_ref");
let use_memo = derive_function_from_js_value(args, "use_memo");
let use_callback = derive_function_from_js_value(args, "use_callback");
CURRENT_DISPATCHER.current = Some(Box::new(Dispatcher::new(
use_state,
use_effect,
use_ref,
use_memo,
use_callback,
)))
}
Next, let's take a look at how we need to modify react-reconciler.
useRef
First, we need to add mount_ref and update_ref in fiber_hooks.rs.
fn mount_ref(initial_value: &JsValue) -> JsValue {
let hook = mount_work_in_progress_hook();
let ref_obj: Object = Object::new();
Reflect::set(&ref_obj, &"current".into(), initial_value);
hook.as_ref().unwrap().borrow_mut().memoized_state =
Some(MemoizedState::MemoizedJsValue(ref_obj.clone().into()));
ref_obj.into()
}
fn update_ref(initial_value: &JsValue) -> JsValue {
let hook = update_work_in_progress_hook();
match hook.unwrap().borrow_mut().memoized_state.clone() {
Some(MemoizedState::MemoizedJsValue(value)) => value,
_ => panic!("ref is none"),
}
}
For useRef, these two methods can be implemented very simply.
Next, following the order of the rendering process, we need to modify begin_work.rs first. Here, we will only handle FiberNode of the Host Component type for now.
fn mark_ref(current: Option<Rc<RefCell<FiberNode>>>, work_in_progress: Rc<RefCell<FiberNode>>) {
let _ref = { work_in_progress.borrow()._ref.clone() };
if (current.is_none() && !_ref.is_null())
|| (current.is_some() && Object::is(¤t.as_ref().unwrap().borrow()._ref, &_ref))
{
work_in_progress.borrow_mut().flags |= Flags::Ref;
}
}
fn update_host_component(
work_in_progress: Rc<RefCell<FiberNode>>,
) -> Option<Rc<RefCell<FiberNode>>> {
...
let alternate = { work_in_progress.borrow().alternate.clone() };
mark_ref(alternate, work_in_progress.clone());
...
}
The handling process is also straightforward. We can mark the FiberNode with a Ref flag based on certain conditions, which will be processed during the commit phase.
Next, we need to add the "layout phase" in the commit_root method in work_loop.rs.
// 1/3: Before Mutation
// 2/3: Mutation
commit_mutation_effects(finished_work.clone(), root.clone());
// Switch Fiber Tree
cloned.borrow_mut().current = finished_work.clone();
// 3/3: Layout
commit_layout_effects(finished_work.clone(), root.clone());
This phase occurs after commit_mutation_effects, which means it happens after modifying the DOM. So we can update the Ref here.
In commit_layout_effects, we can decide whether to update the Ref based on whether the FiberNode contains the Ref flag. We can do this by calling the safely_attach_ref method.
if flags & Flags::Ref != Flags::NoFlags && tag == HostComponent {
safely_attach_ref(finished_work.clone());
finished_work.borrow_mut().flags -= Flags::Ref;
}
In safely_attach_ref, we first retrieve the state_node property from the FiberNode. This property points to the actual node corresponding to the FiberNode. For React DOM, it would be the DOM node.
Next, we handle different cases based on the type of the _ref value.
fn safely_attach_ref(fiber: Rc<RefCell<FiberNode>>) {
let _ref = fiber.borrow()._ref.clone();
if !_ref.is_null() {
let instance = match fiber.borrow().state_node.clone() {
Some(s) => match &*s {
StateNode::Element(element) => {
let node = (*element).downcast_ref::<Node>().unwrap();
Some(node.clone())
}
StateNode::FiberRootNode(_) => None,
},
None => None,
};
if instance.is_none() {
panic!("instance is none")
}
let instance = instance.as_ref().unwrap();
if type_of(&_ref, "function") {
// <div ref={() => {...}} />
_ref.dyn_ref::<Function>()
.unwrap()
.call1(&JsValue::null(), instance);
} else {
// const ref = useRef()
// <div ref={ref} />
Reflect::set(&_ref, &"current".into(), instance);
}
}
}
By now, the implementation of useRef is complete. Let's move on to the other two hooks.
useCallback and useMemo
The implementation of these two hooks becomes simpler. You just need to modify fiber_hooks, and both of them have very similar implementation approaches. Taking useCallback as an example, during the initial render, you only need to save the two arguments passed to useCallback on the Hook node and then return the first argument.
fn mount_callback(callback: Function, deps: JsValue) -> JsValue {
let hook = mount_work_in_progress_hook();
let next_deps = if deps.is_undefined() {
JsValue::null()
} else {
deps
};
let array = Array::new();
array.push(&callback);
array.push(&next_deps);
hook.as_ref().unwrap().clone().borrow_mut().memoized_state =
Some(MemoizedState::MemoizedJsValue(array.into()));
callback.into()
}
When updating, you first retrieve the previously saved second argument and compare it item by item with the new second argument that is passed in. If they are all the same, you return the previously saved first argument. Otherwise, you return the new first argument that was passed in.
fn update_callback(callback: Function, deps: JsValue) -> JsValue {
let hook = update_work_in_progress_hook();
let next_deps = if deps.is_undefined() {
JsValue::null()
} else {
deps
};
if let MemoizedState::MemoizedJsValue(prev_state) = hook
.clone()
.unwrap()
.borrow()
.memoized_state
.as_ref()
.unwrap()
{
if !next_deps.is_null() {
let arr = prev_state.dyn_ref::<Array>().unwrap();
let prev_deps = arr.get(1);
if are_hook_inputs_equal(&next_deps, &prev_deps) {
return arr.get(0);
}
}
let array = Array::new();
array.push(&callback);
array.push(&next_deps);
hook.as_ref().unwrap().clone().borrow_mut().memoized_state =
Some(MemoizedState::MemoizedJsValue(array.into()));
return callback.into();
}
panic!("update_callback, memoized_state is not JsValue");
}
For useMemo, it simply adds an extra step of executing the function, but the other steps remain the same.
With this, the implementation of these two hooks is complete. However, currently, these two hooks don't provide any performance optimization features because we haven't implemented them yet. Let's leave that for the next article.
For the details of this update, please refer to here. Please kindly give me a star!
Top comments (0)