vac-optimizer/src/app/services/frozen-engine.spec.ts

import { afterAll, describe, expect, it } from 'bun:test';
import { frenchPublicHolidaysForYear } from './declared-non-working-days';
import {
	calculatePlanningAsyncInternal,
	calculatePlanningInternal,
	type PlanningArtifacts,
	type PlanningOutcome,
} from './leave-planner.core';
import {
	calculatePlanningAsyncInternal as frozenCalculateAsync,
	calculatePlanningInternal as frozenCalculate,
} from './frozen-engine/leave-planner.core.frozen.mjs';
import {
	DAY,
	generateRequest,
	iso,
	mulberry32,
	weekday,
	type Req,
} from './frozen-engine/random-planning-requests';

/** Rebuild instructions for the frozen engine: docs/benchmark/README.md. */
const BASE_SEED = 280_000;
const CASES = 5000;
const CHUNK = 200;

interface Run {
	outcome: PlanningOutcome;
	artifacts?: Pick<PlanningArtifacts, 'explanation' | 'finalState' | 'explanationError'>;
}

function run(
	calculate: (
		request: unknown,
		c?: undefined,
		a?: (artifacts: PlanningArtifacts) => void,
	) => PlanningOutcome,
	request: Req,
): Run {
	let artifacts: Run['artifacts'];
	const outcome = calculate(structuredClone(request), undefined, (a) => {
		artifacts = {
			explanation: a.explanation,
			finalState: a.finalState,
			explanationError: a.explanationError,
		};
	});
	return { outcome, artifacts };
}

const kinds = new Map<string, number>();
const ruleTypes = new Set<string>();
const count = (key: string) => kinds.set(key, (kinds.get(key) ?? 0) + 1);

let slowest = { seed: 0, ms: 0 };

function crossesYearAndHoliday(request: Req): void {
	const { start, end } = request['horizon'] as { start: string; end: string };
	if (start.slice(0, 4) !== end.slice(0, 4)) count('crosses-year');
	for (let ms = Date.parse(start); ms <= Date.parse(end); ms += DAY)
		if (
			weekday(ms) &&
			frenchPublicHolidaysForYear(new Date(ms).getUTCFullYear()).has(iso(ms))
		) {
			count('holiday-week');
			return;
		}
}

function assertSameAsFrozen(seed: number): void {
	const request = generateRequest(seed);
	const t0 = performance.now();
	const current = run(calculatePlanningInternal, request);
	const frozen = run(frozenCalculate, request);
	const ms = performance.now() - t0;
	if (ms > slowest.ms) slowest = { seed, ms: Math.round(ms) };
	crossesYearAndHoliday(request);
	count(current.outcome.kind);
	for (const rule of request.rules) ruleTypes.add(String(rule['type']).toLowerCase());
	if ((request['declaredDays'] as unknown[]).length) count('declared');
	if (request.rules.some((r) => r['rewardDays'])) count('reward');
	try {
		expect(current).toStrictEqual(frozen);
	} catch (error) {
		throw new Error(
			`Engine differs from frozen engine. seed=${seed}\nrequest=${JSON.stringify(request)}\n${
				(error as Error).message
			}`,
			{ cause: error },
		);
	}
}

type AsyncEngine = typeof calculatePlanningAsyncInternal;

/** Cancels after a seed-derived step count: during preparation, mid-search, or never. */
function cancelAfter(seed: number, total: number): number {
	const rnd = mulberry32(seed ^ 0x5bd1e995);
	const roll = rnd();
	if (roll < 0.1) return Math.floor(rnd() * total);
	if (roll < 0.25) return Number.POSITIVE_INFINITY;
	// Preparation is most of the steps; spend the rest of the budget on the search.
	return Math.floor(total * (0.9 + rnd() * 0.1));
}

async function runCancelled(engine: AsyncEngine, request: Req, limit: number) {
	const calls: unknown[] = [];
	const artifacts: unknown[] = [];
	const explanations: unknown[] = [];
	const outcomes: unknown[] = [];
	let polls = 0;
	const outcome = await engine(
		structuredClone(request),
		(plan, state) => calls.push({ plan, state }),
		({ timings: _timings, ...rest }) => artifacts.push(rest),
		(o) => outcomes.push(o),
		(...args) => explanations.push(args),
		() => ++polls > limit,
	);
	return { outcome, calls, artifacts, explanations, outcomes, polls };
}

const cancelledKinds = new Map<string, number>();
let incumbents = 0;

async function assertCancelledSameAsFrozen(seed: number): Promise<void> {
	const request = generateRequest(seed);
	const total = (await runCancelled(calculatePlanningAsyncInternal, request, Infinity)).polls;
	const limit = cancelAfter(seed, total);
	const current = await runCancelled(calculatePlanningAsyncInternal, request, limit);
	const frozen = await runCancelled(frozenCalculateAsync, request, limit);
	const kind = current.outcome.kind;
	cancelledKinds.set(kind, (cancelledKinds.get(kind) ?? 0) + 1);
	if (kind === 'PlanningIncomplete' && current.calls.length) incumbents++;
	try {
		expect(current).toStrictEqual(frozen);
	} catch (error) {
		throw new Error(
			`Cancelled run differs from frozen engine. seed=${seed} limit=${limit}\nrequest=${JSON.stringify(request)}\n${
				(error as Error).message
			}`,
			{ cause: error },
		);
	}
}

const ONLY_SEED = process.env['FROZEN_ENGINE_SEED'];

describe('frozen engine differential check', () => {
	if (ONLY_SEED) {
		it(`matches the frozen engine for seed ${ONLY_SEED}`, async () => {
			assertSameAsFrozen(Number(ONLY_SEED));
			await assertCancelledSameAsFrozen(Number(ONLY_SEED));
		}, 120_000);
		return;
	}
	for (let from = 0; from < CASES; from += CHUNK) {
		it(`matches the frozen engine for seeds ${BASE_SEED + from}..${BASE_SEED + from + CHUNK - 1}`, () => {
			for (let i = from; i < from + CHUNK; i++) assertSameAsFrozen(BASE_SEED + i);
		}, 120_000);
	}

	for (let from = 0; from < CASES; from += CHUNK * 2) {
		it(`matches the frozen engine when cancelled, seeds ${BASE_SEED + from}.. every 10th`, async () => {
			for (let i = from; i < from + CHUNK * 2; i += 10)
				await assertCancelledSameAsFrozen(BASE_SEED + i);
		}, 300_000);
	}

	it('retains incumbents of cancelled calculations', () => {
		console.info('cancelled coverage', Object.fromEntries(cancelledKinds), { incumbents });
		expect(cancelledKinds.get('PlanningIncomplete') ?? 0).toBeGreaterThan(20);
		expect(incumbents).toBeGreaterThan(10);
	});

	afterAll(() => {
		console.info(
			'frozen-engine coverage of generated cases',
			Object.fromEntries(kinds),
			[...ruleTypes],
			'slowest case',
			slowest,
		);
	});

	it('generates every outcome kind and rule type', () => {
		for (const kind of [
			'Plan',
			'PlanningConflict',
			'InvalidInput',
			'declared',
			'reward',
			'crosses-year',
			'holiday-week',
		])
			expect(kinds.get(kind) ?? 0, kind).toBeGreaterThan(20);
		expect(['if_then', 'optimization', 'time_window'].filter((t) => !ruleTypes.has(t))).toEqual(
			[],
		);
	});
});