Initial project import

tools/v/vlib/hash/huffman/huffman.v

@@ -0,0 +1,225 @@
+// 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
+}