一起写个vite吧！(3) --模块依赖图的开发和毫秒级的热更新

兄弟们我回来了,上一章节一起写个vite吧！(2) --插件机制开发同学们应该消化的差不多了吧，这一节我们又要整新活了,为了比较方便地管理各个模块之间的依赖关系，vite在Dev Server中创建了模块依赖图的数据结构，模块依赖图可以用来缓存，vite的热更新也是依赖于这个模块依赖图的，首先vite的的Dev Server启动后会初始化ModuleGraph(也就是模块依赖图的类):

const moduleGraph: ModuleGraph = new ModuleGraph((url) =>
  container.resolveId(url)
);

接下来我们具体查看ModuleGraph这个类的实现。其中定义了若干个 Map，用来记录模块信息:

// 由原始请求 url 到模块节点的映射，如 /src/index.tsx
urlToModuleMap = new Map<string, ModuleNode>()
// 由模块 id 到模块节点的映射，其中 id 与原始请求 url，为经过 resolveId 钩子解析后的结果
idToModuleMap = new Map<string, ModuleNode>()
// 由文件到模块节点的映射，由于单文件可能包含多个模块，如 .vue 文件，因此 Map 的 value 值为一个集合
fileToModulesMap = new Map<string, Set<ModuleNode>>()

ModuleNode 对象即代表模块节点的具体信息，我们可以来看看它的数据结构:

class ModuleNode {
  // 原始请求 url
  url: string
  // 文件绝对路径 + query
  id: string | null = null
  // 文件绝对路径
  file: string | null = null
  type: 'js' | 'css'
  info?: ModuleInfo
  // resolveId 钩子返回结果中的元数据
  meta?: Record<string, any>
  // 该模块的引用方
  importers = new Set<ModuleNode>()
  // 该模块所依赖的模块
  importedModules = new Set<ModuleNode>()
  // 接受更新的模块
  acceptedHmrDeps = new Set<ModuleNode>()
  // 是否为`接受自身模块`的更新
  isSelfAccepting = false
  // 经过 transform 钩子后的编译结果
  transformResult: TransformResult | null = null
  // SSR 过程中经过 transform 钩子后的编译结果
  ssrTransformResult: TransformResult | null = null
  // SSR 过程中的模块信息
  ssrModule: Record<string, any> | null = null
  // 上一次热更新的时间戳
  lastHMRTimestamp = 0

  constructor(url: string) {
    this.url = url
    this.type = isDirectCSSRequest(url) ? 'css' : 'js'
  }
}

模块依赖图

OK，模块依赖图大概就是这么个东西，当然我们不会跟源码一样实现的那么全，现在我们来实现模块依赖图，我们建个src/node/ModuleGraph.ts，代码如下：

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>();
  // 经过 transform 钩子后的编译结果
  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 };
  }
}

ok,看完上面的代码后，大家应该还记得刚才说的在vite的Dev Server启动后要初始化模块依赖图吧，我们在服务端中加入代码:

src/node/server/index.ts

+ import { ModuleGraph } from "../ModuleGraph";

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

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

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

src/node/server/middlewares/transform.ts

let code = await pluginContainer.load(resolvedResult.id);
if (typeof code === "object" && code !== null) {
  code = code.code;
}
+ const { moduleGraph }  = serverContext;
+ mod = await moduleGraph.ensureEntryFromUrl(url);

然后当我们对 JS?X/ts?X 模块分析完 import 语句之后，需要更新模块之间的依赖关系，就比如我在main.tsx里引用了App.tsx，那我是不是要把它加入 main.tsx 的 importModules (该模块所依赖的模块)里啊，然后去更新模块依赖图，代码如下:

src/node/plugins/importAnalysis.ts

export function importAnalysis() {
  return {
    transform(code: string, id: string) {
      // 省略前面的代码
+      const { moduleGraph } = serverContext;
+      const curMod = moduleGraph.getModuleById(id)!;
+      const importedModules = new Set<string>();
      for(const importInfo of imports) {
        // 省略部分代码
        if (BARE_IMPORT_RE.test(modSource)) {
          // 省略部分代码
+          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 代码
    }
  }
}

现在，一个完整的模块依赖图就能随着 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;
}

ok,模块依赖图的逻辑基本上搞定了，接下来就是令人期待的 hmr 也就是vite的热更新了，大家应该很好奇vite的热更新为什么这么快吧，让我们一探究竟吧。

HMR服务端

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

• 创建文件监听器，以监听文件的变动
• 创建 WebSocket 服务端，负责和客户端进行通信
• 文件变动时，从 ModuleGraph 中定位到需要更新的模块，将更新信息发送给客户端

文件监听器是这么搞的呢，其实用的 chokidar 这个库，可以监听文件的变化，大家应该使用过 nodemon吧，是不是可以监听文件的变化然后重启服务，其实 nodemon 内部也是用的这个库实现的。

ok，那我们创建一个文件监听器

src/node/server/index.ts

import chokidar, { FSWatcher } from "chokidar";

export async function startDevServer() {
  const watcher = chokidar.watch(root, {
    ignored: ["**/node_modules/**", "**/.git/**"],
    ignoreInitial: true,
  });
}

然后我们要先初始化WebSocket 服务端，新建src/node/ws.ts，代码如下:

src/node/ws.ts

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

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();
    },
  };
}

然后我们要src/node/constants.ts在定义 HMR_PORT 常量:

export const HMR_PORT = 24678;

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

src/node/server/index.ts

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() {

+  // WebSocket 对象
+  const ws = createWebSocketServer(app);
  // // 开发服务器上下文
  const serverContext: ServerContext = {
    root: process.cwd(),
    app,
    pluginContainer,
    plugins,
    moduleGraph,
+    ws,
+    watcher
  };
}

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

src/node/hmr.ts

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

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:
            "/" +
            (isWindows
              ? getWindowShortName(file, root)
              : getShortName(file, root)),
          acceptedPath:
            "/" +
            (isWindows
              ? getWindowShortName(file, root)
              : getShortName(file, root)),
        },
      ],
    });
  });
}

然后我们补充一下工具函数：

src/node/utils.ts

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

// window系统下的路径\换成/统一路径，如C:\jjj换成C:/jjj
export function getWindowShortName(file: string, root: string) {
  return file.startsWith(root)
    ? path.posix.relative(slash(root), slash(file))
    : file;
}

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

src/node/server/index.ts

+  import { bindingHMREvents } from "../hmr";

// 开发服务器上下文
const serverContext: ServerContext = {
  root: process.cwd(),
  app,
  pluginContainer,
  plugins,
  moduleGraph,
  ws,
  watcher,
};
+ bindingHMREvents(serverContext);

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: {
    index: "src/node/cli.ts",
+    client: "src/client/client.ts",
  },
});

注: 改动 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，我们通过这个插件将脚本插入客户端，代码如下:

src/node/plugins/clientInject.ts

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

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 { clientInjectPlugin } from './clientInject';

export function resolvePlugins(): Plugin[] {
  return [
+     clientInjectPlugin()
    // 省略其它插件
  ]
}

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

接下来我们在 playground 项目下执行pnpm dev，然后查看页面，可以发现控制台出现了如下的 log 信息，那么证明我们注入客户端的脚本生效了，如果没有的话就去检查一下代码吧

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 { init, parse } from "es-module-lexer";
import {
  BARE_IMPORT_RE,
  PRE_BUNDLE_DIR,
  CLIENT_PUBLIC_PATH,
} from "../constants";
import {
  cleanUrl,
+  getShortName,
  isInternalRequest,
  isJSRequest,
  normalizePath,
} from "../utils";
// magic-string 用来作字符串编辑
import MagicString from "magic-string";
import path from "path";
import { Plugin } from "../plugin";
import { ServerContext } from "../server/index";
import type { PluginContext } from "rollup";

export function importAnalysisPlugin(): Plugin {
  let serverContext: ServerContext;
  return {
    name: "m-vite:import-analysis",
    configureServer(s) {
      // 保存服务端上下文
      serverContext = s;
    },
    async transform(this: PluginContext, code: string, id: string) {
      // 只处理 JS 相关的请求
+      if (!isJSRequest(id) || isInternalRequest(id)) {
        return null;
      }
      await init;
      const importedModules = new Set<string>();
      // 解析 import 语句
      const [imports] = parse(code);
      const ms = new MagicString(code);
      const resolve = async (id: string, importer?: string) => {
+        const resolved = await serverContext.pluginContainer.resolveId(
+          id,
+          importer
+        );

+        if (!resolved) {
+          return;
+        }
+        const cleanedId = cleanUrl(resolved.id);
+        const mod = moduleGraph.getModuleById(cleanedId);
+
+        let resolvedId = `/${getShortName(resolved.id, serverContext.root)}`;
+
+        if (mod && mod.lastHMRTimestamp > 0) {
+          // resolvedId += "?t=" + mod.lastHMRTimestamp;
+        }
+        return resolvedId;
+      };
      const { moduleGraph } = serverContext;
      const curMod = moduleGraph.getModuleById(id)!;
      // 对每一个 import 语句依次进行分析
      for (const importInfo of imports) {
        // 举例说明: const str = `import React from 'react'`
        // str.slice(s, e) => 'react'
        const { s: modStart, e: modEnd, n: modSource } = importInfo;
        if (!modSource || isInternalRequest(modSource)) continue;
        // 静态资源
        if (modSource.endsWith(".svg")) {
          // 加上 ?import 后缀

          // console.log(path.dirname(id));
          // console.log(modSource);

          const resolvedUrl = normalizePath(
            path.relative(
              path.dirname(id),
              path.resolve(path.dirname(id), modSource)
            )
          );
          // console.log(resolvedUrl);
          ms.overwrite(modStart, modEnd, `./${resolvedUrl}?import`);
          continue;
        }
        // 第三方库: 路径重写到预构建产物的路径
        if (BARE_IMPORT_RE.test(modSource)) {
          // const bundlePath = path.join(
          //   serverContext.root,
          //   PRE_BUNDLE_DIR,
          //   `${modSource}.js`
          // )
          const bundlePath = normalizePath(
            path.join("/", PRE_BUNDLE_DIR, `${modSource}.js`)
          );
          ms.overwrite(modStart, modEnd, bundlePath);
          importedModules.add(bundlePath);
        } else if (modSource.startsWith(".") || modSource.startsWith("/")) {
          // 直接调用插件上下文的 resolve 方法，会自动经过路径解析插件的处理
+          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(),
        // 生成 SourceMap
        map: ms.generateMap(),
      };
    },
  };
}

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

src/client/client.ts

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 的模块更新之后回调逻辑。

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

src/client/client.ts

let pending = false;
let queued: Promise<(() => void) | undefined>[] = [];

// 批量任务处理，不与具体的热更新行为挂钩，主要起任务调度作用
async function queueUpdate(p: Promise<(() => void) | undefined>) {
  queued.push(p);
  if (!pending) {
    pending = true;
    await Promise.resolve();
    pending = false;
    const loading = [...queued];
    queued = [];
    (await Promise.all(loading)).forEach((fn) => fn && fn());
  }
}

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 的内容换成下面这样:

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

const App = () => <h1>hello pujie</h1>;

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

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

我们启动项目会看到如下界面:

控制台打印:

ok,然后我们把pujie改成大帅B，然后保存代码我们会看到页面更新了:

同时我们在文件改动后会调用 ModuleGraph 的 invalidateModule 方法，这个方法会清除热更模块以及所有上层引用方模块的编译缓存,如果不清除缓存的话会发生什么呢，如果没清除缓存的话我们再次请求这个模块都是返回缓存里的，也就是无论你怎么改代码保存和刷新浏览器都只会返回缓存的内容，所以我们要清除热更模块以及所有上层引用方模块的编译缓存:

src/client/client.ts

// 方法实现
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 { readFile } from "fs-extra";
import { CLIENT_PUBLIC_PATH } from "../constants";
import { Plugin } from "../plugin";
import { ServerContext } from "../server";
import { getShortName, getWindowShortName, isWindows } from "../utils";

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("/${
          isWindows
            ? getWindowShortName(id, serverContext.root)
            : getShortName(id, serverContext.root)
        }");
import { updateStyle, removeStyle } from "${CLIENT_PUBLIC_PATH}"
  
const id = '${id}';
const css = '${code.replace(/\r\n/g, "")}';

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

这个逻辑比较简单就是用Map将css文件的路径存在map里，如果map里没有就是初始化，有的话就是更新，将style标签里的内容换成新的就更新了。

最后client.ts的完整内容如下:

console.log("[vite] connecting...");

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

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

interface Update {
  type: "js-update" | "css-update";
  path: string;
  acceptedPath: string;
  timestamp: number;
}

// 根据不同的更新类型进行更新
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") {
          console.log(update);

          queueUpdate(fetchUpdate(update));
        }
      });
      break;
  }
}

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

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

const hotModulesMap = new Map<string, HotModule>();
const pruneMap = new Map<string, (data: any) => void | Promise<void>>();

export const createHotContext = (ownerPath: string) => {
  console.log(ownerPath);

  const mod = hotModulesMap.get(ownerPath);
  if (mod) {
    mod.callbacks = [];
  }

  function acceptDeps(deps: string[], callback: any) {
    const mod: HotModule = hotModulesMap.get(ownerPath) || {
      id: ownerPath,
      callbacks: [],
    };
    mod.callbacks.push({
      deps,
      fn: callback,
    });
    hotModulesMap.set(ownerPath, mod);
  }

  return {
    accept(deps: any) {
      if (typeof deps === "function" || !deps) {
        acceptDeps([ownerPath], ([mod]) => deps && deps(mod));
      }
    },
    prune(cb: (data: any) => void) {
      pruneMap.set(ownerPath, cb);
    },
  };
};

let pending = false;
let queued: Promise<(() => void) | undefined>[] = [];

// 批量任务处理，不与具体的热更新行为挂钩，主要起任务调度作用
async function queueUpdate(p: Promise<(() => void) | undefined>) {
  queued.push(p);
  if (!pending) {
    pending = true;
    await Promise.resolve();
    pending = false;
    const loading = [...queued];
    queued = [];
    (await Promise.all(loading)).forEach((fn) => fn && fn());
  }
}

async function fetchUpdate({ path, timestamp }: Update) {
  const mod = hotModulesMap.get(path);
  console.log(path);
  console.log(hotModulesMap);
  console.log(mod);

  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 {
        const newMod = await import(
          path + `?t=${timestamp}${query ? `&${query}` : ""}`
        );
        console.log(newMod);

        moduleMap.set(dep, newMod);
      } catch (e) {
        console.log(e);
      }
    })
  );

  return () => {
    for (const { deps, fn } of mod.callbacks) {
      fn(deps.map((dep: any) => moduleMap.get(dep)));
    }
    console.log(`[vite] hot updated: ${path}`);
  };
}

const sheetsMap = new Map();

export function updateStyle(id: string, content: string) {
  let style = sheetsMap.get(id);
  if (!style) {
    style = document.createElement("style");
    style.setAttribute("type", "text/css");
    style.innerHTML = content;
    document.head.appendChild(style);
  } else {
    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);
}

我们直接重启项目，将修改index.css文件：

将color改成red,我们可以看到页面更新了：

ok，我们实现的mini-vite就告一段落拉，做一下总结吧，其实我也是最近才开始学习vite的原理，刚开始接触的时候确实感觉很难，然后我就想着把这个写成文章逼自己去学，这里不推荐直接去干源码哈，你很牛皮就当我没说，其实这系列文章写下来我对vite的理解深了很多，但是离吃透还是有不远的距离的，大家有问题可以找我一起讨论，我个人是非常喜欢vite的，对开发者实在是太友好了，但它并不是没有缺点，vite在请求数量达到一定量级的时候，no-bundle的服务都会碰到加载性能问题，需要等团队优化，在开发环境下调试后端接口需要频繁刷新界面，还是会影响开发体验，但是vite的构建速度和热更是真滴快，vite3现在也可以支持生产环境开启esbuild打包了，之前开发用esbuild，生产用rollup，会导致依赖的不一致问题，造成bug,在vite3只要添加optimizeDeps.disabled: false的配置就可以在生产环境下使用esbuild了。当然vite3还有其他的更新，这里就不多赘述拉，那就完结撒花咯！！！

往期文章

• 一起写个vite吧！(1) --环境搭建+依赖预构建

• 一起写个vite吧！(2) --插件机制开发

mini-vite仓库地址(里面有对应的提交记录，如果有兴趣可以拉下来跑一跑哇)