fix hmr

packages/myreact-reconciler/src/runtimeFiber/hmr.ts

@@ -5,17 +5,19 @@ import { deleteEffect } from "../dispatchEffect";
 import { listenerMap } from "../renderDispatch";
 import { classComponentUnmount } from "../runtimeComponent";
 import { hookListUnmount } from "../runtimeHook";
-import { fiberToDispatchMap, getCurrentDispatchFromFiber, setRefreshTypeMap } from "../share";
+import { getCurrentDispatchFromFiber, setRefreshTypeMap } from "../share";
 
 import type { MyReactFiberNode } from "./instance";
 import type { MyReactFiberNodeDev } from "./interface";
-import type { MixinMyReactFunctionComponent, MixinMyReactClassComponent } from "@my-react/react";
+import type { MyReactComponentType } from "@my-react/react";
 
-export const hmr = (fiber: MyReactFiberNode, nextType: MixinMyReactFunctionComponent | MixinMyReactClassComponent, forceRefresh?: boolean) => {
+export const hmr = (fiber: MyReactFiberNode, nextType: MyReactComponentType, forceRefresh?: boolean) => {
   if (__DEV__) {
     if (include(fiber.state, STATE_TYPE.__unmount__)) return;
 
-    const element = createElement(nextType as MixinMyReactFunctionComponent, fiber.pendingProps);
+    const renderDispatch = getCurrentDispatchFromFiber(fiber);
+
+    const element = createElement(nextType, fiber.pendingProps);
 
     fiber._installElement(element);
 
@@ -32,22 +34,20 @@ export const hmr = (fiber: MyReactFiberNode, nextType: MixinMyReactFunctionCompo
 
       fiber.updateQueue = null;
 
-      const renderDispatch = fiberToDispatchMap.get(fiber);
-
       const typedFiber = fiber as MyReactFiberNodeDev;
 
       typedFiber._debugHookTypes = [];
 
       // TODO
       deleteEffect(fiber, renderDispatch);
 
+      renderDispatch.commitUnsetRef(fiber);
+
       fiber.state = merge(STATE_TYPE.__create__, STATE_TYPE.__hmr__);
     } else {
       fiber.state = merge(STATE_TYPE.__triggerSync__, STATE_TYPE.__hmr__);
     }
 
-    const renderDispatch = getCurrentDispatchFromFiber(fiber);
-
     listenerMap.get(renderDispatch)?.fiberHMR?.forEach((cb) => cb(fiber));
 
     return fiber;

packages/myreact-reconciler/src/runtimeGenerate/generate.ts

@@ -141,9 +141,11 @@ const getNewFiberWithInitial = (newChild: MaybeArrayMyReactElementNode, parentFi
 export const transformChildrenFiber = (parentFiber: MyReactFiberNode, children: MaybeArrayMyReactElementNode): void => {
   const isUpdate = exclude(parentFiber.state, STATE_TYPE.__create__);
 
+  const isHMR = include(parentFiber.state, STATE_TYPE.__hmr__);
+
   const renderDispatch = currentRenderDispatch.current;
 
-  if (isUpdate) {
+  if (isUpdate || isHMR) {
     const { existingChildrenMap, existingChildrenArray } = getExistingChildren(parentFiber);
 
     parentFiber.child = null;

packages/myreact-refresh/src/RefreshRuntime.ts

@@ -1,6 +1,7 @@
+/* eslint-disable max-lines */
 import { ForwardRef, Memo, STATE_TYPE, TYPEKEY } from "@my-react/react-shared";
 
-import type { forwardRef, memo, MixinMyReactClassComponent, MixinMyReactFunctionComponent, MyReactElementType } from "@my-react/react";
+import type { MixinMyReactClassComponent, MixinMyReactFunctionComponent, MyReactComponentType, MyReactElementType } from "@my-react/react";
 import type { CustomRenderDispatch, CustomRenderDispatchDev, HMR, MyReactFiberNode, MyReactFiberNodeDev } from "@my-react/react-reconciler";
 
 const DISPATCH_FIELD = "__@my-react/dispatch__";
@@ -23,8 +24,6 @@ type Signature = {
   fullKey?: string;
 };
 
-type MyReactComponentType = ReturnType<typeof forwardRef> | ReturnType<typeof memo> | MixinMyReactClassComponent | MixinMyReactFunctionComponent;
-
 type HMRGlobal = {
   [DISPATCH_FIELD]: CustomRenderDispatch[];
   [DEV_REFRESH_FIELD]: (dispatchArray: CustomRenderDispatch[]) => void;
@@ -106,6 +105,26 @@ const haveEqualSignatures = (prevType: MyReactComponentType, nextType: MyReactCo
   return true;
 };
 
+// check memo | forwardRef | function
+const hasEqualType = (prevType: MyReactComponentType, nextType: MyReactComponentType) => {
+  if (prevType === nextType) {
+    return true;
+  }
+
+  if (typeof prevType !== typeof nextType) {
+    return false;
+  } else {
+    if (typeof prevType === "object" && typeof nextType === "object") {
+      return (
+        getProperty(prevType, TYPEKEY) === getProperty(nextType, TYPEKEY) &&
+        hasEqualType(prevType.render as MyReactComponentType, nextType.render as MyReactComponentType)
+      );
+    }
+  }
+
+  return true;
+};
+
 const getRenderTypeFormType = (type: MyReactComponentType): MixinMyReactClassComponent | MixinMyReactFunctionComponent | null => {
   if (!type) {
     console.error(`[@my-react/react-refresh] can not get the real type for current render, it is a bug for @my-react/react-refresh`);
@@ -141,7 +160,7 @@ const canPreserveStateBetween = (prevType: MyReactComponentType, nextType: MyRea
   if (isMyReactClass(prevType) || isMyReactClass(nextType)) {
     return false;
   }
-  if (haveEqualSignatures(prevType, nextType)) {
+  if (hasEqualType(prevType, nextType) && haveEqualSignatures(prevType, nextType)) {
     return true;
   }
   return false;
@@ -263,7 +282,9 @@ export const performReactRefresh = () => {
   const containers: Map<CustomRenderDispatch, boolean> = new Map();
 
   allPending.forEach(([family, _nextType]) => {
-    const prevType = getRenderTypeFormType(family.current);
+    const _prevType = family.current;
+
+    const prevType = getRenderTypeFormType(_prevType);
 
     const nextType = getRenderTypeFormType(_nextType);
 
@@ -282,10 +303,10 @@ export const performReactRefresh = () => {
 
       updatedFamiliesByType.set(_nextType, family);
 
-      family.current = nextType;
+      family.current = _nextType;
 
       if (fibers?.size) {
-        const forceReset = !canPreserveStateBetween(prevType, nextType);
+        const forceReset = !canPreserveStateBetween(_prevType, _nextType);
 
         fibers.forEach((f) => {
           let container: CustomRenderDispatchDev | null = null;
@@ -302,7 +323,7 @@ export const performReactRefresh = () => {
 
           const hasRootUpdate = containers.get(container) || f === container.rootFiber;
 
-          container.__hmr_runtime__?.hmr?.(f, nextType, forceReset);
+          container.__hmr_runtime__?.hmr?.(f, _nextType, forceReset);
 
           containers.set(container, hasRootUpdate);
         });

packages/myreact/src/element/instance.ts

@@ -63,6 +63,8 @@ export type MixinMyReactFunctionComponent<P extends Record<string, unknown> = an
   defaultProps?: Record<string, unknown>;
 };
 
+export type MyReactComponentType = ReturnType<typeof forwardRef> | ReturnType<typeof memo> | MixinMyReactClassComponent | MixinMyReactFunctionComponent;
+
 /**
  * @public
  */

packages/myreact/src/index.ts

@@ -197,6 +197,7 @@ export type {
   MyReactElement,
   MyReactElementType,
   MyReactElementNode,
+  MyReactComponentType,
   ArrayMyReactElementNode,
   ArrayMyReactElementChildren,
   MaybeArrayMyReactElementNode,

ui/vite-ssr-example/src/components/A.tsx

@@ -0,0 +1,11 @@
+import { useState, memo } from "react";
+
+export const A = memo(() => {
+  const [s, setS] = useState(0);
+  return (
+    <div>
+      A component: {s} <br />
+      <button onClick={() => setS(s + 1)}>add A</button>
+    </div>
+  );
+});

ui/vite-ssr-example/src/page/Foo.tsx

@@ -1,3 +1,10 @@
+import { A } from "../components/A";
+
 export default function Foo() {
-  return <div>Foo page</div>;
+  return (
+    <>
+      <div>Foo page</div>
+      <A />
+    </>
+  );
 }