feat(preset-umi): support to preload route chunk files (umijs#12095)

* feat: add bundler info when enable okam

* feat: support to preload route chunk files

* refactor: output scp to templates

* refactor: prefetch to preload and extract as utils

* refactor: preloadMode to preload.mode

* refactor: correct visit logic

* refactor: update enableBy for okam

* refactor: compatible with mako bundler

* refactor: enable by default

* feat: support runtime public path

* refactor: correct logic

* refactor: handle no file route

* refactor: improvements

* refactor: improvements

* refactor: sort files data

packages/preset-umi/.fatherrc.ts

@@ -2,4 +2,18 @@ import { defineConfig } from 'father';
 
 export default defineConfig({
   extends: '../../.fatherrc.base.ts',
+  cjs: {
+    ignores: ['src/client/*'],
+  },
+  umd: {
+    entry: 'src/client/preloadRouteFilesScp.ts',
+    output: 'templates/routePreloadOnLoad',
+    chainWebpack(memo) {
+      memo.output.filename('preloadRouteFilesScp.js');
+      memo.output.delete('libraryTarget');
+      memo.output.iife(true);
+
+      return memo;
+    },
+  },
 });

packages/preset-umi/src/client/preloadRouteFilesScp.ts

@@ -0,0 +1,49 @@
+/**
+ * NOTE: DO NOT USE ADVANCED SYNTAX IN THIS FILE, TO AVOID INSERT HELPERS TO REDUCE SCRIPT SIZE.
+ */
+
+import { getPreloadRouteFiles } from '../features/routePreloadOnLoad/utils';
+
+// always add trailing slash for base
+const basename = '{{basename}}'.replace(/([^/])$/, '$1/');
+const publicPath = '{{publicPath}}';
+const pathname = location.pathname;
+const routePath =
+  pathname.startsWith(basename) &&
+  decodeURI(`/${pathname.slice(basename.length)}`);
+
+// skip preload if basename not match
+if (routePath) {
+  const map = '{{routeChunkFilesMap}}' as any;
+  const doc = document;
+  const head = doc.head;
+  const createElement = doc.createElement.bind(doc);
+  const files = getPreloadRouteFiles(routePath, map, {
+    publicPath,
+  });
+
+  files?.forEach((file) => {
+    const type = file.type;
+    const url = file.url;
+    let tag: HTMLLinkElement | HTMLScriptElement;
+
+    if (type === 'js') {
+      tag = createElement('script');
+      tag.src = url;
+      tag.async = true;
+    } else if (type === 'css') {
+      tag = createElement('link');
+      tag.href = url;
+      tag.rel = 'preload';
+      tag.as = 'style';
+    } else {
+      return;
+    }
+
+    file.attrs.forEach((attr) => {
+      tag.setAttribute(attr[0], attr[1] || '');
+    });
+
+    head.appendChild(tag);
+  });
+}

packages/preset-umi/src/features/okam/okam.ts

@@ -2,11 +2,19 @@ import { checkVersion } from '@umijs/utils';
 import { IApi } from '../../types';
 
 export default (api: IApi) => {
+  api.describe({
+    enableBy: () => Boolean(process.env.OKAM),
+  });
+
   api.onCheck(() => {
     // mako 仅支持 node 16+
     // ref: https://github.com/umijs/mako/issues/300
-    if (api.userConfig.mako) {
-      checkVersion(16, `Node 16 is required when using mako.`);
-    }
+    checkVersion(16, `Node 16 is required when using mako.`);
+  });
+
+  api.modifyAppData((memo) => {
+    memo.bundler = 'mako';
+
+    return memo;
   });
 };

packages/preset-umi/src/features/routePreloadOnLoad/routePreloadOnLoad.ts

@@ -0,0 +1,271 @@
+import type { StatsCompilation } from '@umijs/bundler-webpack/compiled/webpack';
+import { lodash, logger, winPath } from '@umijs/utils';
+import { readFileSync } from 'fs';
+import { dirname, isAbsolute, join, relative } from 'path';
+import { TEMPLATES_DIR } from '../../constants';
+import { createResolver } from '../../libs/scan';
+import type { IApi, IRoute } from '../../types';
+import { PRELOAD_ROUTE_MAP_SCP_TYPE } from './utils';
+
+export interface IRouteChunkFilesMap {
+  /**
+   * script attr prefix (package.json name)
+   */
+  p: string;
+  /**
+   * bundler type
+   */
+  b: string;
+  /**
+   * all chunk files
+   */
+  f: [string, string | number][];
+  /**
+   * route files index map
+   */
+  r: Record<string, number[]>;
+}
+
+/**
+ * forked from: https://github.com/remix-run/react-router/blob/fb0f1f94778f4762989930db209e6a111504aa63/packages/router/utils.ts#L688C1-L719
+ */
+const routeScoreCache = new Map<string, number>();
+
+function computeRouteScore(path: string): number {
+  if (!routeScoreCache.get(path)) {
+    const paramRe = /^:[\w-]+$/;
+    const dynamicSegmentValue = 3;
+    const emptySegmentValue = 1;
+    const staticSegmentValue = 10;
+    const splatPenalty = -2;
+    const isSplat = (s: string) => s === '*';
+    let segments = path.split('/');
+    let initialScore = segments.length;
+    if (segments.some(isSplat)) {
+      initialScore += splatPenalty;
+    }
+
+    routeScoreCache.set(
+      path,
+      segments
+        .filter((s) => !isSplat(s))
+        .reduce(
+          (score, segment) =>
+            score +
+            (paramRe.test(segment)
+              ? dynamicSegmentValue
+              : segment === ''
+              ? emptySegmentValue
+              : staticSegmentValue),
+          initialScore,
+        ),
+    );
+  }
+
+  return routeScoreCache.get(path)!;
+}
+
+export default (api: IApi) => {
+  let routeChunkFilesMap: IRouteChunkFilesMap;
+
+  // enable when package name available
+  // because preload script use package name as attribute prefix value
+  api.describe({
+    enableBy: () => Boolean(api.pkg.name),
+  });
+
+  api.addHTMLHeadScripts(() => {
+    if (api.name === 'build') {
+      // internal tern app use map mode
+      return api.config.tern
+        ? // map mode
+          [
+            {
+              type: PRELOAD_ROUTE_MAP_SCP_TYPE,
+              content: JSON.stringify(routeChunkFilesMap),
+            },
+          ]
+        : // script mode
+          [
+            {
+              content: readFileSync(
+                join(
+                  TEMPLATES_DIR,
+                  'routePreloadOnLoad/preloadRouteFilesScp.js',
+                ),
+                'utf-8',
+              )
+                .replace(
+                  '"{{routeChunkFilesMap}}"',
+                  JSON.stringify(routeChunkFilesMap),
+                )
+                .replace('{{basename}}', api.config.base)
+                .replace(
+                  '"{{publicPath}}"',
+                  `${
+                    // handle runtimePublicPath
+                    api.config.runtimePublicPath ? 'window.publicPath||' : ''
+                  }"${api.config.publicPath}"`,
+                ),
+            },
+          ];
+    }
+
+    return [];
+  });
+
+  api.onBuildComplete(async ({ err, stats }) => {
+    if (!err && !stats.hasErrors()) {
+      const routeModulePath = join(api.paths.absTmpPath, 'core/route.tsx');
+      const routeModuleName = winPath(relative(api.cwd, routeModulePath));
+      const resolver = createResolver({ alias: api.config.alias });
+      const { chunks = [] } = stats.toJson
+        ? // webpack
+          stats.toJson()
+        : // mako
+          (stats.compilation as unknown as StatsCompilation);
+
+      // collect all chunk files and file chunks indexes
+      const chunkFiles: Record<string, { index: number; id: string | number }> =
+        {};
+      const fileChunksMap: Record<
+        string,
+        { files: string[]; indexes?: number[] }
+      > = {};
+      const pickPreloadFiles = (files: string[]) =>
+        files.filter((f) => f.endsWith('.js') || f.endsWith('.css'));
+
+      for (const chunk of chunks) {
+        const routeOrigins = chunk.origins!.filter((origin) =>
+          origin.moduleName?.endsWith(routeModuleName),
+        );
+
+        for (const origin of routeOrigins) {
+          const queue = [chunk.id!].concat(chunk.siblings!);
+          const visited: typeof queue = [];
+          const files: string[] = [];
+          let fileAbsPath: string;
+
+          // resolve route file path
+          try {
+            fileAbsPath = await resolver.resolve(
+              dirname(routeModulePath),
+              origin.request!,
+            );
+          } catch (err) {
+            logger.error(
+              `[routePreloadOnLoad]: route file resolve error, cannot preload for ${origin.request!}`,
+            );
+            continue;
+          }
+
+          // collect all related chunk files for route file
+          while (queue.length) {
+            const currentId = queue.shift()!;
+
+            if (!visited.includes(currentId)) {
+              const currentChunk = chunks.find((c) => c.id === currentId)!;
+
+              // skip sibling entry chunk
+              if (currentChunk.entry) continue;
+
+              // merge files
+              pickPreloadFiles(chunk.files!).forEach((file) => {
+                chunkFiles[file] ??= {
+                  index: Object.keys(chunkFiles).length,
+                  id: currentId,
+                };
+              });
+
+              // merge files
+              files.push(...pickPreloadFiles(currentChunk.files!));
+
+              // continue to search sibling chunks
+              queue.push(...currentChunk.siblings!);
+
+              // mark as visited
+              visited.push(currentId);
+            }
+          }
+
+          fileChunksMap[fileAbsPath] = { files };
+        }
+      }
+
+      // generate indexes for file chunks
+      Object.values(fileChunksMap).forEach((item) => {
+        item.indexes = item.files.map((f) => chunkFiles[f].index);
+      });
+
+      // generate map for path -> files (include parent route files)
+      const routeFilesMap: Record<string, number[]> = {};
+
+      for (const route of Object.values<IRoute>(api.appData.routes)) {
+        // skip redirect route
+        if (!route.file) continue;
+
+        let current = route;
+        const files: string[] = [];
+
+        do {
+          // skip inline function route file
+          if (current.file && !current.file.startsWith('(')) {
+            try {
+              const fileReqPath =
+                isAbsolute(current.file) || current.file.startsWith('@/')
+                  ? current.file
+                  : // a => ./a
+                    // .a => ./.a
+                    current.file.replace(/^([^.]|\.[^./])/, './$1');
+              const fileAbsPath = await resolver.resolve(
+                api.paths.absPagesPath,
+                fileReqPath,
+              );
+
+              files.push(fileAbsPath);
+            } catch {
+              logger.error(
+                `[routePreloadOnLoad]: route file resolve error, cannot preload for ${current.file}`,
+              );
+            }
+          }
+          current = current.parentId && api.appData.routes[current.parentId];
+        } while (current);
+
+        const indexes = files.reduce<number[]>((indexes, file) => {
+          // why fileChunksMap[file] may not existing?
+          // because Mako will merge minimal async chunk into entry chunk
+          // so the merged route chunk does not has to preload
+          return indexes.concat(fileChunksMap[file]?.indexes || []);
+        }, []);
+        const { absPath } = route;
+
+        routeFilesMap[absPath] =
+          // why different route may has same absPath?
+          // because umi implement route.wrappers via nested routes way, the wrapper route will has same absPath with the nested route
+          // so we always select the longest file indexes for the nested route
+          !routeFilesMap[absPath] ||
+          routeFilesMap[absPath].length < indexes.length
+            ? indexes
+            : routeFilesMap[absPath];
+      }
+
+      routeChunkFilesMap = {
+        p: api.pkg.name!,
+        b: api.appData.bundler!,
+        f: Object.entries(chunkFiles)
+          .sort((a, b) => a[1].index - b[1].index)
+          .map(([k, { id }]) => [k, id]),
+        // sort similar to react-router@6
+        r: lodash(routeFilesMap)
+          .toPairs()
+          .sort(
+            ([a]: [string, number[]], [b]: [string, number[]]) =>
+              computeRouteScore(a) - computeRouteScore(b),
+          )
+          .fromPairs()
+          .value() as any,
+      };
+    }
+  });
+};