scriptc/tests/harness/shard.ts

/* Within-file partitioning for the corpus-driven harnesses.
 *
 * vitest's --shard is FILE-granular, and this suite's wall time is dominated
 * by a few monster files (differential.test.ts iterates the whole corpus
 * sequentially in one worker), so a naive file shard leaves one shard nearly
 * as slow as the whole suite. CI therefore sets SCRIPTC_TEST_SHARD="i/n"
 * alongside --shard=i/n, and every harness that globs a big case list runs
 * it through shardSelect below: small files distribute by file, monster
 * files split by case.
 *
 * The assignment hashes each case's stable NAME — never its list index — so
 * a program keeps its shard when the corpus grows around it (compile caches
 * stay warm across corpus additions) and, since the hash maps every key to
 * exactly one shard in [1..n], the union of the n shards runs every case
 * exactly once by construction (shard.test.ts pins the property). Unset =
 * everything: local runs are unchanged.
 */
import { createHash } from "node:crypto";

export interface ShardSpec {
  /** 1-based shard index. */
  index: number;
  count: number;
}

/** Parse a shard spec of the form "i/n" (1 <= i <= n). Undefined when the
 * variable is unset or empty; malformed values throw — a typo silently
 * running everything (or nothing) would defeat the matrix. */
export function parseShardSpec(
  raw: string | undefined = process.env["SCRIPTC_TEST_SHARD"],
): ShardSpec | undefined {
  if (raw === undefined || raw === "") return undefined;
  const m = /^([1-9]\d*)\/([1-9]\d*)$/.exec(raw);
  if (!m) throw new Error(`SCRIPTC_TEST_SHARD must look like "2/3", got "${raw}"`);
  const index = Number(m[1]);
  const count = Number(m[2]);
  if (index > count) throw new Error(`SCRIPTC_TEST_SHARD index out of range: "${raw}"`);
  return { index, count };
}

/** The shard (1-based) a key belongs to among `count` shards: the first four
 * bytes of sha1(key) mod count — stable across runs, platforms, and corpus
 * growth. Total on its domain, so exactly one shard owns any key. */
export function shardOf(key: string, count: number): number {
  return (createHash("sha1").update(key).digest().readUInt32BE(0) % count) + 1;
}

/** The subset of `items` this shard owns. No spec (the local default) or a
 * single shard returns the whole list unchanged. */
export function shardSelect<T>(
  items: readonly T[],
  keyOf: (item: T) => string,
  spec: ShardSpec | undefined = parseShardSpec(),
): T[] {
  if (spec === undefined || spec.count === 1) return [...items];
  return items.filter((item) => shardOf(keyOf(item), spec.count) === spec.index);
}

/** Allocate expensive fixture lists by estimated cost. Longest cases go to
 * the lightest slice first, with key/index tie breakers so every process
 * independently constructs the same disjoint partition. Unlike hash-based
 * sharding, fixture additions can move existing cases between slices. */
export function balancedShardSelect<T>(
  items: readonly T[],
  keyOf: (item: T) => string,
  weightOf: (item: T) => number,
  spec: ShardSpec | undefined = parseShardSpec(),
): T[] {
  if (spec === undefined || spec.count === 1) return [...items];
  const keys = new Set<string>();
  const ranked = items.map((item, index) => {
    const key = keyOf(item);
    const weight = weightOf(item);
    if (keys.has(key)) throw new Error(`duplicate balanced shard key: ${key}`);
    if (!Number.isFinite(weight) || weight <= 0) throw new Error(`invalid balanced shard weight for ${key}: ${weight}`);
    keys.add(key);
    return { index, key, weight };
  }).sort((a, b) => b.weight - a.weight || (a.key < b.key ? -1 : a.key > b.key ? 1 : 0));
  const loads = Array<number>(spec.count).fill(0);
  const owners = Array<number>(items.length);
  for (const item of ranked) {
    let slice = 0;
    for (let candidate = 1; candidate < loads.length; candidate++) {
      if (loads[candidate]! < loads[slice]!) slice = candidate;
    }
    loads[slice]! += item.weight;
    owners[item.index] = slice + 1;
  }
  return items.filter((_, index) => owners[index] === spec.index);
}

/** Suffix for describe titles so shard membership is visible in CI logs:
 * ", shard 2/3" under a spec, "" otherwise. */
export function shardSuffix(spec: ShardSpec | undefined = parseShardSpec()): string {
  return spec === undefined ? "" : `, shard ${spec.index}/${spec.count}`;
}