// Copyright (c) 2019-2024 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.

// Module huffman builds canonical Huffman codes from a per-symbol array of bit
// lengths. This is the assignment step shared by RFC 1951 (DEFLATE, used by
// compress.deflate / zlib / gzip) and RFC 7541 Appendix B (HPACK, used by the
// HTTP/2 code in net.http): given only the lengths, both standards rebuild the
// exact same codes via the `bl_count` / `next_code` algorithm (RFC 1951 §3.2.2).
//
// What the two callers do NOT share — and therefore what this module is
// parameterized by — is the maximum code length (DEFLATE caps at 15 bits and
// decodes via a flat 2^max_bits table; HPACK goes up to 30 bits, where a flat
// table is infeasible) and the bit order (DEFLATE is LSB-first and bit-reverses
// every code, HPACK is MSB-first). The bit I/O loops and codec-specific
// semantics (EOS/padding, end-of-block, extra bits, distance alphabets) stay in
// the callers.
module huffman

// BitOrder selects how a code's bits are laid out in the returned `u32`.
pub enum BitOrder {
	// msb_first keeps the canonical code as-is: the first transmitted bit is
	// the most-significant bit of the code (RFC 7541 / HPACK).
	msb_first
	// lsb_first reverses each code within its length, so the first transmitted
	// bit is the least-significant bit (RFC 1951 / DEFLATE). This is the form a
	// flat LSB-first decode table is indexed by.
	lsb_first
}

// Config parameterizes build(). All fields are required and have no defaults:
// the two callers have intentionally different requirements (15 vs 30 bits,
// LSB vs MSB), so an implicit default would silently fit only one of them.
@[params]
pub struct Config {
pub:
	lengths   []int    @[required] // per-symbol code length in bits; 0 marks an unused symbol
	max_bits  int      @[required] // maximum allowed code length; must be >= every nonzero length
	bit_order BitOrder @[required] // .msb_first (HPACK) or .lsb_first (DEFLATE)
}

// Table is the result of build(): the canonical code for every symbol, plus the
// metadata a caller needs to drive its own bit I/O and to build a decode
// structure (flat_table() for small max_bits, decode_map() otherwise).
@[noinit]
pub struct Table {
pub:
	codes     []u32 // per-symbol code, right-aligned in a u32, in `bit_order`
	lengths   []int // per-symbol bit length (a copy of the input)
	max_bits  int
	bit_order BitOrder
}

// max_flat_bits is the largest max_bits for which flat_table() will
// allocate a table (2^max_bits entries). DEFLATE's 15 fits; HPACK's 30 does
// not and must use decode_map() / a bit-at-a-time decoder instead.
pub const max_flat_bits = 18

// flat_invalid_entry marks a flat_table() slot that no code maps to.
pub const flat_invalid_entry = u32(0xffff_ffff)

// flat_length_bits is how many low bits of a flat_table() entry hold the
// code length; the symbol is stored in the remaining high bits. 5 bits hold
// lengths up to 31, covering every max_bits <= max_flat_bits.
pub const flat_length_bits = 5

// next_codes validates `lengths` against `max_bits` and returns the canonical
// starting code for each length (RFC 1951 §3.2.2): next_code[l] is the code the
// first symbol of length l receives, and callers post-increment it per symbol.
// It also reports whether the code is `complete` (uses the whole code space,
// i.e. the Kraft inequality holds with equality), which lets flat_table() skip
// pre-filling the table. It is the single source of truth for the code
// assignment shared by build() and flat_table(). Errors on max_bits < 1 or
// > 32, a negative length or one exceeding max_bits, or an over-subscribed code
// (Kraft inequality); an incomplete (under-subscribed) code is allowed, as both
// DEFLATE and HPACK use.
fn next_codes(lengths []int, max_bits int) !([]u32, bool) {
	if max_bits < 1 {
		return error('huffman: max_bits must be >= 1, got ${max_bits}')
	}
	if max_bits > 32 {
		return error('huffman: max_bits ${max_bits} exceeds 32 (u32 code storage)')
	}
	mut bl_count := []int{len: max_bits + 1}
	for sym, l in lengths {
		if l < 0 {
			return error('huffman: negative length ${l} for symbol ${sym}')
		}
		if l > max_bits {
			return error('huffman: length ${l} for symbol ${sym} exceeds max_bits ${max_bits}')
		}
		if l > 0 {
			bl_count[l]++
		}
	}
	// Kraft inequality: sum over used symbols of 2^(max_bits - len) must not
	// exceed 2^max_bits, i.e. the code must not be over-subscribed. left == 0 at
	// the end means the code is complete (covers every index of a flat table).
	mut left := u64(1) << max_bits
	for bits in 1 .. max_bits + 1 {
		used := u64(bl_count[bits]) << (max_bits - bits)
		if used > left {
			return error('huffman: over-subscribed code (lengths exceed the code space)')
		}
		left -= used
	}
	mut next_code := []u32{len: max_bits + 1}
	mut c := u32(0)
	for bits in 1 .. max_bits + 1 {
		c = (c + u32(bl_count[bits - 1])) << 1
		next_code[bits] = c
	}
	return next_code, left == 0
}

// build assigns a canonical Huffman code to every symbol from its bit length.
// Symbols with length 0 are unused and get code 0. It returns the codes plus
// the metadata for a decode structure; use it when you need the per-symbol
// codes (e.g. HPACK encoding) and/or decode_map(). Callers that only want a
// flat decode table should use flat_table(), which avoids materializing the
// codes array. See next_codes() for the validation rules.
pub fn build(cfg Config) !Table {
	mut next_code, _ := next_codes(cfg.lengths, cfg.max_bits)!
	mut codes := []u32{len: cfg.lengths.len}
	for sym, l in cfg.lengths {
		if l == 0 {
			continue
		}
		code := next_code[l]
		next_code[l]++
		codes[sym] = if cfg.bit_order == .lsb_first { bit_reverse(code, l) } else { code }
	}
	return Table{
		codes:     codes
		lengths:   cfg.lengths.clone()
		max_bits:  cfg.max_bits
		bit_order: cfg.bit_order
	}
}

// flat_table builds a 2^max_bits lookup table for fast decode of short codes
// (the DEFLATE strategy), directly from the lengths. Each entry is
// `(symbol << flat_length_bits) | length`, or flat_invalid_entry for indices no
// code reaches. The table is indexed by the next max_bits bits read from the
// stream in `bit_order`. It returns an error if max_bits > max_flat_bits, since
// the table would be prohibitively large (use build() + decode_map() instead).
//
// Unlike build(), this allocates no per-symbol codes array and does not copy
// the lengths: it assigns each code inline while filling the table, so a hot
// caller rebuilding a tree per block (compress.deflate) pays no extra
// allocation over hand-rolling the loop.
pub fn flat_table(cfg Config) ![]u32 {
	if cfg.max_bits > max_flat_bits {
		return error('huffman: max_bits ${cfg.max_bits} > max_flat_bits ${max_flat_bits}; use build() + decode_map()')
	}
	mut next_code, complete := next_codes(cfg.lengths, cfg.max_bits)!
	table_size := 1 << cfg.max_bits
	// A complete code writes every table slot, so the invalid pre-fill is dead
	// work; allocate zeroed (vcalloc) and skip it. An incomplete code leaves
	// gaps that must read back as flat_invalid_entry, so it pays the fill.
	mut table := if complete {
		[]u32{len: table_size}
	} else {
		[]u32{len: table_size, init: flat_invalid_entry}
	}
	for sym, l in cfg.lengths {
		if l == 0 {
			continue
		}
		raw := next_code[l]
		next_code[l]++
		entry := (u32(sym) << flat_length_bits) | u32(l)
		// The (max_bits - l) bits beyond the code are don't-cares, so a code of
		// length l fills 2^(max_bits - l) table slots. Where those slots sit
		// depends on bit_order: LSB-first codes occupy every index whose low l
		// bits match the code (stride by 2^l); MSB-first codes occupy a
		// contiguous block whose high l bits match the code.
		if cfg.bit_order == .lsb_first {
			step := 1 << l
			mut idx := int(bit_reverse(raw, l))
			for idx < table_size {
				table[idx] = entry
				idx += step
			}
		} else {
			block := 1 << (cfg.max_bits - l)
			base := int(raw) * block
			for k in 0 .. block {
				table[base + k] = entry
			}
		}
	}
	return table
}

// decode_map builds a map from a packed (length, code) key to its symbol, for
// codecs whose max_bits is too large for a flat table (HPACK's 30 bits). The
// key is `(u64(length) << 32) | code`; a decoder accumulates bits MSB-first and
// looks up after each bit. Only defined for .msb_first tables, where the
// accumulated value matches the stored code; it returns an error otherwise.
pub fn (t Table) decode_map() !map[u64]int {
	if t.bit_order != .msb_first {
		return error('huffman: decode_map requires .msb_first bit order')
	}
	mut m := map[u64]int{}
	for sym, l in t.lengths {
		if l == 0 {
			continue
		}
		m[(u64(l) << 32) | u64(t.codes[sym])] = sym
	}
	return m
}

// bit_reverse reverses the low `n` bits of `v` (used to convert a canonical
// MSB-first code into the LSB-first form a DEFLATE bit reader consumes).
fn bit_reverse(v u32, n int) u32 {
	mut r := u32(0)
	mut val := v
	for _ in 0 .. n {
		r = (r << 1) | (val & 1)
		val >>= 1
	}
	return r
}