Implement React v18 from Scratch Using WASM and Rust – [18] Implement useRef, useCallback, useMemo

By coursity / July 10, 2024

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)
}

Enter fullscreen mode
Exit fullscreen mode

// 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, )))
}

Enter fullscreen mode
Exit fullscreen mode

Next, let’s take a look at how we need to modify react-reconciler.

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"), }
}

Enter fullscreen mode
Exit fullscreen mode

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(&current.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()); ...
}

Enter fullscreen mode
Exit fullscreen mode

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());

Enter fullscreen mode
Exit fullscreen mode

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;
}

Enter fullscreen mode
Exit fullscreen mode

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); } }
}

Enter fullscreen mode
Exit fullscreen mode

By now, the implementation of useRef is complete. Let’s move on to the other two hooks.

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()
}

Enter fullscreen mode
Exit fullscreen mode

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");
}

Enter fullscreen mode
Exit fullscreen mode

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!

Related

Discover more from Coursity

Subscribe to get the latest posts sent to your email.

Leave a Comment Cancel Reply