sitegen/framework/incremental.ts

// Incremental build system using singleton state.
// See `work()`, `compile()`, and `invalidate()` for details.
//
// All state is serializable to allow recovering state across sessions.
// This framework special-cases the asset map, but is otherwise
// agnostic of what it is a compiler for.
let running = false;
let jobs = 0;
let newKeys = 0;
let seenWorks = new Set<string>(); // for detecting conflict vs overwrite
let seenWrites = new Set<string>(); // for detecting conflict vs overwrite
let works = new Map<string, Work>();
let files = new Map<string, TrackedFile>(); // keyed by `toRel` path
let writes = new Map<string, FileWrite>();
let assets = new Map<string, Asset>(); // keyed by hash

export interface Ref<T> {
  /** This method is compatible with `await` syntax */
  then(
    onFulfilled: (value: T) => void,
    onRejected: (error: unknown) => void,
  ): void;
  key: string;
}
type Job<I = any, O = any> = (io: Io, input: I) => Promise<O>;

/**
 * Declare and a unit of work. Return value is memoized and
 * only rebuilt when inputs (declared via `Io`) change. Outputs
 * are written at the end of a compilation (see `compile`).
 *
 * If the returned `Ref` is not awaited or read
 * via io.readWork, the job is never started.
 */
export function work<O>(job: Job<void, O>): Ref<O>;
export function work<I, O>(job: Job<I, O>, input: I): Ref<O>;
export function work<I, O>(job: Job<I, O>, input: I = null as I): Ref<O> {
  const source = JSON.stringify(util.getCallSites(2)[1]);
  const keySource = [source, util.inspect(input)].join(":");
  const key = crypto.createHash("sha1").update(keySource).digest("base64url");
  ASSERT(running);
  ASSERT(
    !seenWorks.has(key),
    `Key '${key}' must be unique during the build. ` +
      `To fix this, provide a manual 'key' argument.`,
  );
  seenWorks.add(key);

  const prev = works.get(key) as Work<O> | null;
  if (prev) {
    return { key, then: (done) => done(prev.value) };
  }

  async function perform() {
    const io = new Io(key);
    jobs += 1;
    newKeys += 1;
    try {
      const value = await job(io, input);
      validateSerializable(value, "");
      const { reads, writes } = io;
      works.set(key, {
        value,
        affects: [],
        reads,
        writes,
        debug: source,
      });
      for (const add of reads.files) {
        const { affects } = UNWRAP(files.get(add));
        ASSERT(!affects.includes(key));
        affects.push(key);
      }
      for (const add of reads.works) {
        const { affects } = UNWRAP(works.get(add));
        ASSERT(!affects.includes(key));
        affects.push(key);
      }
      return value;
    } finally {
      jobs -= 1;
    }
  }

  let cached: Promise<O>;
  return {
    key,
    then: (fufill, reject) => void (cached ??= perform()).then(fufill, reject),
  };
}

export async function compile<T>(compiler: () => Promise<T>) {
  ASSERT(!running, `Cannot run twice at the same time.`);
  try {
    running = true;
    ASSERT(jobs === 0);
    const start = performance.now();
    const timerSpinner = new Spinner({
      text: () =>
        `sitegen! [${
          ((performance.now() - start) / 1000).toFixed(
            1,
          )
        }s]`,
      fps: 10,
    });
    using _endTimerSpinner = { [Symbol.dispose]: () => timerSpinner.stop() };

    const value = await compiler();
    ASSERT(jobs === 0);
    timerSpinner.text = "incremental flush";
    await flush(start);
    timerSpinner.stop();
    return {
      value,
      watchFiles: new Set(files.keys()),
      newOutputs: Array.from(seenWrites).filter(x => x.startsWith('f:')).map(x => x.slice(2)),
      newAssets: !Array.from(seenWrites).some(x => x.startsWith('a:')),
    };
  } finally {
    running = false;
    newKeys = 0;
    seenWrites.clear();
    seenWorks.clear();
  }
}

export async function flush(start: number) {
  // Trim
  const detachedFiles = new Set<string>;
  const referencedAssets = new Set<string>;
  for (const [k, { writes: { assets } }] of works) {
    if (seenWorks.has(k)) {
      for (const asset of assets.values()) referencedAssets.add(asset.hash);
      continue;
    }
    deleteWork(k);
  }
  for (const [k, file] of files) {
    if (file.affects.length > 0) continue;
    files.delete(k);
    detachedFiles.add(k);
  }
  for (const k of assets.keys()) {
    if (!referencedAssets.has(k))
      assets.delete(k);
  }

  const p = [];
  // File writes
  let dist = 0;
  for (const [key, { buffer, size }] of writes) {
    if (buffer) p.push(fs.writeMkdir(path.join(`.clover/o/${key}`), buffer));
    dist += size;
  }
  // Asset map
  {
    const { json, blob } = getAssetManifest();
    const jsonString = Buffer.from(JSON.stringify(json));
    p.push(fs.writeMkdir(".clover/o/static.json", jsonString));
    p.push(fs.writeMkdir(".clover/o/static.blob", blob));
    dist += blob.byteLength + jsonString.byteLength;
  }
  await Promise.all(p);

  // Incremental state
  const serialized = msgpackr.pack(serialize());
  await fs.writeMkdir(".clover/incr.state", serialized);
  const time = (performance.now() - start).toFixed(0);
  console.success(`sitegen! in ${time} ms`);
  console.writeLine(`  - ${works.size} keys (${works.size - newKeys} cached)`);
  console.writeLine(`  - ${assets.size} static assets`);
  console.writeLine(
    `  - dist: ${formatSize(dist)}, incremental: ${
      formatSize(serialized.byteLength)
    }`,
  );
}

export async function restore() {
  let buffer;
  try {
    buffer = await fs.readFile(".clover/incr.state");
  } catch (err: any) {
    if (err.code !== "ENOENT") throw err;
  }
  if (!buffer) return;
  await deserialize(buffer);
}

function forceInvalidate(entry: { affects: string[] }) {
  const queue = [...entry.affects];
  let key;
  while ((key = queue.shift())) {
    const affects = deleteWork(key);
    queue.push(...affects);
  }
}

function deleteWork(key: string) {
  const work = works.get(key);
  if (!work) return [];
  const { reads, affects, writes: w } = work;
  for (const remove of reads.files) {
    const { affects } = UNWRAP(files.get(remove));
    ASSERT(affects.includes(key));
    affects.splice(affects.indexOf(key), 1);
  }
  for (const remove of reads.works) {
    const { affects } = UNWRAP(works.get(remove), remove);
    ASSERT(affects.includes(key));
    affects.splice(affects.indexOf(key), 1);
  }
  for (const remove of affects) {
    const { reads: { works: list } } = UNWRAP(works.get(remove), remove);
    ASSERT(list.has(key));
    list.delete(key);
  }
  for (const file of w.files) {
    if (UNWRAP(writes.get(file)).work === key)
      writes.delete(file);
  }
  // Assets are temporarily kept, trimmed via manual GC after compilation.

  works.delete(key);
  return affects;
}

export function reset() {
  ASSERT(!running);
  works.clear();
  files.clear();
  assets.clear();
}

export function serialize() {
  const fileEntries = Array.from(files, ([k, v]) =>
    [
      k,
      v.type,
      v.type === 'f' ? v.lastModified : v.type === 'd' ? v.contentHash : null,
      ...v.affects,
    ] as const);
  const workEntries = Array.from(works, ([k, v]) =>
    [
      k,
      v.value,
      Array.from(v.reads.files),
      Array.from(v.reads.works),
      Array.from(v.writes.files),
      Array.from(v.writes.assets, ([k, { headers }]) => [k, headers] as const),
      v.affects,
    ] as const);
  const expectedFilesOnDisk = Array.from(
    writes,
    ([k, { size, work }]) => [k, size, work] as const,
  );
  const assetEntries = Array.from(
    assets,
    ([k, asset]) => [k, asset.raw, asset.gzip, asset.zstd] as const,
  );
  return [
    1,
    fileEntries,
    workEntries,
    expectedFilesOnDisk,
    assetEntries,
  ] as const;
}
type SerializedState = ReturnType<typeof serialize>;

/* No-op on failure */
async function deserialize(buffer: Buffer) {
  const decoded = msgpackr.decode(buffer) as SerializedState;
  if (!Array.isArray(decoded)) return false;
  if (decoded[0] !== 1) return false;
  const [, fileEntries, workEntries, expectedFilesOnDisk, assetEntries] =
    decoded;
  for (const [k, type, content, ...affects] of fileEntries) {
    if (type === "f") {
      ASSERT(typeof content === "number");
      files.set(k, { type, affects, lastModified: content });
    } else if (type === 'd') {
      ASSERT(typeof content === "string");
      files.set(k, { type, affects, contentHash: content, contents: [] });
    } else {
      files.set(k, { type, affects });
    }
  }
  for (const entry of workEntries) {
    const [
      k,
      value,
      readFiles,
      readWorks,
      writeFiles,
      writeAssets,
      affects,
    ] = entry;
    works.set(k, {
      value,
      reads: {
        files: new Set(readFiles),
        works: new Set(readWorks),
      },
      writes: {
        files: new Set(writeFiles),
        assets: new Map(Array.from(writeAssets, ([k, headers]) => [k, {
          hash: JSON.parse(UNWRAP(headers.etag)),
          headers,
        }])),
      },
      affects,
    });
  }
  const statFiles = await Promise.all(expectedFilesOnDisk
    .map(([k, size, work]) =>
      fs.stat(path.join(".clover/o", k))
        .catch((err) => {
          if (err.code === "ENOENT") return null;
          throw err;
        })
        .then((stat) => ({ k, size, work, stat }))
    ));
  for (const { k, stat, work, size } of statFiles) {
    if (stat?.size === size) {
      writes.set(k, {
        size: size,
        buffer: null,
        work,
      });
    } else {
      forceInvalidate({ affects: [work] });
    }
  }
  for (const [hash, raw, gzip, zstd] of assetEntries) {
    assets.set(hash, { raw, gzip, zstd });
  }

  await Promise.all(Array.from(files, ([key, file]) => invalidateEntry(key, file)));
}

export async function invalidate(filePath: string): Promise<boolean> {
  const key = toRel(toAbs(filePath));
  const file = UNWRAP(files.get(key), `Untracked file '${key}'`)
  return invalidateEntry(key, file)
}

export async function invalidateEntry(key: string, file: TrackedFile): Promise<boolean> {
  try {
    if (file.type === "d") {
      const contents = file.contents = await fs.readdir(key);
      contents.sort();
      const contentHash = crypto
        .createHash("sha1")
        .update(contents.join("\0"))
        .digest("base64url");
      if (file.contentHash !== contentHash) {
        file.contentHash = contentHash;
        throw new Error();
      }
    } else if (file.type === 'f') {
      const lastModified = await fs.stat(key)
        .then(x => Math.floor(x.mtimeMs), () => 0);
      if (file.lastModified !== lastModified) {
        file.lastModified = lastModified;
        throw new Error();
      }
    } else {
      file.type satisfies 'null';
      const stat = await fs.stat(key).catch(() => null);
      if (stat)  throw new Error();
    }
    return false;
  } catch (e) {
    forceInvalidate(file);
    hot.unload(toAbs(key));
    if (file.type === 'null') files.delete(key);
    return true;
  }
}

export function getAssetManifest() {
  const writer = new BufferWriter();
  const asset = Object.fromEntries(
    Array.from(works, (work) => work[1].writes.assets)
      .filter((map) => map.size > 0)
      .flatMap((map) =>
        Array.from(map, ([key, { hash, headers }]) => {
          const { raw, gzip, zstd } = UNWRAP(
            assets.get(hash),
            `Asset ${key} (${hash})`,
          );
          return [key, {
            raw: writer.write(raw, "raw:" + hash),
            gzip: writer.write(gzip, "gzip:" + hash),
            zstd: writer.write(zstd, "zstd:" + hash),
            headers,
          }] as const;
        })
      ),
  ) satisfies BuiltAssetMap;
  return { json: asset, blob: writer.get() };
}

/* Input/Output with automatic tracking.
 * - Inputs read with Io are tracked to know when to rebuild
 * - Outputs written with Io are deleted when abandoned.
 */
export class Io {
  constructor(public key: string) {}
  reads: Reads = { files: new Set(), works: new Set() };
  writes: Writes = { files: new Set(), assets: new Map() };

  #trackFs(file: string) {
    const resolved = toAbs(file);
    const key = toRel(resolved);
    this.reads.files.add(key);
    return { resolved, key };
  }
  async readWork<T>(ref: Ref<T>): Promise<T> {
    this.reads.works.add(ref.key);
    return await ref;
  }
  /** Track a file in the compilation without reading it. */
  async trackFile(file: string) {
    const { key, resolved } = this.#trackFs(file);
    if (!files.get(key)) {
      let lastModified: number = 0;
      try {
        lastModified = Math.floor((await fs.stat(file)).mtimeMs);
        files.set(key, { type: "f", lastModified, affects: [] });
      } catch {
        files.set(key, { type: "null", affects: [] });
      }
    }
    return resolved;
  }
  async readFile(file: string) {
    return fs.readFile(await this.trackFile(file), "utf-8");
  }
  async readJson<T>(file: string) {
    return JSON.parse(await this.readFile(file)) as T;
  }
  async readDir(dir: string) {
    const { key, resolved } = this.#trackFs(dir);
      const existing = files.get(key);
    try {
      if (existing?.type === 'd') return existing.contents;
      const contents = await fs.readdir(resolved);
      contents.sort();
      const contentHash = crypto
        .createHash("sha1")
        .update(contents.join("\0"))
        .digest("base64url");
      files.set(key, {
        type: "d",
        affects: [],
        contentHash,
        contents,
      });
      return contents;
    } catch (err) {
      if (!existing) files.set(key, { type: "null", affects: [] });
      throw err;
    }
  }
  async readDirRecursive(dir: string): Promise<string[]> {
    const dirs = await this.readDir(dir);
    return (
      await Promise.all(
        dirs.map(async (child) => {
          const abs = path.join(dir, child);
          const stat = await fs.stat(abs);
          if (stat.isDirectory()) {
            return (await this.readDirRecursive(abs)).map((grand) =>
              path.join(child, grand)
            );
          } else {
            return child;
          }
        }),
      )
    ).flat();
  }
  /* Track all dependencies of a module. */
  async import<T>(file: string): Promise<T> {
    const { resolved } = this.#trackFs(file);
    try {
      return require(resolved) as T;
    } finally {
      const queue = [resolved];
      const seen = new Set<string>();
      let current;
      while ((current = queue.shift())) {
        const stat = hot.getFileStat(current);
        if (!stat) continue;
        const { key } = this.#trackFs(current);
        if (!files.get(key)) {
          files.set(key, {
            type: "f",
            affects: [],
            lastModified: stat?.lastModified ?? 0,
          });
        }
        for (const imp of stat.imports) {
          if (!seen.has(imp)) {
            seen.add(imp);
            queue.push(imp);
          }
        }
      }
    }
  }
  async writeAsset(
    pathname: string,
    blob: string | Buffer,
    headersOption?: HeadersInit,
  ) {
    ASSERT(pathname.startsWith("/"));
    ASSERT(!seenWrites.has("a:" + pathname));

    const buffer = typeof blob === "string" ? Buffer.from(blob) : blob;

    const headers = new Headers(headersOption ?? {});
    const hash = crypto.createHash("sha1").update(buffer).digest("hex");
    if (!headers.has("Content-Type")) {
      headers.set("Content-Type", mime.contentTypeFor(pathname));
    }
    headers.set("ETag", JSON.stringify(hash));
    this.writes.assets.set(pathname, {
      hash,
      // @ts-expect-error TODO
      headers: Object.fromEntries(headers),
    });
    if (!assets.has(hash)) {
      jobs += 1;
      assets.set(hash, undefined!);
      const [gzipBuffer, zstdBuffer] = await Promise.all([
        gzip(buffer),
        zstdCompress(buffer),
      ]);
      assets.set(hash, {
        raw: buffer,
        gzip: gzipBuffer,
        zstd: zstdBuffer,
      });
      jobs -= 1;
    }
  }
  writeFile(subPath: string, blob: string | Buffer) {
    ASSERT(!subPath.startsWith("/"));
    ASSERT(
      !seenWrites.has("f:" + subPath),
      `File overwritten: ${JSON.stringify(subPath)}`,
    );
    seenWrites.add("f:" + subPath);
    const buffer = Buffer.isBuffer(blob) ? blob : Buffer.from(blob);
    writes.set(subPath, {
      buffer,
      size: buffer.byteLength,
      work: this.key,
    });
  }
}

class BufferWriter {
  size = 0;
  seen = new Map<string, BufferView>();
  buffers: Buffer[] = [];

  write(buffer: Buffer, hash: string): BufferView {
    let view = this.seen.get(hash);
    if (view) return view;
    view = [this.size, this.size += buffer.byteLength];
    this.seen.set(hash, view);
    this.buffers.push(buffer);
    return view;
  }

  get() {
    return Buffer.concat(this.buffers);
  }
}

export function validateSerializable(value: unknown, key: string) {
  if (typeof value === "string") {
    if (value.includes(hot.projectRoot)) {
      throw new Error(
        `Return value must not contain the CWD for portability, found at ${key}`,
      );
    }
  } else if (value && typeof value === "object") {
    if (Array.isArray(value)) {
      value.forEach((item, i) => validateSerializable(item, `${key}[${i}]`));
    } else if (Object.getPrototypeOf(value) === Object.prototype || Buffer.isBuffer(value)) {
      Object.entries(value).forEach(([k, v]) =>
        validateSerializable(v, `${key}.${k}`)
      );
    } else {
      throw new Error(
        `Return value must be a plain JS object, found ${
          Object.getPrototypeOf(value).constructor.name
        } at ${key}`,
      );
    }
  } else if (["bigint", "function", "symbol"].includes(typeof value)) {
    throw new Error(
      `Return value must be a plain JS object, found ${typeof value} at ${key}`,
    );
  }
}

export function toAbs(absPath: string) {
  return path.resolve(hot.projectRoot, absPath);
}

export function toRel(absPath: string) {
  return path.relative(hot.projectRoot, absPath).replaceAll("\\", "/");
}

export type BufferView = [start: number, end: number];
interface Reads {
  files: Set<string>;
  works: Set<string>;
}
interface FileWrite {
  buffer: Buffer | null;
  size: number;
  work: string;
}
interface Writes {
  files: Set<string>;
  assets: Map<string, {
    hash: string;
    headers: Record<string, string>;
  }>;
}
interface Asset {
  raw: Buffer;
  gzip: Buffer;
  zstd: Buffer;
}
interface Work<T = unknown> {
  debug?: string;
  value: T;
  reads: Reads;
  writes: Writes;
  affects: string[];
}
type TrackedFile =
  & {
    affects: string[];
  }
  & (
    | { type: "f"; lastModified: number }
    | { type: "d"; contentHash: string; contents: string[] }
    | { type: "null"; }
  );
export interface BuiltAssetMap {
  [route: string]: BuiltAsset;
}
export interface BuiltAsset {
  raw: BufferView;
  gzip: BufferView;
  zstd: BufferView;
  headers: Record<string, string>;
}

const gzip = util.promisify(zlib.gzip);
const zstdCompress = util.promisify(zlib.zstdCompress);

import * as fs from "#sitegen/fs";
import * as path from "node:path";
import * as hot from "./hot.ts";
import * as util from "node:util";
import * as crypto from "node:crypto";
import * as mime from "#sitegen/mime";
import * as zlib from "node:zlib";
import * as console from "@paperclover/console";
import { Spinner } from "@paperclover/console/Spinner";
import { formatSize } from "@/file-viewer/format.ts";
import * as msgpackr from "msgpackr";