fix #22 (#63)

Author: lxsmnsyc

src/babel/core/create-registry.ts

@@ -4,7 +4,6 @@ import { IMPORT_REFRESH, IMPORT_REGISTRY } from './constants';
 import { getHotIdentifier } from './get-hot-identifier';
 import { getImportIdentifier } from './get-import-identifier';
 import { getRootStatementPath } from './get-root-statement-path';
-import { generateViteHMRRequirement } from './get-vite-hmr-requirement';
 import type { StateContext } from './types';
 
 const REGISTRY = 'REGISTRY';
@@ -34,22 +33,21 @@ export function createRegistry(
     )[0],
   );
   const pathToHot = getHotIdentifier(state);
-  const statements: t.Statement[] = [
-    t.expressionStatement(
-      t.callExpression(getImportIdentifier(state, path, IMPORT_REFRESH), [
-        t.stringLiteral(state.bundler),
-        pathToHot,
-        identifier,
-      ]),
-    ),
-  ];
-
-  generateViteHMRRequirement(state, statements, pathToHot);
-
   (
     path.scope.getProgramParent().path as babel.NodePath<t.Program>
   ).pushContainer('body', [
-    t.ifStatement(pathToHot, t.blockStatement(statements)),
+    t.ifStatement(
+      pathToHot,
+      t.blockStatement([
+        t.expressionStatement(
+          t.callExpression(getImportIdentifier(state, path, IMPORT_REFRESH), [
+            t.stringLiteral(state.bundler),
+            pathToHot,
+            identifier,
+          ]),
+        ),
+      ]),
+    ),
   ]);
   state.imports.set(REGISTRY, identifier);
   return identifier;

src/babel/core/get-hmr-decline-call.ts

@@ -1,23 +1,43 @@
 import type * as babel from '@babel/core';
 import * as t from '@babel/types';
-import type { StateContext } from './types';
+import { IMPORT_DECLINE } from './constants';
 import { getHotIdentifier } from './get-hot-identifier';
 import { getImportIdentifier } from './get-import-identifier';
-import { IMPORT_DECLINE } from './constants';
-import { generateViteHMRRequirement } from './get-vite-hmr-requirement';
+import type { StateContext } from './types';
 
 export function getHMRDeclineCall(state: StateContext, path: babel.NodePath) {
   const pathToHot = getHotIdentifier(state);
-  const statements = [
-    t.expressionStatement(
-      t.callExpression(getImportIdentifier(state, path, IMPORT_DECLINE), [
-        t.stringLiteral(state.bundler),
-        pathToHot,
-      ]),
-    ),
-  ];
-
-  generateViteHMRRequirement(state, statements, pathToHot);
 
-  return t.ifStatement(pathToHot, t.blockStatement(statements));
+  if (state.bundler === 'vite') {
+    return t.ifStatement(
+      pathToHot,
+      t.blockStatement([
+        t.expressionStatement(
+          t.callExpression(
+            t.memberExpression(pathToHot, t.identifier('accept')),
+            [
+              t.arrowFunctionExpression(
+                [],
+                t.callExpression(
+                  t.memberExpression(pathToHot, t.identifier('invalidate')),
+                  [],
+                ),
+              ),
+            ],
+          ),
+        ),
+      ]),
+    );
+  }
+  return t.ifStatement(
+    pathToHot,
+    t.blockStatement([
+      t.expressionStatement(
+        t.callExpression(getImportIdentifier(state, path, IMPORT_DECLINE), [
+          t.stringLiteral(state.bundler),
+          pathToHot,
+        ]),
+      ),
+    ]),
+  );
 }