跳到主要内容

官方模拟版

2023年06月11日
柏拉文
越努力，越幸运

一、模块依赖图开发​

模块依赖图在 no-bundle 构建服务中是一个不可或缺的数据结构，一方面可以存储各个模块的信息，用于记录编译缓存，另一方面也可以记录各个模块间的依赖关系，用于实现 HMR。

接下来我们来实现模块依赖图，即 ModuleGraph 类，新建src/node/ModuleGraph.ts，内容如下:

import { PartialResolvedId, TransformResult } from "rollup";
import { cleanUrl } from "./utils";

export class ModuleNode {
  // 资源访问 url
  url: string;
  // 资源绝对路径
  id: string | null = null;
  importers = new Set<ModuleNode>();
  importedModules = new Set<ModuleNode>();
  transformResult: TransformResult | null = null;
  lastHMRTimestamp = 0;
  constructor(url: string) {
    this.url = url;
  }
}

export class ModuleGraph {
  // 资源 url 到 ModuleNode 的映射表
  urlToModuleMap = new Map<string, ModuleNode>();
  // 资源绝对路径到 ModuleNode 的映射表
  idToModuleMap = new Map<string, ModuleNode>();

  constructor(
    private resolveId: (url: string) => Promise<PartialResolvedId | null>
  ) {}

  getModuleById(id: string): ModuleNode | undefined {
    return this.idToModuleMap.get(id);
  }

  async getModuleByUrl(rawUrl: string): Promise<ModuleNode | undefined> {
    const { url } = await this._resolve(rawUrl);
    return this.urlToModuleMap.get(url);
  }

  async ensureEntryFromUrl(rawUrl: string): Promise<ModuleNode> {
    const { url, resolvedId } = await this._resolve(rawUrl);
    // 首先检查缓存
    if (this.urlToModuleMap.has(url)) {
      return this.urlToModuleMap.get(url) as ModuleNode;
    }
    // 若无缓存，更新 urlToModuleMap 和 idToModuleMap
    const mod = new ModuleNode(url);
    mod.id = resolvedId;
    this.urlToModuleMap.set(url, mod);
    this.idToModuleMap.set(resolvedId, mod);
    return mod;
  }

  async updateModuleInfo(
    mod: ModuleNode,
    importedModules: Set<string | ModuleNode>
  ) {
    const prevImports = mod.importedModules;
    for (const curImports of importedModules) {
      const dep =
        typeof curImports === "string"
          ? await this.ensureEntryFromUrl(cleanUrl(curImports))
          : curImports;
      if (dep) {
        mod.importedModules.add(dep);
        dep.importers.add(mod);
      }
    }
    // 清除已经不再被引用的依赖
    for (const prevImport of prevImports) {
      if (!importedModules.has(prevImport.url)) {
        prevImport.importers.delete(mod);
      }
    }
  }

  // HMR 触发时会执行这个方法
  invalidateModule(file: string) {
    const mod = this.idToModuleMap.get(file);
    if (mod) {
      // 更新时间戳
      mod.lastHMRTimestamp = Date.now();
      mod.transformResult = null;
      mod.importers.forEach((importer) => {
        this.invalidateModule(importer.id!);
      });
    }
  }

  private async _resolve(
    url: string
  ): Promise<{ url: string; resolvedId: string }> {
    const resolved = await this.resolveId(url);
    const resolvedId = resolved?.id || url;
    return { url, resolvedId };
  }
}

接着我们看看如何将这个 ModuleGraph 接入到目前的架构中。

首先在服务启动前，我们需要初始化 ModuleGraph 实例:

// src/node/server/index.ts
export interface ServerContext {
  root: string;
  pluginContainer: PluginContainer;
  app: connect.Server;
  plugins: Plugin[];
  moduleGraph: ModuleGraph;
}

export async function startDevServer() {
  const app = connect();
  const root = process.cwd();
  const startTime = Date.now();

  const plugins = resolvePlugins();
  const pluginContainer = createPluginContainer(plugins);
  const moduleGraph = new ModuleGraph((url) => pluginContainer.resolveId(url));
  const serverContext: ServerContext = {
    root: process.cwd(),
    app,
    pluginContainer,
    plugins,
    moduleGraph
  };

  for (const plugin of plugins) {
    if (plugin.configureServer) {
      await plugin.configureServer(serverContext);
    }
  }


  app.use(transformMiddleware(serverContext));
  app.use(indexHtmlMiddleware(serverContext));
  app.use(staticMiddleware(serverContext.root));

  app.listen(3000, async () => {
    await optimize(root);

    console.log(
      green('🚀 No-Bundle 服务已经成功启动!'),
      `耗时: ${Date.now() - startTime}ms`
    );
    console.log(`> 本地访问路径: ${blue('http://localhost:3000')}`);
  });
}

然后在加载完模块后，也就是调用插件容器的 load 方法后，我们需要通过 ensureEntryFromUrl 方法注册模块:

// src/node/server/middlewares/transform.ts
export async function transformRequest(
  url: string,
  serverContext: ServerContext
) {
  const { moduleGraph, pluginContainer } = serverContext;
  
  url = cleanUrl(url);

  let mod = await moduleGraph.getModuleByUrl(url);
  if (mod && mod.transformResult) {
    return mod.transformResult;
  }
  

  const resolvedResult = await pluginContainer.resolveId(url);
  let transformResult;

  if (resolvedResult?.id) {
    let code = await pluginContainer.load(resolvedResult.id);
    if (typeof code === 'object' && code !== null) {
      code = code.code;
    }

    mod = await moduleGraph.ensureEntryFromUrl(url);

    if (code) {
      transformResult = await pluginContainer.transform(
        code as string,
        resolvedResult?.id
      );
    }
  }

  if (mod) {
    mod.transformResult = transformResult;
  }
  
  return transformResult;
}

当我们对 JS 模块分析完 import 语句之后，需要更新模块之间的依赖关系:

// src/node/plugins/importAnalysis.ts
export function importAnalysisPlugin(): Plugin {
  let serverContext: ServerContext;
  return {
    name: 'm-vite:import-analysis',
    configureServer(s) {
      serverContext = s;
    },
    async transform(code: string, id: string) {
      if (!isJSRequest(id)) {
        return null;
      }
      await init;
      const importedModules = new Set<string>();
      const [imports] = parse(code);
      const ms = new MagicString(code);

      const resolve = async (id: string, importer?: string) => {
        const resolved = await serverContext.pluginContainer.resolveId(
          id,
          normalizePath(importer)
        );
        if (!resolved) {
          return;
        }
        let resolvedId = `/${getShortName(resolved.id, serverContext.root)}`;
        return resolvedId;
      };

      const { moduleGraph } = serverContext;
      const curMod = moduleGraph.getModuleById(id)!;
      
      for (const importInfo of imports) {
        const { s: modStart, e: modEnd, n: modSource } = importInfo;
        if (!modSource) continue;

        if (modSource.endsWith('.svg')) {
          const resolvedUrl = await resolve(modSource, id);
          ms.overwrite(modStart, modEnd, `${resolvedUrl}?import`);
          continue;
        }

        if (BARE_IMPORT_RE.test(modSource)) {
          const bundlePath = normalizePath(
            path.join('/', PRE_BUNDLE_DIR, `${modSource}.js`)
          );
          ms.overwrite(modStart, modEnd, bundlePath);
          importedModules.add(bundlePath);
        } else if (modSource.startsWith('.') || modSource.startsWith('/')) {
          const resolved = await resolve(modSource, id);
          if (resolved) {
            ms.overwrite(modStart, modEnd, resolved);
            importedModules.add(resolved);
          }
        }
      }

      moduleGraph.updateModuleInfo(curMod, importedModules);
      
      return {
        code: ms.toString(),
        map: ms.generateMap()
      };
    }
  };
}

现在，一个完整的模块依赖图就能随着 JS 请求的到来而不断建立起来了。另外，基于现在的模块依赖图，我们也可以记录模块编译后的产物，并进行缓存。让我们回到 transform 中间件中:

export async function transformRequest(
  url: string,
  serverContext: ServerContext
) {
  const { moduleGraph, pluginContainer } = serverContext;
  url = cleanUrl(url);

  let mod = await moduleGraph.getModuleByUrl(url);
  if (mod && mod.transformResult) {
    return mod.transformResult;
  }

  const resolvedResult = await pluginContainer.resolveId(url);
  let transformResult;

  if (resolvedResult?.id) {
    let code = await pluginContainer.load(resolvedResult.id);
    if (typeof code === 'object' && code !== null) {
      code = code.code;
    }

    mod = await moduleGraph.ensureEntryFromUrl(url);

    if (code) {
      transformResult = await pluginContainer.transform(
        code as string,
        resolvedResult?.id
      );
    }
  }

  if (mod) {
    mod.transformResult = transformResult;
  }
  
  return transformResult;
}

在搭建好模块依赖图之后，我们把目光集中到最重要的部分——HMR 上面。

二、HMR 服务端开发​

HMR 在服务端需要完成如下的工作:

  1. 创建文件监听器，以监听文件的变动

  2. 创建 WebSocket 服务端，负责和客户端进行通信

  3. 文件变动时，从 ModuleGraph 中定位到需要更新的模块，将更新信息发送给客户端

首先，我们来创建文件监听器:

// src/node/server/index.ts
import connect from 'connect';
import { Plugin } from "../plugin";
import { blue, green } from 'picocolors';
import { resolvePlugins } from '../plugins';
import { ModuleGraph } from "../ModuleGraph";
import { optimize } from '../optimizer/index';
import chokidar, { FSWatcher } from "chokidar";
import { staticMiddleware } from "./middlewares/static";
import { transformMiddleware } from './middlewares/transform';
import { indexHtmlMiddleware } from "./middlewares/indexHtml";
import { createPluginContainer, PluginContainer } from '../pluginContainer';

export interface ServerContext {
  root: string;
  pluginContainer: PluginContainer;
  app: connect.Server;
  plugins: Plugin[];
  moduleGraph: ModuleGraph;
}

export async function startDevServer() {
  const app = connect();
  const root = process.cwd();
  const startTime = Date.now();

  const plugins = resolvePlugins();
  const pluginContainer = createPluginContainer(plugins);
  const watcher = chokidar.watch(root, {
    ignored: ["**/node_modules/**", "**/.git/**"],
    ignoreInitial: true,
  });
  

  const moduleGraph = new ModuleGraph((url) => pluginContainer.resolveId(url));
  const serverContext: ServerContext = {
    root: process.cwd(),
    app,
    pluginContainer,
    plugins,
    moduleGraph
  };

  for (const plugin of plugins) {
    if (plugin.configureServer) {
      await plugin.configureServer(serverContext);
    }
  }


  ……
  app.listen(3000, async () => {
    await optimize(root);

    console.log(
      green('🚀 No-Bundle 服务已经成功启动!'),
      `耗时: ${Date.now() - startTime}ms`
    );
    console.log(`> 本地访问路径: ${blue('http://localhost:3000')}`);
  });
}

接着初始化 WebSocket 服务端，新建src/node/ws.ts，内容如下:

import connect from 'connect';
import { red } from 'picocolors';
import { HMR_PORT } from './constants';
import { WebSocketServer, WebSocket } from 'ws';

export function createWebSocketServer(server: connect.Server): {
  send: (msg: string) => void;
  close: () => void;
} {
  let wss: WebSocketServer;
  wss = new WebSocketServer({ port: HMR_PORT });
  wss.on('connection', socket => {
    socket.send(JSON.stringify({ type: 'connected' }));
  });

  wss.on('error', (e: Error & { code: string }) => {
    if (e.code !== 'EADDRINUSE') {
      console.error(red(`WebSocket server error:\n${e.stack || e.message}`));
    }
  });

  return {
    send(payload: Object) {
      const stringified = JSON.stringify(payload);
      wss.clients.forEach(client => {
        if (client.readyState === WebSocket.OPEN) {
          client.send(stringified);
        }
      });
    },

    close() {
      wss.close();
    }
  };
}

同时定义 HMR_PORT 常量:

// src/node/constants.ts
export const HMR_PORT = 24678;

接着我们将 WebSocket 服务端实例加入 no-bundle 服务中:

// src/node/server/index.ts
import connect from 'connect';
import { Plugin } from "../plugin";
import { blue, green } from 'picocolors';
import { resolvePlugins } from '../plugins';
import { ModuleGraph } from "../ModuleGraph";
import { createWebSocketServer } from '../ws';
import { optimize } from '../optimizer/index';
import chokidar, { FSWatcher } from "chokidar";
import { staticMiddleware } from "./middlewares/static";
import { transformMiddleware } from './middlewares/transform';
import { indexHtmlMiddleware } from "./middlewares/indexHtml";
import { createPluginContainer, PluginContainer } from '../pluginContainer';

export interface ServerContext {
  root: string;
  pluginContainer: PluginContainer;
  app: connect.Server;
  plugins: Plugin[];
  moduleGraph: ModuleGraph;
  ws: { send: (data: any) => void; close: () => void };
  watcher: FSWatcher;
}

export async function startDevServer() {
  const app = connect();
  const root = process.cwd();
  const startTime = Date.now();

  const plugins = resolvePlugins();
  const pluginContainer = createPluginContainer(plugins);
  const watcher = chokidar.watch(root, {
    ignored: ["**/node_modules/**", "**/.git/**"],
    ignoreInitial: true,
  });
  const ws = createWebSocketServer(app);

  const moduleGraph = new ModuleGraph((url) => pluginContainer.resolveId(url));
  const serverContext: ServerContext = {
    root: process.cwd(),
    app,
    pluginContainer,
    plugins,
    moduleGraph,
    ws,
    watcher
  };

  for (const plugin of plugins) {
    if (plugin.configureServer) {
      await plugin.configureServer(serverContext);
    }
  }


  app.use(transformMiddleware(serverContext));
  app.use(indexHtmlMiddleware(serverContext));
  app.use(staticMiddleware(serverContext.root));

  app.listen(3000, async () => {
    await optimize(root);

    console.log(
      green('🚀 No-Bundle 服务已经成功启动!'),
      `耗时: ${Date.now() - startTime}ms`
    );
    console.log(`> 本地访问路径: ${blue('http://localhost:3000')}`);
  });
}

下面我们来实现当文件变动时，服务端具体的处理逻辑，新建 src/node/hmr.ts:

import { getShortName } from "./utils";
import { blue, green } from "picocolors";
import { ServerContext } from "./server/index";

export function bindingHMREvents(serverContext: ServerContext) {
  const { watcher, ws, root } = serverContext;

  watcher.on("change", async (file) => {
    console.log(`✨${blue("[hmr]")} ${green(file)} changed`);
    const { moduleGraph } = serverContext;
    // 清除模块依赖图中的缓存
    await moduleGraph.invalidateModule(file);
    // 向客户端发送更新信息
    ws.send({
      type: "update",
      updates: [
        {
          type: "js-update",
          timestamp: Date.now(),
          path: "/" + getShortName(file, root),
          acceptedPath: "/" + getShortName(file, root),
        },
      ],
    });
  });
}

注意补充一下缺失的工具函数:

// src/node/utils.ts
export function getShortName(file: string, root: string) {
  return file.startsWith(root + "/") ? path.posix.relative(root, file) : file;
}

接着我们在服务中添加如下代码:

// src/node/server/index.ts
import connect from 'connect';
import { Plugin } from "../plugin";
import { normalizePath } from "../utils";
import { blue, green } from 'picocolors';
import { bindingHMREvents } from "../hmr";
import { resolvePlugins } from '../plugins';
import { ModuleGraph } from "../ModuleGraph";
import { createWebSocketServer } from '../ws';
import { optimize } from '../optimizer/index';
import chokidar, { FSWatcher } from "chokidar";
import { staticMiddleware } from "./middlewares/static";
import { transformMiddleware } from './middlewares/transform';
import { indexHtmlMiddleware } from "./middlewares/indexHtml";
import { createPluginContainer, PluginContainer } from '../pluginContainer';

export interface ServerContext {
  root: string;
  pluginContainer: PluginContainer;
  app: connect.Server;
  plugins: Plugin[];
  moduleGraph: ModuleGraph;
  ws: { send: (data: any) => void; close: () => void };
  watcher: FSWatcher;
}

export async function startDevServer() {
  const app = connect();
  const root = process.cwd();
  const startTime = Date.now();

  const plugins = resolvePlugins();
  const pluginContainer = createPluginContainer(plugins);
  const watcher = chokidar.watch(root, {
    ignored: ["**/node_modules/**", "**/.git/**"],
    ignoreInitial: true,
  });
  const ws = createWebSocketServer(app);

  const moduleGraph = new ModuleGraph((url) => pluginContainer.resolveId(url));
  const serverContext: ServerContext = {
    root: normalizePath(process.cwd()),
    app,
    pluginContainer,
    plugins,
    moduleGraph,
    ws,
    watcher
  };

  bindingHMREvents(serverContext);

  for (const plugin of plugins) {
    if (plugin.configureServer) {
      await plugin.configureServer(serverContext);
    }
  }


  app.use(transformMiddleware(serverContext));
  app.use(indexHtmlMiddleware(serverContext));
  app.use(staticMiddleware(serverContext.root));

  app.listen(3000, async () => {
    await optimize(root);

    console.log(
      green('🚀 No-Bundle 服务已经成功启动!'),
      `耗时: ${Date.now() - startTime}ms`
    );
    console.log(`> 本地访问路径: ${blue('http://localhost:3000')}`);
  });
}

三、HMR 客户端开发​

HMR 客户端指的是我们向浏览器中注入的一段 JS 脚本，这段脚本中会做如下的事情:

  • 创建 WebSocket 客户端，用于和服务端通信

  • 在收到服务端的更新信息后，通过动态 import 拉取最新的模块内容，执行 accept 更新回调

  • 暴露 HMR 的一些工具函数，比如 import.meta.hot 对象的实现

首先我们来开发客户端的脚本内容，你可以新建src/client/client.ts文件，然后在 tsup.config.ts 中增加如下的配置:

import { defineConfig } from "tsup";

export default defineConfig({
  // 后续会增加 entry
  entry: {
    index: "src/node/cli.ts",
    client: "src/client/client.ts",
  },
  // 产物格式，包含 esm 和 cjs 格式
  format: ["esm", "cjs"],
  // 目标语法
  target: "es2020",
  // 生成 sourcemap
  sourcemap: true,
  // 没有拆包的需求，关闭拆包能力
  splitting: false,
});

注: 改动 tsup 配置之后，为了使最新配置生效，你需要在 mini-vite 项目中执行 pnpm start 重新进行构建。

客户端脚本的具体实现如下:

// src/client/client.ts
console.log("[vite] connecting...");

// 1. 创建客户端 WebSocket 实例
// 其中的 __HMR_PORT__ 之后会被 no-bundle 服务编译成具体的端口号
const socket = new WebSocket(`ws://localhost:__HMR_PORT__`, "vite-hmr");

// 2. 接收服务端的更新信息
socket.addEventListener("message", async ({ data }) => {
  handleMessage(JSON.parse(data)).catch(console.error);
});

// 3. 根据不同的更新类型进行更新
async function handleMessage(payload: any) {
  switch (payload.type) {
    case "connected":
      console.log(`[vite] connected.`);
      // 心跳检测
      setInterval(() => socket.send("ping"), 1000);
      break;

    case "update":
      // 进行具体的模块更新
      payload.updates.forEach((update: Update) => {
        if (update.type === "js-update") {
          // 具体的更新逻辑，后续来开发
        }
      });
      break;
  }
}

关于客户端具体的 JS 模块更新逻辑和工具函数的实现，你暂且不用过于关心。我们先把这段比较简单的 HMR 客户端代码注入到浏览器中，首先在新建 src/node/plugins/clientInject.ts，内容如下:

import path from 'path';
import fs from 'fs-extra';
import { Plugin } from '../plugin';
import { ServerContext } from '../server/index';
import { CLIENT_PUBLIC_PATH, HMR_PORT } from '../constants';

export function clientInjectPlugin(): Plugin {
  let serverContext: ServerContext;
  return {
    name: 'm-vite:client-inject',
    configureServer(s) {
      serverContext = s;
    },
    resolveId(id) {
      if (id === CLIENT_PUBLIC_PATH) {
        return { id };
      }
      return null;
    },
    async load(id) {
      // 加载 HMR 客户端脚本
      if (id === CLIENT_PUBLIC_PATH) {
        const realPath = path.join(
          serverContext.root,
          'node_modules',
          'mini-vite',
          'dist',
          'client.mjs'
        );
        const code = await fs.readFile(realPath, 'utf-8');
        return {
          // 替换占位符
          code: code.replace('__HMR_PORT__', JSON.stringify(HMR_PORT))
        };
      }
    },
    transformIndexHtml(raw) {
      // 插入客户端脚本
      // 即在 head 标签后面加上 <script type="module" src="/@vite/client"></script>
      // 注: 在 indexHtml 中间件里面会自动执行 transformIndexHtml 钩子
      return raw.replace(
        /(<head[^>]*>)/i,
        `$1<script type="module" src="${CLIENT_PUBLIC_PATH}"></script>`
      );
    }
  };
}

同时添加相应的常量声明:

// src/node/constants.ts
export const CLIENT_PUBLIC_PATH = "/@vite/client";

接着我们来注册这个插件:

// src/node/plugins/index.ts
import { cssPlugin } from './css';
import { Plugin } from '../plugin';
import { assetPlugin } from "./assets";
import { resolvePlugin } from './resolve';
import { esbuildTransformPlugin } from './esbuild';
import { clientInjectPlugin } from './clientInject';
import { importAnalysisPlugin } from './importAnalysis';

export function resolvePlugins(): Plugin[] {
  return [
    clientInjectPlugin(),
    resolvePlugin(),
    esbuildTransformPlugin(),
    importAnalysisPlugin(),
    cssPlugin(),
    assetPlugin(),
  ];
}

需要注意的是，clientInject插件最好放到最前面的位置，以免后续插件的 load 钩子干扰客户端脚本的加载。

接下来你可以在 test 项目下执行pnpm dev，然后查看页面，可以发现控制台出现了如下的 log 信息:

Preview

查看网络面板，也能发现客户端脚本的请求被正常响应:

Preview

OK，接下来我们就来继续完善客户端脚本的具体实现。

值得一提的是，之所以我们可以在代码中编写类似import.meta.hot.xxx之类的方法，是因为 Vite 帮我们在模块最顶层注入了import.meta.hot对象，而这个对象由createHotContext来实现，具体的注入代码如下所示:

import { createHotContext as __vite__createHotContext } from "/@vite/client";
import.meta.hot = __vite__createHotContext("/src/App.tsx");

下面我们在 import 分析插件中做一些改动，实现插入这段代码的功能:

// /src/node/plugins/importAnalysis.ts
import path from 'path';
import resolve from 'resolve';
import { Plugin } from '../plugin';
import { pathExists } from 'fs-extra';
import MagicString from 'magic-string';
import { init, parse } from 'es-module-lexer';
import { ServerContext } from '../server/index';
import {
  BARE_IMPORT_RE,
  PRE_BUNDLE_DIR,
  DEFAULT_EXTERNALS,
  CLIENT_PUBLIC_PATH
} from '../constants';
import { cleanUrl, getShortName, isInternalRequest, isJSRequest, normalizePath } from '../utils';

export function importAnalysisPlugin(): Plugin {
  let serverContext: ServerContext;
  return {
    name: 'm-vite:import-analysis',
    configureServer(s) {
      serverContext = s;
    },
    async transform(code: string, id: string) {
      if (!isJSRequest(id) || isInternalRequest(id)) {
        return null;
      }
      await init;
      const importedModules = new Set<string>();
      const [imports] = parse(code);
      const ms = new MagicString(code);

      const resolve = async (id: string, importer?: string) => {
        const resolved = await serverContext.pluginContainer.resolveId(
          id,
          normalizePath(importer)
        );
        if (!resolved) {
          return;
        }
        let resolvedId = `/${getShortName(resolved.id, serverContext.root)}`;
        return resolvedId;
      };

      const { moduleGraph } = serverContext;
      const curMod = moduleGraph.getModuleById(id)!;

      for (const importInfo of imports) {
        const { s: modStart, e: modEnd, n: modSource } = importInfo;
        if (!modSource) continue;

        if (modSource.endsWith('.svg')) {
          const resolvedUrl = await resolve(modSource, id);
          ms.overwrite(modStart, modEnd, `${resolvedUrl}?import`);
          continue;
        }

        if (BARE_IMPORT_RE.test(modSource)) {
          const bundlePath = normalizePath(
            path.join('/', PRE_BUNDLE_DIR, `${modSource}.js`)
          );
          ms.overwrite(modStart, modEnd, bundlePath);
          importedModules.add(bundlePath);
        } else if (modSource.startsWith('.') || modSource.startsWith('/')) {
          const resolved = await resolve(modSource, id);
          if (resolved) {
            ms.overwrite(modStart, modEnd, resolved);
            importedModules.add(resolved);
          }
        }
      }

      if (!id.includes('node_modules')) {
        // 注入 HMR 相关的工具函数
        ms.prepend(
          `import { createHotContext as __vite__createHotContext } from "${CLIENT_PUBLIC_PATH}";` +
            `import.meta.hot = __vite__createHotContext(${JSON.stringify(
              cleanUrl(curMod.url)
            )});`
        );
      }

      moduleGraph.updateModuleInfo(curMod, importedModules);

      return {
        code: ms.toString(),
        map: ms.generateMap()
      };
    }
  };
}

接着在 /src/node/utils.ts 补充工具函数:

import os from 'os';
import path from 'path';
import { QUERY_RE, HASH_RE, JS_TYPES_RE, CLIENT_PUBLIC_PATH } from './constants';

export function slash(p: string): string {
  return p.replace(/\\/g, '/');
}

export const isWindows = os.platform() === 'win32';
export const INTERNAL_LIST = [CLIENT_PUBLIC_PATH, "/@react-refresh"];

export function normalizePath(id: string): string {
  return path.posix.normalize(isWindows ? slash(id) : id);
}

export const isJSRequest = (id: string): boolean => {
  id = cleanUrl(id);
  if (JS_TYPES_RE.test(id)) {
    return true;
  }
  if (!path.extname(id) && !id.endsWith('/')) {
    return true;
  }
  return false;
};

export const cleanUrl = (url: string): string =>
  url.replace(HASH_RE, '').replace(QUERY_RE, '');

export const isCSSRequest = (id: string): boolean =>
  cleanUrl(id).endsWith('.css');

export function isImportRequest(url: string): boolean {
  return url.endsWith('?import');
}

export function removeImportQuery(url: string): string {
  return url.replace(/\?import$/, '');
}

export function getShortName(file: string, root: string) {
  return file.startsWith(root + '/') ? path.posix.relative(root, file) : file;
}

export function isInternalRequest(url: string): boolean {
  return INTERNAL_LIST.includes(url);
}

接着在 /src/node/constants.ts 补充常量:

export const CLIENT_PUBLIC_PATH = "/@vite/client";

接着启动 test，打开页面后你可以发现 import.meta.hot 的实现代码已经被成功插入:

Preview

现在，我们回到客户端脚本的实现中，来开发createHotContext 这个工具方法:

interface HotModule {
  id: string;
  callbacks: HotCallback[];
}

interface HotCallback {
  deps: string[];
  fn: (modules: object[]) => void;
}

// HMR 模块表
const hotModulesMap = new Map<string, HotModule>();
// 不在生效的模块表
const pruneMap = new Map<string, (data: any) => void | Promise<void>>();

export const createHotContext = (ownerPath: string) => {
  const mod = hotModulesMap.get(ownerPath);
  if (mod) {
    mod.callbacks = [];
  }

  function acceptDeps(deps: string[], callback: any) {
    const mod: HotModule = hotModulesMap.get(ownerPath) || {
      id: ownerPath,
      callbacks: [],
    };
    // callbacks 属性存放 accept 的依赖、依赖改动后对应的回调逻辑
    mod.callbacks.push({
      deps,
      fn: callback,
    });
    hotModulesMap.set(ownerPath, mod);
  }

  return {
    accept(deps: any, callback?: any) {
      // 这里仅考虑接受自身模块更新的情况
      // import.meta.hot.accept()
      if (typeof deps === "function" || !deps) {
        acceptDeps([ownerPath], ([mod]) => deps && deps(mod));
      }
    },
    // 模块不再生效的回调
    // import.meta.hot.prune(() => {})
    prune(cb: (data: any) => void) {
      pruneMap.set(ownerPath, cb);
    },
  };
};

在 accept 方法中，我们会用hotModulesMap这张表记录该模块所 accept 的模块，以及 accept 的模块更新之后回调逻辑。

接着，我们来开发客户端热更新的具体逻辑，也就是服务端传递更新内容之后客户端如何来派发更新。实现代码如下:

async function fetchUpdate({ path, timestamp }: Update) {
  const mod = hotModulesMap.get(path);
  if (!mod) return;

  const moduleMap = new Map();
  const modulesToUpdate = new Set<string>();
  modulesToUpdate.add(path);

  await Promise.all(
    Array.from(modulesToUpdate).map(async (dep) => {
      const [path, query] = dep.split(`?`);
      try {
        // 通过动态 import 拉取最新模块
        const newMod = await import(
          path + `?t=${timestamp}${query ? `&${query}` : ""}`
        );
        moduleMap.set(dep, newMod);
      } catch (e) {}
    })
  );

  return () => {
    // 拉取最新模块后执行更新回调
    for (const { deps, fn } of mod.callbacks) {
      fn(deps.map((dep: any) => moduleMap.get(dep)));
    }
    console.log(`[vite] hot updated: ${path}`);
  };
}

现在，我们可以来初步测试一下 HMR 的功能，你可以暂时将 main.tsx 的内容换成下面这样, 注意: 先把 App 组件放到 main.ts 中, 因为只实现了修改自身刷新:

import React from "react";
import ReactDOM from "react-dom";
import "./index.css";

const App = () => <div>hello 123123</div>;

ReactDOM.render(<App />, document.getElementById("root"));

// @ts-ignore
import.meta.hot.accept(() => {
  ReactDOM.render(<App />, document.getElementById("root"));
});

启动 test，然后打开浏览器，可以看到如下的文本:

Preview

现在回到编辑器中，修改文本内容，然后保存，你可以发现页面内容也跟着发生了变化，并且网络面板发出了拉取最新模块的请求，说明 HMR 已经成功生效:

Preview

同时，当你再次刷新页面，看到的仍然是最新的页面内容。这一点非常重要，之所以能达到这样的效果，是因为我们在文件改动后会调用 ModuleGraph 的 invalidateModule 方法，这个方法会清除热更模块以及所有上层引用方模块的编译缓存:

// 方法实现
invalidateModule(file: string) {
  const mod = this.idToModuleMap.get(file);
  if (mod) {
    mod.lastHMRTimestamp = Date.now();
    mod.transformResult = null;
    mod.importers.forEach((importer) => {
      this.invalidateModule(importer.id!);
    });
  }
}

这样每次经过 HMR 后，再次刷新页面，渲染出来的一定是最新的模块内容。

当然，我们也可以对 CSS 实现热更新功能，在客户端脚本中添加如下的工具函数:

// /src/client/client.ts
const sheetsMap = new Map();

export function updateStyle(id: string, content: string) {
  let style = sheetsMap.get(id);
  if (!style) {
    // 添加 style 标签
    style = document.createElement("style");
    style.setAttribute("type", "text/css");
    style.innerHTML = content;
    document.head.appendChild(style);
  } else {
    // 更新 style 标签内容
    style.innerHTML = content;
  }
  sheetsMap.set(id, style);
}

export function removeStyle(id: string): void {
  const style = sheetsMap.get(id);
  if (style) {
    document.head.removeChild(style);
  }
  sheetsMap.delete(id);
}

紧接着我们调整一下 CSS 编译插件的代码:

// /src/node/plugins/css.ts
import { Plugin } from '../plugin';
import { readFile } from 'fs-extra';
import { getShortName } from '../utils';
import { ServerContext } from '../server';
import { CLIENT_PUBLIC_PATH } from '../constants';

export function cssPlugin(): Plugin {
  let serverContext: ServerContext;
  return {
    name: 'm-vite:css',
    configureServer(s) {
      serverContext = s;
    },
    load(id) {
      if (id.endsWith('.css')) {
        return readFile(id, 'utf-8');
      }
    },
    async transform(code, id) {
      if (id.endsWith('.css')) {
        // 包装成 JS 模块
        const jsContent = `
import { createHotContext as __vite__createHotContext } from "${CLIENT_PUBLIC_PATH}";
import.meta.hot = __vite__createHotContext("/${getShortName(
          id,
          serverContext.root
        )}");

import { updateStyle, removeStyle } from "${CLIENT_PUBLIC_PATH}"
  
const id = '${id}';
const css = '${code.replace(/\n/g, '')}';

updateStyle(id, css);
import.meta.hot.accept();
export default css;
import.meta.hot.prune(() => removeStyle(id));`.trim();
        return {
          code: jsContent
        };
      }
      return null;
    }
  };
}

最后，你可以重启 test 项目，本地尝试修改 CSS 代码，可以看到类似如下的热更新效果:

Preview