oxc-project__oxc/packages/codegen/tsdown_plugins/unmap_writes.ts

import { parseSync, Visitor } from "oxc-parser";

import type { Plugin } from "rolldown";

/**
 * Plugin to turn every mapped write back into a plain one, for builds without source maps.
 *
 * ```ts
 * // Original code
 * writeWithMap(state, "declare ", CAT_OTHER, node.start, node.end, node);
 *
 * // After transform
 * write(state, "declare ", CAT_OTHER);
 * ```
 *
 * The mapped writes and `markMap*` exist to record a source mapping at the offsets they are given.
 * A build without source map support has nothing to record, so the call becomes the plain write
 * it would otherwise be, and the offsets and the node go with it - they would still be read,
 * for a function which ignores them.
 *
 * `remove` is how many trailing arguments are dropped from the call.
 *
 * The import is rewritten to match, which both keeps the plain name in scope and leaves the mapped functions
 * unreferenced for the minifier to remove.
 *
 * `printString` and `printNonNegativeFloat` take the offsets and the node only to hand them on to their
 * mapped writes. So they are rewritten the same way, but keep their names - the arguments come off every call.
 * The parameters also come off their declarations which, unlike the mapped writes, survive into the build.
 *
 * Only valid where `SOURCEMAPS` is `false`, which is where the config includes it.
 *
 * The checks afterwards fail the build rather than let a mismatch through. A rewrite which dropped
 * the wrong argument, or left a call to a name nothing imports, produces a printer that is broken
 * everywhere or subtly wrong everywhere, and nothing downstream names the plugin as the cause.
 */
const REWRITES = {
  // Every mapped write ends `..., start, end, node`, so 3 arguments come off.
  // `write` takes the `last` category before them, `writeNoLast` does not.
  writeWithMap: { arity: 6, remove: 3, rename: "write" },
  writeWithMapNamed: { arity: 5, remove: 3, rename: "writeIdent" },
  writeWithMapNoLast: { arity: 5, remove: 3, rename: "writeNoLast" },
  writeWithMapNamedNoLast: { arity: 5, remove: 3, rename: "writeNoLast" },
  writeWithMapNamedJSXNoLast: { arity: 5, remove: 3, rename: "writeNoLast" },
  // A private identifier writes its own `#`, so its plain form is not `write`
  writeWithMapNamedPrivate: { arity: 5, remove: 3, rename: "writePrivate" },
  writeWithMapEnd: { arity: 6, remove: 3, rename: "write" },
  // `rename: null` because a standalone mark has no non-sourcemap equivalent.
  // Everything but `state` comes off - `markMapAtStartOffset`'s column offset included.
  markMapStart: { arity: 4, remove: 3, rename: null },
  markMapAfter: { arity: 4, remove: 3, rename: null },
  markMapAtStartOffset: { arity: 5, remove: 4, rename: null },
  markMapEnd: { arity: 4, remove: 3, rename: null },
  markMapNamed: { arity: 7, remove: 6, rename: null },
  // `rename: null` to transform the function declarations, removing the dropped params
  printString: { arity: 5, remove: 3, rename: null },
  printNonNegativeFloat: { arity: 5, remove: 3, rename: null },
} as const;

const WRITE_MODULE = "./write.ts";

type MappedName = keyof typeof REWRITES;

/**
 * Create the plugin for one build flavour.
 *
 * @param sourcemaps - `true` to strip only `node`, `false` to unmap the writes entirely
 * @param debug - `true` if is a debug build
 * @returns The plugin
 */
export default function unmapWritesPlugin(sourcemaps: boolean, debug: boolean): Plugin {
  return {
    name: "unmap-writes",
    transform: {
      // Only process TS files in `src-js/print` directory
      filter: { id: /\/src-js\/print\/.+\.ts$/ },

      handler(code, path, meta) {
        // Don't alter anything in debug sourcemap builds.
        // No functions are renamed and no arguments to those functions are removed.
        if (sourcemaps && debug) return;

        // Parse file
        const magicString = meta.magicString!;
        const { program, errors } = parseSync(path, code);
        if (errors.length !== 0) throw new Error(`Failed to parse ${path}: ${errors[0].message}`);

        // Imports from `write.ts` which have been replaced, and which need to be replaced.
        // The 2 are compared and ensured equal at the end of the visitation pass.
        const importReplacements: string[] = [];
        const neededImportReplacements = new Set<string>();
        let importIsTransformed = false;

        new Visitor({
          // Remove the dropped args from calls, and rename callees where required
          CallExpression(node) {
            const { callee } = node;
            if (callee.type !== "Identifier") return;

            const name = callee.name as MappedName;
            const rewrite = REWRITES[name];
            if (rewrite === undefined) return;

            const { arity, rename } = rewrite;
            const args = node.arguments;
            if (args.length !== arity) {
              throw new Error(
                `\`${callee.name}\` takes ${arity} arguments, found ${args.length}: `
                  + `${path}:${node.start}`,
              );
            }

            // Remove the dropped arguments, from the end of the last kept one to the end of the last.
            // In sourcemaps builds, remove only the last argument (`node`).
            const remove = sourcemaps ? 1 : rewrite.remove;
            magicString.remove(args[arity - remove - 1].end, args[arity - 1].end);

            // If callee needs to be renamed, rename it, and record that it needs to be imported under new name.
            // Nothing is renamed in sourcemap builds.
            if (!sourcemaps && rename !== null) {
              magicString.overwrite(callee.start, callee.end, rename);
              neededImportReplacements.add(name);
            }
          },

          // Remove the dropped params from declarations
          FunctionDeclaration(node) {
            const { id } = node;
            if (id === null) return;

            const name = id.name as MappedName;
            const rewrite = REWRITES[name];
            if (rewrite === undefined) return;

            const { arity } = rewrite;
            const { params } = node;
            if (params.length !== arity) {
              throw new Error(
                `\`${id.name}\` defined with ${params.length} parameters, expected ${arity}: `
                  + `${path}:${node.start}`,
              );
            }

            // Remove the dropped params, from the end of the last kept one to the end of the last.
            // In sourcemaps builds, remove only the last param (`node`).
            const remove = sourcemaps ? 1 : rewrite.remove;
            magicString.remove(params[arity - remove - 1].end, params[arity - 1].end);
          },

          // Replace import of functions whose calls are renamed or removed
          ImportDeclaration(node) {
            // Skip imports not from `write.ts`, and type imports
            if (node.source.value !== WRITE_MODULE || node.importKind === "type") return;

            // For simplicity, we only support a single import from `write.ts` (not including type imports)
            if (importIsTransformed) {
              throw new Error(`Multiple imports from \`${WRITE_MODULE}\`: ${path}`);
            }
            importIsTransformed = true;

            // Collect all specifiers, replacing any that need to be
            const specifierNames = new Set<string>();
            for (const specifier of node.specifiers) {
              if (
                specifier.type !== "ImportSpecifier"
                || specifier.importKind === "type"
                || specifier.imported.type !== "Identifier"
              ) {
                throw new Error(
                  `Only simple imports from \`${WRITE_MODULE}\` are supported: ${path}`,
                );
              }

              const name = specifier.imported.name as MappedName;

              // Only rename imports in no-sourcemap builds
              if (!sourcemaps) {
                const rewrite = REWRITES[name];
                if (rewrite !== undefined) {
                  const { rename } = rewrite;
                  if (rename !== null) {
                    specifierNames.add(rename);
                    importReplacements.push(name);
                    continue;
                  }
                }
              }

              specifierNames.add(name);
            }

            if (importReplacements.length > 0) {
              magicString.overwrite(
                node.start,
                node.end,
                `import { ${[...specifierNames].join(", ")} } from "${WRITE_MODULE}";`,
              );
            }
          },
        }).visit(program);

        // Check that replaced callees also had the corresponding import replaced
        if (importReplacements.length !== neededImportReplacements.size) {
          const missing = [];
          for (const name of neededImportReplacements) {
            if (!importReplacements.includes(name)) missing.push(name);
          }

          if (missing.length > 0) {
            throw new Error(
              `Missing import of ${missing.map((name) => `\`${name}\``).join(", ")}: ${path}`,
            );
          }
        }

        return { code: magicString };
      },
    },
  };
}