From 025473353fbc217a9c979c4d07df482705a0b396 Mon Sep 17 00:00:00 2001
From: Hiroshi Ogawa <hi.ogawa.zz@gmail.com>
Date: Thu, 19 Oct 2023 14:46:18 +0900
Subject: [PATCH] fix: fix memo effect

---
 packages/tiny-react/src/compat/index.ts | 63 ++++++++++++++++++-------
 packages/tiny-react/src/index.test.ts   | 34 -------------
 2 files changed, 45 insertions(+), 52 deletions(-)

diff --git a/packages/tiny-react/src/compat/index.ts b/packages/tiny-react/src/compat/index.ts
index e71e4b7b..e5293e95 100644
--- a/packages/tiny-react/src/compat/index.ts
+++ b/packages/tiny-react/src/compat/index.ts
@@ -6,6 +6,7 @@ import {
   EMPTY_NODE,
   type FC,
   NODE_TYPE_CUSTOM,
+  type VCustom,
   type VNode,
 } from "../virtual-dom";
 
@@ -47,30 +48,56 @@ export function createRoot(container: Element) {
 
 // https://react.dev/reference/react/memo
 export function memo<P extends {}>(
-  fc: FC<P>,
-  propsAreEqual: (
-    prevProps: Readonly<P>,
-    nextProps: Readonly<P>
-  ) => boolean = objectShallowEqual
+  Fc: FC<P>,
+  isEqualProps: (prev: {}, next: {}) => boolean = objectShallowEqual
 ): FC<P> {
   function Memo(props: P) {
-    const prev = useRef<{ props: Readonly<P>; result: VNode } | undefined>(
-      undefined
-    );
-    if (!prev.current || !propsAreEqual(prev.current.props, props)) {
-      prev.current = {
+    const [manager] = useState(() => new MemoManager(Fc, props, isEqualProps));
+    manager.update(props);
+    return manager.vnode;
+  }
+  Object.defineProperty(Memo, "name", { value: `Memo(${Fc.name})` });
+  return Memo;
+}
+
+class MemoManager {
+  public vnode: VCustom;
+
+  // TODO
+  // currently single `VCustom.render` identity is required, but `FcOnce` must be invalidated on update.
+  // this cache is not necessary if "identical vnode" optimization is implemented.
+  // with that optimization, simple this should work:
+  //   vnode = {
+  //     type: NODE_TYPE_CUSTOM,
+  //     render: Fc,
+  //     props,
+  //   }
+  private FcOnce: FC<{}>;
+  private render = (props: {}) => this.FcOnce(props);
+
+  constructor(
+    private Fc: FC<any>,
+    props: {},
+    private isEqualProps: (prev: {}, next: {}) => boolean
+  ) {
+    this.vnode = {
+      type: NODE_TYPE_CUSTOM,
+      render: this.render,
+      props,
+    };
+    this.FcOnce = once(this.Fc);
+  }
+
+  update(props: {}) {
+    if (!this.isEqualProps(this.vnode.props, props)) {
+      this.vnode = {
+        type: NODE_TYPE_CUSTOM,
+        render: this.render,
         props,
-        result: {
-          type: NODE_TYPE_CUSTOM,
-          render: once(fc) as FC<any>,
-          props,
-        },
       };
+      this.FcOnce = once(this.Fc);
     }
-    return prev.current.result;
   }
-  Object.defineProperty(Memo, "name", { value: `Memo(${fc.name})` });
-  return Memo;
 }
 
 // from preact https://github.com/preactjs/preact/blob/4b1a7e9276e04676b8d3f8a8257469e2f732e8d4/compat/src/util.js#L19-L23
diff --git a/packages/tiny-react/src/index.test.ts b/packages/tiny-react/src/index.test.ts
index 5947b356..8a58f7e7 100644
--- a/packages/tiny-react/src/index.test.ts
+++ b/packages/tiny-react/src/index.test.ts
@@ -1731,7 +1731,6 @@ describe(memo, () => {
         </div>
       </main>
     `);
-    // TODO: it shouldn't dispose
     expect(mockFn.mock.calls).toMatchInlineSnapshot(`
       [
         [
@@ -1746,15 +1745,9 @@ describe(memo, () => {
         [
           "[effect] y-hi 0",
         ],
-        [
-          "[dispose] x-hi 0",
-        ],
         [
           "[render] x-hello 0",
         ],
-        [
-          "[effect] x-hello 0",
-        ],
       ]
     `);
 
@@ -1793,24 +1786,12 @@ describe(memo, () => {
         [
           "[effect] y-hi 0",
         ],
-        [
-          "[dispose] x-hi 0",
-        ],
         [
           "[render] x-hello 0",
         ],
-        [
-          "[effect] x-hello 0",
-        ],
-        [
-          "[dispose] x-hello 0",
-        ],
         [
           "[render] x-hi 0",
         ],
-        [
-          "[effect] x-hi 0",
-        ],
       ]
     `);
 
@@ -1848,27 +1829,12 @@ describe(memo, () => {
         [
           "[effect] y-hi 0",
         ],
-        [
-          "[dispose] x-hi 0",
-        ],
         [
           "[render] x-hello 0",
         ],
-        [
-          "[effect] x-hello 0",
-        ],
-        [
-          "[dispose] x-hello 0",
-        ],
         [
           "[render] x-hi 0",
         ],
-        [
-          "[effect] x-hi 0",
-        ],
-        [
-          "[dispose] x-hi 0",
-        ],
         [
           "[render] x-hi 1",
         ],