sitegen/framework/hot.ts

// This implements the ability to use TS, TSX, and more plugins
// in Node.js. It cannot be built on the ES module loader,
// because there is no exposed way to replace modules when
// needed (see nodejs#49442).
//
// It also allows using a simple compile cache, which is used by
// the site generator to determine when code changes.
export const projectRoot = path.resolve(import.meta.dirname, "../");
export const projectSrc = path.resolve(projectRoot, "src");

// Create a project-relative require. For convenience, it is generic-typed.
export const load = createRequire(
  pathToFileURL(path.join(projectRoot, "run.js")).toString(),
) as {
  <T = unknown>(id: string): T;
  extensions: NodeJS.Dict<(mod: NodeJS.Module, file: string) => unknown>;
  cache: NodeJS.Dict<NodeJS.Module>;
  resolve: (id: string, o?: { paths: string[] }) => string;
};
export const { cache } = load;

// Register extensions by overwriting `require.extensions`
const require = load;
const exts = require.extensions;
exts[".ts"] = loadEsbuild;
exts[".tsx"] = loadEsbuild;
exts[".jsx"] = loadEsbuild;
exts[".marko"] = loadMarko;
exts[".mdx"] = loadMdx;
exts[".css"] = loadCss;

// Intercept all module load calls to track CSS imports + file times.
export interface FileStat {
  cssImportsRecursive: string[] | null;
  lastModified: number;
  imports: string[];
}
const fileStats = new Map<string, FileStat>();

export function getFileStat(filepath: string) {
  return fileStats.get(path.resolve(filepath));
}

function shouldTrackPath(filename: string) {
  return !filename.includes("node_modules") &&
    !filename.includes(import.meta.dirname);
}

const Module = load<typeof import("node:module")>("node:module");
const ModulePrototypeUnderscoreCompile = Module.prototype._compile;
Module.prototype._compile = function (
  content: string,
  filename: string,
  format: "module" | "commonjs",
) {
  const result = ModulePrototypeUnderscoreCompile.call(
    this,
    content,
    filename,
    format,
  );
  const stat = fs.statSync(filename);
  if (shouldTrackPath(filename)) {
    const cssImportsMaybe: string[] = [];
    const imports: string[] = [];
    for (const { filename: file, cloverClientRefs } of this.children) {
      if (file.endsWith(".css")) cssImportsMaybe.push(file);
      else {
        const child = fileStats.get(file);
        if (!child) continue;
        const { cssImportsRecursive } = child;
        if (cssImportsRecursive) cssImportsMaybe.push(...cssImportsRecursive);
        imports.push(file);
        if (cloverClientRefs && cloverClientRefs.length > 0) {
          (this.cloverClientRefs ??= [])
            .push(...cloverClientRefs);
        }
      }
    }
    fileStats.set(filename, {
      cssImportsRecursive: cssImportsMaybe.length > 0
        ? Array.from(new Set(cssImportsMaybe))
        : null,
      imports,
      lastModified: stat.mtimeMs,
    });
  }
  return result;
};

// Implement @/ prefix
const ModuleUnderscoreResolveFilename = Module._resolveFilename;
Module._resolveFilename = (...args) => {
  if (args[0].startsWith("@/")) {
    const replacedPath = "." + args[0].slice(1);
    try {
      return require.resolve(replacedPath, { paths: [projectSrc] });
    } catch (err: any) {
      if (err.code === "MODULE_NOT_FOUND" && (err?.requireStack?.length ?? 0) <= 1) {
        err.message.replace(replacedPath, args[0]);
      }
    }
  }
  return ModuleUnderscoreResolveFilename(...args);
};

function loadEsbuild(module: NodeJS.Module, filepath: string) {
  return loadEsbuildCode(module, filepath, fs.readFileSync(filepath, "utf8"));
}

interface LoadOptions {
  scannedClientRefs?: string[];
}
function loadEsbuildCode(
  module: NodeJS.Module,
  filepath: string,
  src: string,
  opt: LoadOptions = {},
) {
  if (filepath === import.meta.filename) {
    module.exports = self;
    return;
  }
  let loader: any = "tsx";
  if (filepath.endsWith(".ts")) loader = "ts";
  else if (filepath.endsWith(".jsx")) loader = "jsx";
  else if (filepath.endsWith(".js")) loader = "js";
  if (opt.scannedClientRefs) {
    module.cloverClientRefs = opt.scannedClientRefs;
  } else {
    let { code, refs } = resolveClientRefs(src, filepath);
    module.cloverClientRefs = refs;
    src = code;
  }
  if (src.includes("import.meta")) {
    src = `
      import.meta.url = ${JSON.stringify(pathToFileURL(filepath).toString())};
      import.meta.dirname = ${JSON.stringify(path.dirname(filepath))};
      import.meta.filename = ${JSON.stringify(filepath)};
    `.trim().replace(/\n/g, "") + src;
  }
  src = esbuild.transformSync(src, {
    loader,
    format: "cjs",
    target: "esnext",
    jsx: "automatic",
    jsxImportSource: "#ssr",
    jsxDev: true,
    sourcefile: filepath,
  }).code;
  return module._compile(src, filepath, "commonjs");
}

function resolveClientRef(sourcePath: string, ref: string) {
  const filePath = resolveFrom(sourcePath, ref);
  if (
    !filePath.endsWith(".client.ts") &&
    !filePath.endsWith(".client.tsx")
  ) {
    throw new Error("addScript must be a .client.ts or .client.tsx");
  }
  return path.relative(projectSrc, filePath);
}

// TODO: extract the marko compilation tools out, lazy load them
export interface MarkoCacheEntry {
  src: string;
  scannedClientRefs: string[];
}
export const markoCache = new Map<string, MarkoCacheEntry>();
function loadMarko(module: NodeJS.Module, filepath: string) {
  let cache = markoCache.get(filepath);
  if (!cache) {
    let src = fs.readFileSync(filepath, "utf8");
    // A non-standard thing here is Clover Sitegen implements
    // its own client side scripting stuff, so it overrides
    // bare client import statements to it's own usage.
    const scannedClientRefs = new Set<string>();
    if (src.match(/^\s*client\s+import\s+["']/m)) {
      src = src.replace(
        /^\s*client\s+import\s+("[^"]+"|'[^']+')[^\n]+/m,
        (_, src) => {
          const ref = JSON.parse(`"${src.slice(1, -1)}"`);
          const resolved = resolveClientRef(filepath, ref);
          scannedClientRefs.add(resolved);
          return `<CloverScriptInclude=${
            JSON.stringify(getScriptId(resolved))
          } />`;
        },
      ) + '\nimport { addScript as CloverScriptInclude } from "#sitegen";\n';
    }

    src = marko.compileSync(src, filepath).code;
    src = src.replace("marko/debug/html", "#ssr/marko");
    cache = { src, scannedClientRefs: Array.from(scannedClientRefs) };
    markoCache.set(filepath, cache);
  }

  const { src, scannedClientRefs } = cache;
  return loadEsbuildCode(module, filepath, src, {
    scannedClientRefs,
  });
}

function loadMdx(module: NodeJS.Module, filepath: string) {
  const input = fs.readFileSync(filepath);
  const out = mdx.compileSync(input, { jsxImportSource: "#ssr" }).value;
  const src = typeof out === "string" ? out : Buffer.from(out).toString("utf8");
  return loadEsbuildCode(module, filepath, src);
}

function loadCss(module: NodeJS.Module, _filepath: string) {
  module.exports = {};
}

export function reloadRecursive(filepath: string) {
  filepath = path.resolve(filepath);
  const existing = cache[filepath];
  if (existing) deleteRecursiveInner(filepath, existing);
  fileStats.clear();
  return require(filepath);
}

export function unload(filepath: string) {
  filepath = path.resolve(filepath);
  const existing = cache[filepath];
  if (existing) delete cache[filepath];
  fileStats.delete(filepath);
}

function deleteRecursiveInner(id: string, module: any) {
  if (id.includes(path.sep + "node_modules" + path.sep)) {
    return;
  }
  delete cache[id];
  for (const child of module.children) {
    if (child.filename.includes("/engine/")) return;
    const existing = cache[child.filename];
    if (existing === child) deleteRecursiveInner(child.filename, existing);
  }
}

export function getCssImports(filepath: string) {
  filepath = path.resolve(filepath);
  if (!require.cache[filepath]) throw new Error(filepath + " was never loaded");
  return fileStats.get(filepath)?.cssImportsRecursive ?? [];
}

export function getClientScriptRefs(filepath: string) {
  filepath = path.resolve(filepath);
  const module = require.cache[filepath];
  if (!module) throw new Error(filepath + " was never loaded");
  return module.cloverClientRefs ?? [];
}

export function getSourceCode(filepath: string) {
  filepath = path.resolve(filepath);
  const module = require.cache[filepath];
  if (!module) throw new Error(filepath + " was never loaded");
  if (!module.cloverSourceCode) {
    throw new Error(filepath + " did not record source code");
  }
  return module.cloverSourceCode;
}

export function resolveFrom(src: string, dest: string) {
  try {
    return createRequire(src).resolve(dest);
  } catch (err: any) {
    if (err.code === "MODULE_NOT_FOUND" && err.requireStack.length <= 1) {
      err.message = err.message.split("\n")[0] + " from '" + src + "'";
    }
    throw err;
  }
}

const importRegExp =
  /import\s+(\*\sas\s([a-zA-Z0-9$_]+)|{[^}]+})\s+from\s+(?:"#sitegen"|'#sitegen')/s;
const getSitegenAddScriptRegExp = /addScript(?:\s+as\s+([a-zA-Z0-9$_]+))?/;
interface ResolvedClientRefs {
  code: string;
  refs: string[];
}
export function resolveClientRefs(
  code: string,
  filepath: string,
): ResolvedClientRefs {
  // This match finds a call to 'import ... from "#sitegen"'
  const importMatch = code.match(importRegExp);
  if (!importMatch) return { code, refs: [] };
  const items = importMatch[1];
  let identifier = "";
  if (items.startsWith("{")) {
    const clauseMatch = items.match(getSitegenAddScriptRegExp);
    if (!clauseMatch) return { code, refs: [] }; // did not import
    identifier = clauseMatch[1] || "addScript";
  } else if (items.startsWith("*")) {
    identifier = importMatch[2] + "\\s*\\.\\s*addScript";
  } else {
    throw new Error("Impossible");
  }
  identifier = identifier.replaceAll("$", "\\$"); // only needed escape
  const findCallsRegExp = new RegExp(
    `\\b(${identifier})\\s*\\(("[^"]+"|'[^']+')\\)`,
    "gs",
  );
  const scannedClientRefs = new Set<string>();
  code = code.replace(findCallsRegExp, (_, call, arg) => {
    const ref = JSON.parse(`"${arg.slice(1, -1)}"`);
    const resolved = resolveClientRef(filepath, ref);
    scannedClientRefs.add(resolved);
    return `${call}(${JSON.stringify(getScriptId(resolved))})`;
  });
  return { code, refs: Array.from(scannedClientRefs) };
}

export function getScriptId(file: string) {
  return (path.isAbsolute(file) ? path.relative(projectSrc, file) : file)
    .replaceAll("\\", "/");
}

declare global {
  namespace NodeJS {
    interface Module {
      cloverClientRefs?: string[];
      cloverSourceCode?: string;

      _compile(
        this: NodeJS.Module,
        content: string,
        filepath: string,
        format: "module" | "commonjs",
      ): unknown;
    }
  }
}
declare module "node:module" {
  export function _resolveFilename(
    id: string,
    parent: NodeJS.Module,
  ): unknown;
}

import * as fs from "./lib/fs.ts";
import * as path from "node:path";
import { pathToFileURL } from "node:url";
import * as esbuild from "esbuild";
import * as marko from "@marko/compiler";
import { createRequire } from "node:module";
import * as mdx from "@mdx-js/mdx";
import * as self from "./hot.ts";
import { Buffer } from "node:buffer";