Implement React v18 from Scratch Using WASM and Rust - [10] Implement Update for Single Node.

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

The previous article mentioned that we haven't fully implemented the update process yet. So, in this article, we will implement it.

Let's continue using the previous example:

function App() {
  const [name, setName] = useState(() => 'ayou')
  setTimeout(() => {
    setName('ayouayou')
  }, 1000)
  return (
    <div>
      <Comp>{name}</Comp>
    </div>
  )
}

When we call setName('ayouayou'), it triggers the update process. The setName method is returned in the mount_state during the initial render. This method attaches a Hook node to the memoized_state of the current FiberNode. If there are multiple hooks, they form a linked list. The Hook node has an update_queue, which is clearly an update queue. It also has a memoized_state property that records the current state of the Hook.

fn mount_state(initial_state: &JsValue) -> Result<Vec<JsValue>, JsValue> {
  // Add hook to current FiberNode memoized_state
  let hook = mount_work_in_progress_hook();
  let memoized_state: JsValue;

  if initial_state.is_function() {
    memoized_state = initial_state
        .dyn_ref::<Function>()
        .unwrap()
        .call0(&JsValue::null())?;
  } else {
      memoized_state = initial_state.clone();
  }

  hook.as_ref().unwrap().clone().borrow_mut().memoized_state =
      Some(MemoizedState::JsValue(memoized_state.clone()));
  let queue = create_update_queue();
  hook.as_ref().unwrap().clone().borrow_mut().update_queue = Some(queue.clone());
  ...
}

mount_state ultimately returns initial_state and a function:

let q_rc = Rc::new(queue.clone());
let q_rc_cloned = q_rc.clone();
let fiber = unsafe {
    CURRENTLY_RENDERING_FIBER.clone().unwrap()
};
let closure = Closure::wrap(Box::new(move |action: &JsValue| unsafe {
    dispatch_set_state(
        fiber.clone(),
        (*q_rc_cloned).clone(),
        action,
    )
}) as Box<dyn Fn(&JsValue)>);
let function = closure.as_ref().unchecked_ref::<Function>().clone();
closure.forget();

queue.clone().borrow_mut().dispatch = Some(function.clone());

Ok(vec![memoized_state, function.into()])

It's a bit strange here with q_rc_cloned in the closure. queue is already of type Rc, so why is there an additional layer of Rc on the outside? This is because if we change (*q_rc_cloned).clone() to queue.clone(), it will result in the following error:

error[E0382]: borrow of moved value: `queue`
   --> packages/react-reconciler/src/fiber_hooks.rs:251:5
    |
233 |     let queue = create_update_queue();
    |         ----- move occurs because `queue` has type `Rc<RefCell<UpdateQueue>>`, which does not implement the `Copy` trait
...
240 |     let closure = Closure::wrap(Box::new(move |action: &JsValue| unsafe {
    |                                          ----------------------- value moved into closure here
...
243 |             queue.clone(),
    |             ----- variable moved due to use in closure
...
251 |     queue.clone().borrow_mut().dispatch = Some(function.clone());
    |     ^^^^^ value borrowed here after move

The reason is that the ownership of the value of queue has already been moved into the closure, so it can no longer be used outside. Can we remove the move? Let's try, and we find that it results in this error:

error[E0597]: `queue` does not live long enough
   --> packages/react-reconciler/src/fiber_hooks.rs:243:13
    |
240 |       let closure = Closure::wrap(Box::new(|action: &JsValue| unsafe {
    |                                   -        ------------------ value captured here
    |  _________________________________|
    | |
241 | |         dispatch_set_state(
242 | |             fiber.clone(),
243 | |             queue.clone(),
    | |             ^^^^^ borrowed value does not live long enough
...   |
246 | |         )
247 | |     }) as Box<dyn Fn(&JsValue)>);
    | |______- cast requires that `queue` is borrowed for `'static`
...
254 |   }
    |   - `queue` dropped here while still borrowed

The reason is that if we don't move it in, queue will be deallocated after mount_state is executed, but it is still borrowed inside the closure, which is obviously not allowed.

It is often said that the steep learning curve of Rust lies in the fact that you are constantly fighting with the compiler. However, this is the philosophy of Rust: to discover most issues during compilation, which leads to a much higher efficiency in fixing them compared to discovering and fixing them after deployment. Moreover, the Rust compiler is quite intelligent and provides clear problem descriptions.

Let's get back to the error of using move and queue. Analyzing the situation, since queue has been moved, we can't use queue afterwards. So, if we move some other value, wouldn't that work? That's why we have queue_rc, and the memory models of the two are compared as shown below:

Another point worth mentioning is that we attach this closure function to the dispatch property of the queue of each Hook node:

queue.clone().borrow_mut().dispatch = Some(function.clone());

This is done to return the same function during update_state:

fn update_state(initial_state: &JsValue) -> Result<Vec<JsValue>, JsValue> {
   ...
    Ok(vec![
        hook.clone().unwrap().clone()
            .borrow()
            .memoized_state
            .clone()
            .unwrap()
            .js_value()
            .unwrap().clone(),
        queue.clone().unwrap().borrow().dispatch.clone().into(),
    ])
}

However, I feel that having dispatch as an attribute of Hook is more appropriate. At least for now, it doesn't seem to have any direct association with queue.

Returning to the code, when dispatch is called, it eventually invokes dispatch_set_state:

fn dispatch_set_state(
    fiber: Rc<RefCell<FiberNode>>,
    update_queue: Rc<RefCell<UpdateQueue>>,
    action: &JsValue,
) {
    let update = create_update(action.clone());
    enqueue_update(update_queue.clone(), update);
    unsafe {
        WORK_LOOP
            .as_ref()
            .unwrap()
            .clone()
            .borrow()
            .schedule_update_on_fiber(fiber.clone());
    }
}

Its purpose is to update the update_queue of the Hook node with the provided action and initiate a new round of update process. At this point, the state of the App node looks as shown in the following diagram:

Next, the process is similar to the initial rendering. First, let's look at the "begin work" phase. During the update process, the "begin work" phase primarily handles the child nodes of the FiberNode. It generates new child FiberNode by comparing the existing child FiberNode in the Fiber Tree with the newly generated ReactElement (referred to as children in the code). This is commonly known as the diffing process:

The way children are generated differs based on the type of FiberNode:

• HostRoot: Values are taken from memoized_state.
• HostComponent: Values are taken from pending_props.
• FunctionComponent: Obtained by executing the Function pointed to by the type.
• HostText: This process is not applicable and can be ignored.

There are two scenarios for generating these new child FiberNode:

• When the key and type of the diffing ReactElement and FiberNode are the same. The FiberNode is reused, and the pending_props of the FiberNode are updated with the props from the ReactElement:

• In other cases, a new FiberNode is created, and the parent node is marked with the ChildDeletion flag. The old FiberNode is added to the deletions list:

I won't provide the code here, but you can refer to the child_fiber file in this commit.

Since generating children for FunctionComponent is a bit more complex, let's go back and look at the changes made in the render_with_hooks method. The main changes are:

pub fn render_with_hooks(work_in_progress: Rc<RefCell<FiberNode>>) -> Result<JsValue, JsValue> {
  ...
  if current.is_some() {
      // log!("还未实现update时renderWithHooks");
      update_hooks_to_dispatcher(true);
  } else {
      update_hooks_to_dispatcher(false);
  }
  ...
}

fn update_hooks_to_dispatcher(is_update: bool) {
    let object = Object::new();

    let closure = Closure::wrap(Box::new(if is_update { update_state } else { mount_state })
        as Box<dyn Fn(&JsValue) -> Result<Vec<JsValue>, JsValue>>);
    let function = closure.as_ref().unchecked_ref::<Function>().clone();
    closure.forget();
    Reflect::set(&object, &"use_state".into(), &function).expect("TODO: panic set use_state");

    updateDispatcher(&object.into());
}

During the update, the use_state in the dispatcher is replaced with the update_state method. The update_state method primarily calculates the new memoized_state based on the update_queue and memoized_state of the Hooks and returns it. It also returns the dispatch function.

fn update_state(initial_state: &JsValue) -> Result<Vec<JsValue>, JsValue> {
    let hook = update_work_in_progress_hook();

    let hook_cloned = hook.clone().unwrap().clone();
    let queue = hook_cloned.borrow().update_queue.clone();
    let base_state = hook_cloned.borrow().memoized_state.clone();

    unsafe {
        hook_cloned.borrow_mut().memoized_state = process_update_queue(
            base_state,
            queue.clone(),
            CURRENTLY_RENDERING_FIBER.clone().unwrap(),
        );
    }

    Ok(vec![
        hook.clone().unwrap().clone()
            .borrow()
            .memoized_state
            .clone()
            .unwrap()
            .js_value()
            .unwrap().clone(),
        queue.clone().unwrap().borrow().dispatch.clone().into(),
    ])
}

That's all for the "begin work" phase. Next, let's take a look at the "complete work" phase, which is relatively simpler. In this phase, nodes are marked with the Update flag, and the logic for handling HostText and HostComponent is modified.

WorkTag::HostText => {
  if current.is_some() && work_in_progress_cloned.borrow().state_node.is_some() {
      // reuse FiberNode
      let old_text = derive_from_js_value(&current.clone().unwrap().clone().borrow().memoized_props, "content");
      let new_test = derive_from_js_value(&new_props, "content");
      if !Object::is(&old_text, &new_test) {
          CompleteWork::mark_update(work_in_progress.clone());
      }
  } else {
      let text_instance = self.host_config.create_text_instance(
          Reflect::get(&new_props, &JsValue::from_str("content"))
              .unwrap()
              .as_string()
              .unwrap(),
      );
      work_in_progress.clone().borrow_mut().state_node =
          Some(Rc::new(StateNode::Element(text_instance.clone())));
  }

  self.bubble_properties(work_in_progress.clone());
  None
}，
WorkTag::HostComponent => {
  if current.is_some() && work_in_progress_cloned.borrow().state_node.is_some() {
    // reuse FiberNode
    log!("TODO: update properties")
  } else {
      let instance = self.host_config.create_instance(
          work_in_progress
              .clone()
              .borrow()
              ._type
              .as_ref()
              .as_string()
              .unwrap(),
      );
      self.append_all_children(instance.clone(), work_in_progress.clone());
      work_in_progress.clone().borrow_mut().state_node =
          Some(Rc::new(StateNode::Element(instance.clone())));
  }

  self.bubble_properties(work_in_progress.clone());
  None
}

Finally, we have the "commit" phase, which mainly involves adding handling for Update and ChildDeletion in the commit_mutation_effects_on_fiber function.

fn commit_mutation_effects_on_fiber(&self, finished_work: Rc<RefCell<FiberNode>>) {
  ...
  if flags.contains(Flags::ChildDeletion) {
      let deletions = finished_work.clone().borrow().deletions.clone();
      if deletions.is_some() {
          let deletions = deletions.unwrap();
          for child_to_delete in deletions {
              self.commit_deletion(child_to_delete);
          }
      }
      finished_work.clone().borrow_mut().flags -= Flags::ChildDeletion;
  }

  if flags.contains(Flags::Update) {
      self.commit_update(finished_work.clone());
      finished_work.clone().borrow_mut().flags -= Flags::Update;
  }
  ...

In the Update part, only HostText is currently handled, which is relatively simple, so we won't go into detail. Let's directly look at the code. Here, I'll focus on explaining ChildDeletion.

In the "begin work" phase, we mentioned that child nodes marked for deletion are added to the deletions list of their parent node. So here, we iterate over this list and call commit_deletion. This function traverses the subtree rooted at child_to_delete in a pre-order manner (prioritizing the traversal of the root node). It executes the relevant side effects on these nodes, such as invoking the componentWillUnmount method or the destroy method returned by useEffect. From this, we can observe that the side effects of the parent component are executed first.

For example, consider the following example:

The traversal order is div -> p -> i -> span. Additionally, the first node encountered during traversal is recorded, which in this case is div. The deletion operation is then performed on this node.

Alright, the single-node update process is now complete. In summary:

• In the "begin work" phase, mark child nodes for deletion or insertion.
• In the "complete work" phase, mark nodes for update.
• In the commit phase, perform a depth-first traversal of the Fiber Tree, processing the marked nodes. For nodes marked as ChildDeletion, a pre-order traversal is performed on the subtree rooted at that node.

For more details, please refer to this update.

Please kindly give me a star!