Blog 22 - Feat memo

examples/hello-world/src/App.tsx

@@ -1 +1 @@
-export {default} from './useContext/index2'
+export {default} from './memo'

examples/hello-world/src/memo/index.tsx

@@ -0,0 +1,27 @@
+import {useState, memo} from 'react'
+
+export default function App() {
+  const [num, update] = useState(0)
+  console.log('App render ', num)
+  return (
+    <div onClick={() => update(num + 1)}>
+      <Cpn num={num} name={'cpn1'} />
+      <Cpn num={0} name={'cpn2'} />
+    </div>
+  )
+}
+
+const Cpn = memo(function ({num, name}) {
+  console.log('render ', name)
+  return (
+    <div>
+      {name}: {num}
+      <Child />
+    </div>
+  )
+})
+
+function Child() {
+  console.log('Child render')
+  return <p>i am child</p>
+}

packages/react-reconciler/src/begin_work.rs

@@ -1,10 +1,10 @@
 use std::cell::RefCell;
 use std::rc::Rc;
 
-use wasm_bindgen::JsValue;
+use wasm_bindgen::{JsCast, JsValue};
 
-use shared::{derive_from_js_value, is_dev, log};
-use web_sys::js_sys::Object;
+use shared::{derive_from_js_value, is_dev, log, shallow_equal};
+use web_sys::js_sys::{Function, Object};
 
 use crate::child_fiber::{clone_child_fiblers, mount_child_fibers, reconcile_child_fibers};
 use crate::fiber::{FiberNode, MemoizedState};
@@ -118,11 +118,51 @@ pub fn begin_work(
         WorkTag::HostText => Ok(None),
         WorkTag::ContextProvider => Ok(update_context_provider(
             work_in_progress.clone(),
-            render_lane,
+            render_lane.clone(),
         )),
+        WorkTag::MemoComponent => update_memo_component(work_in_progress.clone(), render_lane),
     };
 }
 
+fn update_memo_component(
+    work_in_progress: Rc<RefCell<FiberNode>>,
+    render_lane: Lane,
+) -> Result<Option<Rc<RefCell<FiberNode>>>, JsValue> {
+    let current = { work_in_progress.borrow().alternate.clone() };
+    let next_props = { work_in_progress.borrow().pending_props.clone() };
+
+    if current.is_some() {
+        let current = current.unwrap();
+        let prev_props = current.borrow().memoized_props.clone();
+        if !check_scheduled_update_or_context(current.clone(), render_lane.clone()) {
+            let mut props_equal = false;
+            let compare = derive_from_js_value(&work_in_progress.borrow()._type, "compare");
+            if compare.is_function() {
+                let f = compare.dyn_ref::<Function>().unwrap();
+                props_equal = f
+                    .call2(&JsValue::null(), &prev_props, &next_props)
+                    .unwrap()
+                    .as_bool()
+                    .unwrap();
+            } else {
+                props_equal = shallow_equal(&prev_props, &next_props);
+            }
+
+            if props_equal && Object::is(&current.borrow()._ref, &work_in_progress.borrow()._ref) {
+                unsafe { DID_RECEIVE_UPDATE = false };
+                work_in_progress.borrow_mut().pending_props = prev_props;
+                work_in_progress.borrow_mut().lanes = current.borrow().lanes.clone();
+                return Ok(bailout_on_already_finished_work(
+                    work_in_progress.clone(),
+                    render_lane,
+                ));
+            }
+        }
+    }
+    let Component = { derive_from_js_value(&work_in_progress.borrow()._type, "type") };
+    update_function_component(work_in_progress.clone(), Component, render_lane)
+}
+
 fn update_context_provider(
     work_in_progress: Rc<RefCell<FiberNode>>,
     render_lane: Lane,

packages/react-reconciler/src/commit_work.rs

@@ -12,9 +12,7 @@ use crate::fiber::{FiberNode, FiberRootNode, StateNode};
 use crate::fiber_flags::{get_mutation_mask, get_passive_mask, Flags};
 use crate::fiber_hooks::Effect;
 use crate::work_tags::WorkTag;
-use crate::work_tags::WorkTag::{
-    ContextProvider, FunctionComponent, HostComponent, HostRoot, HostText,
-};
+use crate::work_tags::WorkTag::{FunctionComponent, HostComponent, HostRoot, HostText};
 use crate::HOST_CONFIG;
 
 static mut NEXT_EFFECT: Option<Rc<RefCell<FiberNode>>> = None;
@@ -314,7 +312,6 @@ fn commit_deletion(child_to_delete: Rc<RefCell<FiberNode>>, root: Rc<RefCell<Fib
             FunctionComponent => {
                 commit_passive_effect(unmount_fiber.clone(), root.clone(), "unmount");
             }
-            HostRoot => {}
             HostComponent => {
                 if cloned.borrow().is_none() {
                     *cloned.borrow_mut() = Some(unmount_fiber.clone());
@@ -325,10 +322,7 @@ fn commit_deletion(child_to_delete: Rc<RefCell<FiberNode>>, root: Rc<RefCell<Fib
                     *cloned.borrow_mut() = Some(unmount_fiber.clone());
                 }
             }
-            HostRoot => todo!(),
-            HostComponent => todo!(),
-            HostText => todo!(),
-            ContextProvider => todo!(),
+            _ => todo!(),
         };
     });
 

packages/react-reconciler/src/complete_work.rs

@@ -5,7 +5,7 @@ use std::rc::Rc;
 use wasm_bindgen::JsValue;
 use web_sys::js_sys::{Object, Reflect};
 
-use shared::{derive_from_js_value, log};
+use shared::derive_from_js_value;
 
 use crate::fiber::{FiberNode, StateNode};
 use crate::fiber_context::pop_provider;
@@ -141,16 +141,6 @@ impl CompleteWork {
         let current = { work_in_progress_cloned.borrow().alternate.clone() };
         let tag = { work_in_progress_cloned.borrow().tag.clone() };
         match tag {
-            WorkTag::FunctionComponent => {
-                self.bubble_properties(work_in_progress.clone());
-                // log!("bubble_properties function {:?}", work_in_progress.clone());
-                None
-            }
-            WorkTag::HostRoot => {
-                self.bubble_properties(work_in_progress.clone());
-                // log!("bubble_properties HostRoot {:?}", work_in_progress.clone());
-                None
-            }
             WorkTag::HostComponent => {
                 if current.is_some() && work_in_progress_cloned.borrow().state_node.is_some() {
                     // todo compare
@@ -222,6 +212,10 @@ impl CompleteWork {
                 self.bubble_properties(work_in_progress.clone());
                 None
             }
+            _ => {
+                self.bubble_properties(work_in_progress.clone());
+                None
+            }
         }
     }
 }

packages/react-reconciler/src/fiber.rs

@@ -9,7 +9,7 @@ use scheduler::Task;
 use wasm_bindgen::JsValue;
 use web_sys::js_sys::Reflect;
 
-use shared::{derive_from_js_value, log, type_of, REACT_PROVIDER_TYPE};
+use shared::{derive_from_js_value, log, type_of, REACT_MEMO_TYPE, REACT_PROVIDER_TYPE};
 
 use crate::fiber_context::ContextItem;
 use crate::fiber_flags::Flags;
@@ -128,6 +128,17 @@ impl Debug for FiberNode {
                 )
                 .expect("print error");
             }
+            WorkTag::MemoComponent => {
+                write!(
+                    f,
+                    "{:?}(flags:{:?}, subtreeFlags:{:?}, lanes:{:?})",
+                    self._type.as_ref(),
+                    self.flags,
+                    self.subtree_flags,
+                    self.lanes
+                )
+                .expect("print error");
+            }
         })
     }
 }
@@ -167,12 +178,17 @@ impl FiberNode {
         let mut fiber_tag = WorkTag::FunctionComponent;
         if _type.is_string() {
             fiber_tag = WorkTag::HostComponent
-        } else if type_of(&_type, "object")
-            && derive_from_js_value(&_type, "$$typeof") == REACT_PROVIDER_TYPE
-        {
-            fiber_tag = WorkTag::ContextProvider;
+        } else if type_of(&_type, "object") {
+            let _typeof = derive_from_js_value(&_type, "$$typeof");
+            if _typeof == REACT_PROVIDER_TYPE {
+                fiber_tag = WorkTag::ContextProvider;
+            } else if _typeof == REACT_MEMO_TYPE {
+                fiber_tag = WorkTag::MemoComponent;
+            } else {
+                log!("Unsupported type {:?}", _type);
+            }
         } else if !type_of(&_type, "function") {
-            log!("Unsupported type {:?}", ele);
+            log!("Unsupported type {:?}", _type);
         }
 
         let mut fiber = FiberNode::new(fiber_tag, props, key, _ref);

packages/react-reconciler/src/work_tags.rs

@@ -5,4 +5,5 @@ pub enum WorkTag {
     HostComponent = 5,
     HostText = 6,
     ContextProvider = 8,
+    MemoComponent = 15,
 }

packages/react/src/lib.rs

@@ -1,7 +1,10 @@
 use js_sys::{Array, Object, Reflect, JSON};
 use wasm_bindgen::prelude::*;
 
-use shared::{derive_from_js_value, REACT_CONTEXT_TYPE, REACT_ELEMENT_TYPE, REACT_PROVIDER_TYPE};
+use shared::{
+    derive_from_js_value, REACT_CONTEXT_TYPE, REACT_ELEMENT_TYPE, REACT_MEMO_TYPE,
+    REACT_PROVIDER_TYPE,
+};
 
 use crate::current_dispatcher::CURRENT_DISPATCHER;
 
@@ -166,3 +169,27 @@ pub unsafe fn create_context(default_value: &JsValue) -> JsValue {
     Reflect::set(&context, &"Provider".into(), &provider);
     context.into()
 }
+
+#[wasm_bindgen]
+pub unsafe fn memo(_type: &JsValue, compare: &JsValue) -> JsValue {
+    let fiber_type = Object::new();
+
+    Reflect::set(
+        &fiber_type,
+        &"$$typeof".into(),
+        &JsValue::from_str(REACT_MEMO_TYPE),
+    );
+    Reflect::set(&fiber_type, &"type".into(), _type);
+
+    let null = JsValue::null();
+    Reflect::set(
+        &fiber_type,
+        &"compare".into(),
+        if compare.is_undefined() {
+            &null
+        } else {
+            compare
+        },
+    );
+    fiber_type.into()
+}

packages/shared/src/lib.rs

@@ -1,10 +1,11 @@
-use web_sys::js_sys::Reflect;
 use web_sys::js_sys::JSON::stringify;
-use web_sys::wasm_bindgen::JsValue;
+use web_sys::js_sys::{Object, Reflect};
+use web_sys::wasm_bindgen::{JsCast, JsValue};
 
 pub static REACT_ELEMENT_TYPE: &str = "react.element";
 pub static REACT_CONTEXT_TYPE: &str = "react.context";
 pub static REACT_PROVIDER_TYPE: &str = "react.provider";
+pub static REACT_MEMO_TYPE: &str = "react.memo";
 
 #[macro_export]
 macro_rules! log {
@@ -67,3 +68,35 @@ pub fn to_string(js_value: &JsValue) -> String {
         }
     })
 }
+
+pub fn shallow_equal(a: &JsValue, b: &JsValue) -> bool {
+    if Object::is(a, b) {
+        return true;
+    }
+
+    if !type_of(a, "object") || a.is_null() || !type_of(b, "object") || b.is_null() {
+        return false;
+    }
+
+    let a = a.dyn_ref::<Object>().unwrap();
+    let b = b.dyn_ref::<Object>().unwrap();
+    let keys_a = Object::keys(a);
+    let keys_b = Object::keys(b);
+
+    if keys_a.length() != keys_b.length() {
+        return false;
+    }
+
+    for key in keys_a {
+        if !Object::has_own_property(&b, &key)
+            || !Object::is(
+                &Reflect::get(&a, &key).unwrap(),
+                &Reflect::get(&b, &key).unwrap(),
+            )
+        {
+            return false;
+        }
+    }
+
+    true
+}

readme.md

@@ -1,7 +1,7 @@
 # How to start?
 
 - install wasm-pack
-- `npm run build`
+- `npm run build:dev`
 - cd `examples/hello-world`
 - `pnpm install`
 - `npm run dev`
@@ -43,3 +43,9 @@
 [从零实现 React v18，但 WASM 版 - [17] 实现 Concurrent 模式](https://www.paradeto.com/2024/06/19/big-react-wasm-17/)
 
 [从零实现 React v18，但 WASM 版 - [18] 实现 useRef, useCallback, useMemo](https://www.paradeto.com/2024/07/10/big-react-wasm-18/)
+
+[从零实现 React v18，但 WASM 版 - [19] 性能优化之 bailout 和 eagerState](https://www.paradeto.com/2024/07/19/big-react-wasm-19/)
+
+[从零实现 React v18，但 WASM 版 - [20] 实现 Context](https://www.paradeto.com/2024/07/26/big-react-wasm-20/)
+
+[从零实现 React v18，但 WASM 版 - [21] 性能优化支持 Context](https://www.paradeto.com/2024/07/26/big-react-wasm-21/)