Sample Sort

import type { NumericTrace } from './trace';

function paddedLength(length: number) {
	let size = 1;
	while (size < length) size *= 2;
	return size;
}

function networkCompare(
	trace: NumericTrace,
	values: number[],
	left: number,
	right: number,
	ascending: boolean
) {
	const beforeLeft = values[left];
	const beforeRight = values[right];
	if (left < trace.array.length && right < trace.array.length) {
		trace.compareValues(left, right, beforeLeft, beforeRight);
	} else {
		trace.note(`Compare wire ${left} with padded wire ${right}`);
	}
	if (ascending ? beforeLeft > beforeRight : beforeLeft < beforeRight) {
		values[left] = beforeRight;
		values[right] = beforeLeft;
		if (left < trace.array.length)
			trace.write(left, values[left], `Network exchange at wire ${left}`);
		if (right < trace.array.length)
			trace.write(right, values[right], `Network exchange at wire ${right}`);
	}
}

export function bitonicSort(trace: NumericTrace) {
	const n = trace.array.length;
	if (n < 2) return;
	const size = paddedLength(n);
	const values = [...trace.array, ...new Array<number>(size - n).fill(Infinity)];
	for (let block = 2; block <= size; block *= 2) {
		trace.note(`Build bitonic sequences of length ${block}`);
		for (let stride = block / 2; stride > 0; stride = Math.floor(stride / 2)) {
			for (let index = 0; index < size; index++) {
				const peer = index ^ stride;
				if (peer > index) {
					const ascending = (index & block) === 0;
					networkCompare(trace, values, index, peer, ascending);
				}
			}
		}
	}
}

export function batcherOddEvenMergeSort(trace: NumericTrace) {
	const n = trace.array.length;
	if (n < 2) return;
	const size = paddedLength(n);
	const values = [...trace.array, ...new Array<number>(size - n).fill(Infinity)];
	function merge(start: number, length: number, gap: number) {
		const step = gap * 2;
		if (step < length) {
			merge(start, length, step);
			merge(start + gap, length, step);
			for (let index = start + gap; index + gap < start + length; index += step) {
				networkCompare(trace, values, index, index + gap, true);
			}
		} else {
			networkCompare(trace, values, start, start + gap, true);
		}
	}
	function sort(start: number, length: number) {
		if (length < 2) return;
		const half = length / 2;
		sort(start, half);
		sort(start + half, half);
		trace.note(`Odd-even merge network joins two runs of length ${half}`);
		merge(start, length, 1);
	}
	sort(0, size);
}

/** Sample-sort stages are independent partitions that can run on separate workers. */
export function sampleSort(trace: NumericTrace) {
	const n = trace.array.length;
	if (n < 2) return;
	const bucketCount = Math.min(4, Math.max(2, Math.ceil(Math.sqrt(n))));
	const sampleCount = Math.min(n, bucketCount * 4);
	const sample = Array.from(
		{ length: sampleCount },
		(_, index) => trace.array[Math.floor((index * n) / sampleCount)]
	);
	sample.sort((a, b) => a - b);
	const splitters = Array.from(
		{ length: bucketCount - 1 },
		(_, index) =>
			sample[Math.min(sample.length - 1, Math.floor(((index + 1) * sample.length) / bucketCount))]
	);
	trace.note(`Choose sample splitters ${splitters.join(', ')}`);
	const buckets = Array.from({ length: bucketCount }, () => [] as number[]);
	for (let index = 0; index < n; index++) {
		const value = trace.array[index];
		let bucket = 0;
		while (bucket < splitters.length && value > splitters[bucket]) bucket++;
		buckets[bucket].push(value);
		trace.note(`Send ${value} to sample bucket ${bucket}`, { [index]: 'bg-vis-write' });
	}
	let output = 0;
	for (let bucket = 0; bucket < bucketCount; bucket++) {
		trace.note(`Sort sample bucket ${bucket}`);
		const values = buckets[bucket];
		function mergeSort(start: number, end: number) {
			if (end - start < 2) return;
			const middle = start + Math.floor((end - start) / 2);
			mergeSort(start, middle);
			mergeSort(middle, end);
			const source = values.slice(start, end);
			let left = 0;
			let right = middle - start;
			for (let position = start; position < end; position++) {
				if (left < middle - start && (right >= end - start || source[left] <= source[right])) {
					if (right < end - start)
						trace.note(`Compare ${source[left]} with ${source[right]} in bucket ${bucket}`);
					values[position] = source[left++];
				} else {
					if (left < middle - start)
						trace.note(`Compare ${source[left]} with ${source[right]} in bucket ${bucket}`);
					values[position] = source[right++];
				}
			}
		}
		mergeSort(0, values.length);
		for (const value of values) trace.write(output++, value);
	}
}