fix(tiny-react): fix memo

.changeset/polite-months-cough.md

@@ -0,0 +1,5 @@
+---
+"@hiogawa/tiny-react": patch
+---
+
+fix: fix `memo`

packages/tiny-react/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@hiogawa/tiny-react",
-  "version": "0.0.1",
+  "version": "0.0.2-pre.1",
   "homepage": "https://github.com/hi-ogawa/js-utils/tree/main/packages/tiny-react",
   "repository": {
     "type": "git",

packages/tiny-react/src/compat/index.ts

@@ -1,6 +1,14 @@
+import { once } from "@hiogawa/utils";
 import { useEffect, useRef, useState } from "../hooks";
 import { render } from "../reconciler";
-import { type BNode, EMPTY_NODE, type FC, type VNode } from "../virtual-dom";
+import {
+  type BNode,
+  EMPTY_NODE,
+  type FC,
+  NODE_TYPE_CUSTOM,
+  type VCustom,
+  type VNode,
+} from "../virtual-dom";
 
 // non comprehensive compatibility features
 
@@ -39,23 +47,43 @@ export function createRoot(container: Element) {
 }
 
 // https://react.dev/reference/react/memo
-export function memo<P extends object>(
-  fc: FC<P>,
-  propsAreEqual: (
-    prevProps: Readonly<P>,
-    nextProps: Readonly<P>
-  ) => boolean = objectShallowEqual
+export function memo<P extends {}>(
+  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 = { props, result: fc(props) };
+    // we need to make `VCustom.render` stable,
+    // but `once(Fc)` has to be invalidated on props change.
+    // after "identical vnode" optimization is implemented,
+    // it can be simplified to
+    //   {
+    //     type: NODE_TYPE_CUSTOM,
+    //     render: Fc,
+    //     props,
+    //   }
+    const [state] = useState(() => {
+      const state: {
+        render: FC;
+        current?: { vnode: VCustom; onceFc: FC };
+      } = {
+        render: (props: any) => state.current!.onceFc(props),
+      };
+      return state;
+    });
+
+    if (!state.current || !isEqualProps(state.current.vnode.props, props)) {
+      state.current = {
+        vnode: {
+          type: NODE_TYPE_CUSTOM,
+          render: state.render,
+          props,
+        },
+        onceFc: once(Fc),
+      };
     }
-    return prev.current.result;
+    return state.current.vnode;
   }
-  Object.defineProperty(Memo, "name", { value: `Memo(${fc.name})` });
+  Object.defineProperty(Memo, "name", { value: `Memo(${Fc.name})` });
   return Memo;
 }
 

packages/tiny-react/src/index.test.ts

@@ -12,7 +12,7 @@ import {
 } from "./hooks";
 import { render, updateCustomNode } from "./reconciler";
 import { sleepFrame } from "./test-utils";
-import { getBNodeSlot } from "./virtual-dom";
+import { EMPTY_NODE, getBNodeSlot } from "./virtual-dom";
 
 describe(render, () => {
   it("basic", () => {
@@ -1613,14 +1613,23 @@ describe("ref", () => {
 });
 
 describe(memo, () => {
-  it("basic", () => {
+  it("basic", async () => {
     const mockFn = vi.fn();
+    const mockFnSnapshot = () => mockFn.mock.calls.map((args) => args[0]);
 
     const Custom = memo(function Custom(props: {
       label: string;
       value: number;
     }) {
-      mockFn(props.label, props.value);
+      mockFn(`[render] ${props.label} ${props.value}`);
+
+      useEffect(() => {
+        mockFn(`[effect] ${props.label} ${props.value}`);
+        return () => {
+          mockFn(`[dispose] ${props.label} ${props.value}`);
+        };
+      }, []);
+
       return h.div({}, props.label, props.value);
     });
 
@@ -1634,6 +1643,7 @@ describe(memo, () => {
         h(Custom, { label: "y-hi", value: 0 })
       )
     );
+    await sleepFrame();
     expect(parent).toMatchInlineSnapshot(`
       <main>
         <div>
@@ -1646,19 +1656,16 @@ describe(memo, () => {
         </div>
       </main>
     `);
-    expect(mockFn.mock.calls).toMatchInlineSnapshot(`
+    expect(mockFnSnapshot()).toMatchInlineSnapshot(`
       [
-        [
-          "x-hi",
-          0,
-        ],
-        [
-          "y-hi",
-          0,
-        ],
+        "[render] x-hi 0",
+        "[render] y-hi 0",
+        "[effect] x-hi 0",
+        "[effect] y-hi 0",
       ]
     `);
 
+    mockFn.mockReset();
     root.render(
       h(
         Fragment,
@@ -1667,6 +1674,7 @@ describe(memo, () => {
         h(Custom, { label: "y-hi", value: 0 })
       )
     );
+    await sleepFrame();
     expect(parent).toMatchInlineSnapshot(`
       <main>
         <div>
@@ -1679,19 +1687,9 @@ describe(memo, () => {
         </div>
       </main>
     `);
-    expect(mockFn.mock.calls).toMatchInlineSnapshot(`
-      [
-        [
-          "x-hi",
-          0,
-        ],
-        [
-          "y-hi",
-          0,
-        ],
-      ]
-    `);
+    expect(mockFnSnapshot()).toMatchInlineSnapshot("[]");
 
+    mockFn.mockReset();
     root.render(
       h(
         Fragment,
@@ -1700,6 +1698,7 @@ describe(memo, () => {
         h(Custom, { label: "y-hi", value: 0 })
       )
     );
+    await sleepFrame();
     expect(parent).toMatchInlineSnapshot(`
       <main>
         <div>
@@ -1712,23 +1711,13 @@ describe(memo, () => {
         </div>
       </main>
     `);
-    expect(mockFn.mock.calls).toMatchInlineSnapshot(`
+    expect(mockFnSnapshot()).toMatchInlineSnapshot(`
       [
-        [
-          "x-hi",
-          0,
-        ],
-        [
-          "y-hi",
-          0,
-        ],
-        [
-          "x-hello",
-          0,
-        ],
+        "[render] x-hello 0",
       ]
     `);
 
+    mockFn.mockReset();
     root.render(
       h(
         Fragment,
@@ -1737,6 +1726,7 @@ describe(memo, () => {
         h(Custom, { label: "y-hi", value: 0 })
       )
     );
+    await sleepFrame();
     expect(parent).toMatchInlineSnapshot(`
       <main>
         <div>
@@ -1749,27 +1739,13 @@ describe(memo, () => {
         </div>
       </main>
     `);
-    expect(mockFn.mock.calls).toMatchInlineSnapshot(`
+    expect(mockFnSnapshot()).toMatchInlineSnapshot(`
       [
-        [
-          "x-hi",
-          0,
-        ],
-        [
-          "y-hi",
-          0,
-        ],
-        [
-          "x-hello",
-          0,
-        ],
-        [
-          "x-hi",
-          0,
-        ],
+        "[render] x-hi 0",
       ]
     `);
 
+    mockFn.mockReset();
     root.render(
       h(
         Fragment,
@@ -1790,28 +1766,19 @@ describe(memo, () => {
         </div>
       </main>
     `);
-    expect(mockFn.mock.calls).toMatchInlineSnapshot(`
+    expect(mockFnSnapshot()).toMatchInlineSnapshot(`
       [
-        [
-          "x-hi",
-          0,
-        ],
-        [
-          "y-hi",
-          0,
-        ],
-        [
-          "x-hello",
-          0,
-        ],
-        [
-          "x-hi",
-          0,
-        ],
-        [
-          "x-hi",
-          1,
-        ],
+        "[render] x-hi 1",
+      ]
+    `);
+
+    mockFn.mockReset();
+    root.render(EMPTY_NODE);
+    expect(parent).toMatchInlineSnapshot("<main />");
+    expect(mockFnSnapshot()).toMatchInlineSnapshot(`
+      [
+        "[dispose] x-hi 0",
+        "[dispose] y-hi 0",
       ]
     `);
   });