Exploration of building a module graph 'quickly' using esbuild and enhanced-resolve

graphBuilder.ts

import {
  CachedInputFileSystem,
  FileSystem,
  ResolveContext,
  ResolverFactory,
} from 'enhanced-resolve';
import { build, ImportKind, Loader } from 'esbuild';
import * as Fs from 'fs';
import Module from 'module';
import * as Path from 'path';
import { invariant } from '../invariant';

interface ResolvedEdge {
  fromId: string;
  toId: string;
  kind: ImportKind;
}

interface UnresolvedEdge {
  fromId: string;
  fromContext: string;
  toSpec: string;
  kind: ImportKind;
}

interface File {
  id: string;
  content: string;
}

export interface BuildGraphOptions {
  conditionNames?: string[];
  define?: Record<string, string>;
  fs?: FileSystem;
  ignoreModules?: string[];
  resolveExtensions?: string[];
  rootDir?: string;
}

async function buildGraph(
  entrypoint: string[],
  options: BuildGraphOptions = {}
) {
  const fileSystem = new CachedInputFileSystem(options.fs ?? Fs, 4000);
  const resolver = ResolverFactory.createResolver({
    pnpApi: null,
    // TODO: Configurable condition names
    conditionNames: options.conditionNames ?? ['import', 'require', 'default'],
    useSyncFileSystemCalls: false,
    fileSystem,
    // TODO: Configurable file extensions
    extensions: options.resolveExtensions ?? [
      '.js',
      '.jsx',
      '.json',
      '.ts',
      '.tsx',
    ],
  });
  const rootDir = options.rootDir ?? process.cwd();
  /** Dependencies that we will ignore */
  const ignoreEdge = new Set(Module.builtinModules);

  /** Queue of unresolved edges needing to be resolved */
  const unresolvedEdgeQueue: UnresolvedEdge[] = [];
  /** Edges that have already been recognized */
  // TODO: Is it worth trying to dedupe edges?
  // const seenEdge = new MapMapSet<string, ImportKind, string>();

  /** Queue of files needing to be parsed */
  const unparsedFileQueue: string[] = [];
  /** Files that have been seen */
  const seenFiles = new Set<string>();

  const edges = new Set<ResolvedEdge>();
  const files = new Map<string, File>();

  const readFile = (fileName: string): Promise<string> => {
    return new Promise<string>((resolve, reject) => {
      return fileSystem.readFile(
        fileName,
        { encoding: 'utf8' },
        (err, result) => {
          if (err) {
            return reject(err);
          }

          return resolve(result as string);
        }
      );
    });
  };

  const resolve = (
    spec: string,
    fromId: string
  ): Promise<{ resolved: string | false | undefined; ctx: ResolveContext }> => {
    const ctx: ResolveContext = {};

    return new Promise((resolve, reject) =>
      resolver.resolve({}, fromId, spec, ctx, (err, resolved) => {
        if (err) {
          return reject(err);
        }

        return resolve({ ctx, resolved });
      })
    );
  };

  // We can't let this throw because nothing will register a catch handler.
  const parseFile = async (fileName: string): Promise<void> => {
    try {
      const sourceContent = await readFile(fileName);
      const loader = loaderForPath(fileName);

      if (!loader) {
        // We can't figure out a loader so let's just treat it as a leaf node
        files.set(fileName, { content: sourceContent, id: fileName });
        return;
      }

      const buildResult = await build({
        bundle: true,
        define: {
          // TODO: Dynamic env vars
          'proess.env.NODE_ENV': JSON.stringify('development'),
        },
        format: 'esm',
        metafile: true,
        platform: 'neutral',
        plugins: [
          {
            name: 'capture-edges',
            setup: (build) => {
              build.onResolve({ filter: /.*/ }, (args) => {
                if (!ignoreEdge.has(args.path)) {
                  unresolvedEdgeQueue.push({
                    fromId: fileName,
                    fromContext: args.resolveDir,
                    kind: args.kind,
                    toSpec: args.path,
                  });
                }

                // Mark everythign as external. We're only using esbuild to transform
                // on a file-by-file basis and capture dependencies.
                return { external: true };
              });
            },
          },
        ],
        // sourcemap: true,
        // sourcesContent: true,
        stdin: {
          contents: sourceContent,
          resolveDir: Path.dirname(fileName),
          loader: loaderForPath(fileName),
          sourcefile: fileName,
        },
        // TODO: Dynamic target
        target: 'node14',
        treeShaking: true,
        write: false,
      });

      const content = buildResult.outputFiles[0].text;

      files.set(fileName, { content, id: fileName });
    } catch (err) {
      console.error('parseFile error', err);
    }
  };

  // We can't let this throw because nothing will register a catch handler.
  const resolveEdge = async (edge: UnresolvedEdge): Promise<void> => {
    try {
      const { resolved } = await resolve(edge.toSpec, edge.fromContext);

      // TODO: We need special handling for `false` files and proper error handling
      // for files that failed to resolve.
      invariant(resolved, 'All files must successfully resolve (for now)');

      // Record the resolved edge.
      // TODO: Make sure we don't record the same logical edge twice.
      edges.add({
        fromId: edge.fromId,
        toId: resolved,
        kind: edge.kind,
      });
      unparsedFileQueue.push(resolved);
    } catch (err) {
      console.error('resolveEdge error', err);
    }
  };

  const promises = new Set<Promise<unknown>>();
  const track = (op: Promise<unknown>) => {
    promises.add(op);
    op.finally(() => promises.delete(op));
  };

  for (const entrypointSpec of entrypoint) {
    unresolvedEdgeQueue.push({
      fromId: '<root>',
      fromContext: rootDir,
      toSpec: entrypointSpec,
      kind: 'entry-point',
    });
  }

  while (unparsedFileQueue.length || unresolvedEdgeQueue.length) {
    while (unparsedFileQueue.length) {
      const unparsedFile = unparsedFileQueue.shift()!;

      if (!seenFiles.has(unparsedFile)) {
        seenFiles.add(unparsedFile);

        track(parseFile(unparsedFile));
      }
    }

    while (unresolvedEdgeQueue.length) {
      const unresolvedEdge = unresolvedEdgeQueue.shift()!;

      track(resolveEdge(unresolvedEdge));
    }

    while (promises.size) {
      await Promise.race(promises);
    }
  }

  return { edges, files };
}

function loaderForPath(fileName: string): Loader | undefined {
  const ext = Path.extname(fileName).slice(1);
  switch (ext) {
    case 'js':
    case 'mjs':
    case 'cjs':
      return 'js';
    case 'jsx':
    case 'ts':
    case 'tsx':
    case 'json':
    case 'css':
      return ext;
  }
}

// class MapMapSet<K, K2, V> {
//   private items = new Map<K, Map<K2, Set<V>>>();

//   add(key: K, subkey: K2, value: V) {
//     let itemValues = this.items.get(key);

//     if (!itemValues) {
//       itemValues = new Map();
//       this.items.set(key, itemValues);
//     }

//     let subkeyValues = itemValues.get(subkey);

//     if (!subkeyValues) {
//       subkeyValues = new Set();
//       itemValues.set(subkey, subkeyValues);
//     }

//     subkeyValues.add(value);

//     return this;
//   }

//   has(key: K, subkey: K2, value: V): boolean {
//     return this.items.get(key)?.get(subkey)?.has(value) === true;
//   }
// }

if (require.main === module)
  (async () => {
    console.time('buildGraph');
    const graph = await buildGraph(['./src/dev'], {});
    console.timeEnd('buildGraph');

    console.log(
      'Found %d modules with %d edges',
      graph.files.size,
      graph.edges.size
    );
  })();