Huffman coding.

// Huffman coding over bytes: count each byte, then keep merging the two lightest trees
// until one is left. A byte's code length is how deep it ends up in that tree.
export const MAX_INPUT = 1024; // a teaching bound; it keeps every code well within 15 bits
export const MAX_LENGTH = 15; // DEFLATE's longest code (RFC 1951)

export type HuffmanErrorCode = 'too-long' | 'bad-lengths' | 'missing-code' | 'bad-bits';

export class HuffmanError extends Error {
	readonly code: HuffmanErrorCode;
	constructor(code: HuffmanErrorCode, message: string) {
		super(message);
		this.name = 'HuffmanError';
		this.code = code;
	}
}

// A leaf's id is its byte value. Each merge makes a node with the next id from 256.
export type TreeNode = { id: number; weight: number; left: number | null; right: number | null };

export type Code = {
	counts: number[]; // how often each byte 0–255 appears
	nodes: TreeNode[]; // the leaves in byte order, then one node per merge, in merge order
	lengths: number[]; // each byte's code length, 0 if it never appears
};

// Lighter first. Equal weights go to the smaller id, so leaves come before merged nodes.
const lighter = (a: TreeNode, b: TreeNode) =>
	a.weight < b.weight || (a.weight === b.weight && a.id < b.id);

export function buildCode(input: Uint8Array): Code {
	if (input.length > MAX_INPUT)
		throw new HuffmanError('too-long', `input is limited to ${MAX_INPUT} bytes`);
	const counts = Array<number>(256).fill(0);
	for (const byte of input) counts[byte]++;
	const nodes: TreeNode[] = [];
	for (let byte = 0; byte < 256; byte++)
		if (counts[byte] > 0) nodes.push({ id: byte, weight: counts[byte], left: null, right: null });
	const lengths = Array<number>(256).fill(0);
	if (nodes.length === 1) {
		lengths[nodes[0].id] = 1; // a lone byte still needs one bit each time it appears
		return { counts, nodes, lengths };
	}
	const leaves = nodes.length;
	const heap: TreeNode[] = [];
	for (const node of nodes) push(heap, node);
	while (heap.length > 1) {
		const left = pop(heap); // the lightest tree gets bit 0
		const right = pop(heap); // the next lightest gets bit 1
		const merged = {
			id: 256 + nodes.length - leaves,
			weight: left.weight + right.weight,
			left: left.id,
			right: right.id
		};
		nodes.push(merged);
		push(heap, merged);
	}
	// Walk down from the last node made, the root. A leaf's depth is its code length.
	const byId = new Map(nodes.map((node) => [node.id, node]));
	const stack: [TreeNode, number][] = nodes.length > 0 ? [[nodes[nodes.length - 1], 0]] : [];
	while (stack.length > 0) {
		const [node, depth] = stack.pop()!;
		if (node.left === null || node.right === null) lengths[node.id] = depth;
		else stack.push([byId.get(node.left)!, depth + 1], [byId.get(node.right)!, depth + 1]);
	}
	return { counts, nodes, lengths };
}

// RFC 1951 §3.2.2: shorter codes come first, and codes of one length follow byte order.
// The lengths alone decide every code, so a decoder needs only the lengths.
export function canonicalCodes(lengths: readonly number[]): string[] {
	checkLengths(lengths);
	const count = Array<number>(MAX_LENGTH + 1).fill(0);
	for (const length of lengths) if (length > 0) count[length]++;
	const next = Array<number>(MAX_LENGTH + 1).fill(0);
	let code = 0;
	for (let bits = 1; bits <= MAX_LENGTH; bits++) {
		code = (code + count[bits - 1]) << 1;
		next[bits] = code;
	}
	return lengths.map((length) =>
		length === 0 ? '' : (next[length]++).toString(2).padStart(length, '0')
	);
}

// Huffman coding over bytes: count each byte, then keep merging the two lightest trees
// until one is left. A byte's code length is how deep it ends up in that tree.
const (
	MaxInput  = 1024 // a teaching bound; it keeps every code well within 15 bits
	MaxLength = 15   // DEFLATE's longest code (RFC 1951)
)

type HuffmanError struct {
	Code    string // "too-long", "bad-lengths", "missing-code", or "bad-bits"
	Message string
}

func (e *HuffmanError) Error() string { return e.Code + ": " + e.Message }

// A leaf's ID is its byte value. Each merge makes a node with the next ID from 256.
// Left and Right are -1 for a leaf.
type Node struct {
	ID, Weight, Left, Right int
}

type Code struct {
	Counts  [256]int // how often each byte appears
	Nodes   []Node   // the leaves in byte order, then one node per merge, in merge order
	Lengths [256]int // each byte's code length, 0 if it never appears
}

func BuildCode(input []byte) (Code, error) {
	var code Code
	if len(input) > MaxInput {
		return code, &HuffmanError{"too-long", fmt.Sprintf("input is limited to %d bytes", MaxInput)}
	}
	for _, b := range input {
		code.Counts[b]++
	}
	for b, count := range code.Counts {
		if count > 0 {
			code.Nodes = append(code.Nodes, Node{ID: b, Weight: count, Left: -1, Right: -1})
		}
	}
	if len(code.Nodes) == 1 {
		code.Lengths[code.Nodes[0].ID] = 1 // a lone byte still needs one bit each time it appears
		return code, nil
	}
	leaves := len(code.Nodes)
	trees := &forest{}
	for _, node := range code.Nodes {
		heap.Push(trees, node)
	}
	for trees.Len() > 1 {
		left := heap.Pop(trees).(Node)  // the lightest tree gets bit 0
		right := heap.Pop(trees).(Node) // the next lightest gets bit 1
		merged := Node{ID: 256 + len(code.Nodes) - leaves, Weight: left.Weight + right.Weight, Left: left.ID, Right: right.ID}
		code.Nodes = append(code.Nodes, merged)
		heap.Push(trees, merged)
	}
	// Walk down from the last node made, the root. A leaf's depth is its code length.
	byID := map[int]Node{}
	for _, node := range code.Nodes {
		byID[node.ID] = node
	}
	type visit struct{ node, depth int }
	stack := []visit{}
	if len(code.Nodes) > 0 {
		stack = append(stack, visit{code.Nodes[len(code.Nodes)-1].ID, 0})
	}
	for len(stack) > 0 {
		top := stack[len(stack)-1]
		stack = stack[:len(stack)-1]
		node := byID[top.node]
		if node.Left < 0 {
			code.Lengths[node.ID] = top.depth
		} else {
			stack = append(stack, visit{node.Left, top.depth + 1}, visit{node.Right, top.depth + 1})
		}
	}
	return code, nil
}

// CanonicalCodes follows RFC 1951 §3.2.2: shorter codes come first, and codes of one length
// follow byte order. The lengths alone decide every code, so a decoder needs only the lengths.
func CanonicalCodes(lengths [256]int) ([256]string, error) {
	var codes [256]string
	if err := checkLengths(lengths); err != nil {
		return codes, err
	}
	var count, next [MaxLength + 1]int
	for _, length := range lengths {
		if length > 0 {
			count[length]++
		}
	}
	code := 0
	for bits := 1; bits <= MaxLength; bits++ {
		code = (code + count[bits-1]) << 1
		next[bits] = code
	}
	for b, length := range lengths {
		if length > 0 {
			codes[b] = fmt.Sprintf("%0*b", length, next[length])
			next[length]++
		}
	}
	return codes, nil
}