Merge pull request #55 from spicyjpeg/psxmdec

Full MDEC support, C library refactors, cleanups and bugfixes (v0.20)

1 files changed, 297 insertions, 0 deletions

diff --git a/libpsn00b/psxpress/generate_lookup_table.py b/libpsn00b/psxpress/generate_lookup_table.py
new file mode 100644
index 0000000..b40771f
--- /dev/null
+++ b/libpsn00b/psxpress/generate_lookup_table.py
@@ -0,0 +1,297 @@
+#!/usr/bin/env python3
+# Huffman lookup table generator script for psxpress
+# (C) 2022 spicyjpeg - MPL licensed
+
+import sys, json
+from array     import array
+from itertools import repeat
+from argparse  import ArgumentParser, FileType
+
+HUFFMAN_TREE = {
+	"10":			0xfe00, # End of block
+	"11":			(  0,  1 ),
+	"01": {
+		"1":		(  1,  1 ),
+		"00":		(  0,  2 ),
+		"01":		(  2,  1 )
+	},
+	"001": {
+		"01":		(  0,  3 ),
+		"10":		(  4,  1 ),
+		"11":		(  3,  1 ),
+		"00000":	( 13,  1 ),
+		"00001":	(  0,  6 ),
+		"00010":	( 12,  1 ),
+		"00011":	( 11,  1 ),
+		"00100":	(  3,  2 ),
+		"00101":	(  1,  3 ),
+		"00110":	(  0,  5 ),
+		"00111":	( 10,  1 )
+	},
+	"0001": {
+		"00":		(  7,  1 ),
+		"01":		(  6,  1 ),
+		"10":		(  1,  2 ),
+		"11":		(  5,  1 )
+	},
+	"00001": {
+		"00":		(  2,  2 ),
+		"01":		(  9,  1 ),
+		"10":		(  0,  4 ),
+		"11":		(  8,  1 )
+	},
+	"0000001": {
+		"000":		( 16,  1 ),
+		"001":		(  5,  2 ),
+		"010":		(  0,  7 ),
+		"011":		(  2,  3 ),
+		"100":		(  1,  4 ),
+		"101":		( 15,  1 ),
+		"110":		( 14,  1 ),
+		"111":		(  4,  2 )
+	},
+	"00000001": {
+		"0000":		(  0, 11 ),
+		"0001":		(  8,  2 ),
+		"0010":		(  4,  3 ),
+		"0011":		(  0, 10 ),
+		"0100":		(  2,  4 ),
+		"0101":		(  7,  2 ),
+		"0110":		( 21,  2 ),
+		"0111":		( 20,  1 ),
+		"1000":		(  0,  9 ),
+		"1001":		( 19,  1 ),
+		"1010":		( 18,  1 ),
+		"1011":		(  1,  5 ),
+		"1100":		(  3,  3 ),
+		"1101":		(  0,  8 ),
+		"1110":		(  6,  2 ),
+		"1111":		( 17,  1 )
+	},
+	"000000001": {
+		"0000":		( 10,  2 ),
+		"0001":		(  9,  2 ),
+		"0010":		(  5,  3 ),
+		"0011":		(  3,  4 ),
+		"0100":		(  2,  5 ),
+		"0101":		(  1,  7 ),
+		"0110":		(  1,  6 ),
+		"0111":		(  0, 15 ),
+		"1000":		(  0, 14 ),
+		"1001":		(  0, 13 ),
+		"1010":		(  0, 12 ),
+		"1011":		( 26,  1 ),
+		"1100":		( 25,  1 ),
+		"1101":		( 24,  1 ),
+		"1110":		( 23,  1 ),
+		"1111":		( 22,  1 )
+	},
+	"0000000001": {
+		"0000":		(  0, 31 ),
+		"0001":		(  0, 30 ),
+		"0010":		(  0, 29 ),
+		"0011":		(  0, 28 ),
+		"0100":		(  0, 27 ),
+		"0101":		(  0, 26 ),
+		"0110":		(  0, 25 ),
+		"0111":		(  0, 24 ),
+		"1000":		(  0, 23 ),
+		"1001":		(  0, 22 ),
+		"1010":		(  0, 21 ),
+		"1011":		(  0, 20 ),
+		"1100":		(  0, 19 ),
+		"1101":		(  0, 18 ),
+		"1110":		(  0, 17 ),
+		"1111":		(  0, 16 )
+	},
+	"00000000001": {
+		"0000":		(  0, 40 ),
+		"0001":		(  0, 39 ),
+		"0010":		(  0, 38 ),
+		"0011":		(  0, 37 ),
+		"0100":		(  0, 36 ),
+		"0101":		(  0, 35 ),
+		"0110":		(  0, 34 ),
+		"0111":		(  0, 33 ),
+		"1000":		(  0, 32 ),
+		"1001":		(  1, 14 ),
+		"1010":		(  1, 13 ),
+		"1011":		(  1, 12 ),
+		"1100":		(  1, 11 ),
+		"1101":		(  1, 10 ),
+		"1110":		(  1,  9 ),
+		"1111":		(  1,  8 )
+	},
+	"000000000001": {
+		"0000":		(  1, 18 ),
+		"0001":		(  1, 17 ),
+		"0010":		(  1, 16 ),
+		"0011":		(  1, 15 ),
+		"0100":		(  6,  3 ),
+		"0101":		( 16,  2 ),
+		"0110":		( 15,  2 ),
+		"0111":		( 14,  2 ),
+		"1000":		( 13,  2 ),
+		"1001":		( 12,  2 ),
+		"1010":		( 11,  2 ),
+		"1011":		( 31,  1 ),
+		"1100":		( 30,  1 ),
+		"1101":		( 29,  1 ),
+		"1110":		( 28,  1 ),
+		"1111":		( 27,  1 )
+	}
+}
+
+## Utilities
+
+def to_int10(value):
+	clamped = min(max(int(value), -0x200), 0x1ff)
+	return clamped + (0 if clamped >= 0 else 0x400)
+
+def uint32_to_lines(data, indent = "\t", columns = 6):
+	for offset in range(0, len(data), columns):
+		line = f"{indent}0x{data[offset]:08x}"
+
+		for item in data[(offset + 1):(offset + columns)]:
+			line += f", 0x{item:08x}"
+
+		yield line
+
+## Table generation
+
+def iterate_tree(tree):
+	for code, value in tree.items():
+		if type(value) is dict:
+			# Iterate through any subtree recursively.
+			for suffix, _value in iterate_tree(value):
+				yield f"{code}{suffix}", _value
+
+		elif type(value) is tuple:
+			run_length, ac = value
+			yield f"{code}0", (run_length << 10) | to_int10(ac)
+			yield f"{code}1", (run_length << 10) | to_int10(-ac)
+
+		else:
+			yield code, value
+
+def generate_table(codes, table_bits, prefix_bits = 0):
+	table = array("I", repeat(0, 2 ** table_bits))
+
+	for code, value in codes:
+		used_bits = len(code)
+		free_bits = table_bits - (used_bits - prefix_bits)
+		index     = int(code[prefix_bits:], 2) << free_bits
+
+		# Fill out every entry in the table whose index starts with the same
+		# string of bits.
+		for combo in range(2 ** free_bits):
+			table[index | combo] = (used_bits << 16) | value
+
+	return table
+
+def compress_table(table):
+	values     = []
+	last_value = table[0]
+	run_length = 0
+
+	for value in table[1:]:
+		if value == last_value and run_length < 0x7ff:
+			run_length += 1
+			continue
+
+		# The run length is stored in the top 11 bits of each value, which are
+		# otherwise unused.
+		values.append((run_length << 21) | last_value)
+		last_value = value
+		run_length = 0
+
+	values.append((run_length << 21) | last_value)
+	return array("I", values)
+
+## Main
+
+UNCOMPRESSED_TEMPLATE = """static const DECDCTTAB {name} = {{
+	.lut = {{
+{short}
+	}},
+	.lut00 = {{
+{long}
+	}}
+}};
+"""
+COMPRESSED_TEMPLATE = """static const uint32_t {name}[{length}] = {{
+{table}
+}};
+"""
+
+def get_args():
+	parser = ArgumentParser(
+		description = "Generates a Huffman lookup table structure, to be used by DecDCTvlc2()."
+	)
+	parser.add_argument(
+		"-c", "--compress",
+		action = "store_true",
+		help   = "generate run-length compressed data instead of a DECDCTTAB struct"
+	)
+	parser.add_argument(
+		"-n", "--name",
+		type    = str,
+		default = "_default_huffman_table",
+		help    = "set the symbol name in the generated C source",
+		metavar = "file"
+	)
+	parser.add_argument(
+		"-t", "--tree",
+		type    = FileType("rt"),
+		help    = "use a custom Huffman tree from the specified JSON file",
+		metavar = "json_file"
+	)
+	parser.add_argument(
+		"-o", "--output",
+		type    = FileType("wt"),
+		default = sys.stdout,
+		help    = "where to output generated table (stdout by default)",
+		metavar = "file"
+	)
+
+	return parser.parse_args()
+
+def main():
+	args = get_args()
+	tree = json.load(args.tree) if args.tree else HUFFMAN_TREE
+
+	short_codes, short_bits = [], 0
+	long_codes, long_bits   = [], 0
+
+	for pair in iterate_tree(tree):
+		if (code := pair[0]).startswith("00000000"):
+			long_codes.append(pair)
+			long_bits = max(long_bits, len(code) - 8)
+		else:
+			short_codes.append(pair)
+			short_bits = max(short_bits, len(code))
+
+	short_table = generate_table(short_codes, short_bits, 0)
+	long_table  = generate_table(long_codes,  long_bits,  8)
+
+	if args.compress:
+		short_table.extend(long_table)
+		table = compress_table(short_table)
+
+		source = COMPRESSED_TEMPLATE.format(
+			name   = args.name,
+			length = len(table),
+			table  = ",\n".join(uint32_to_lines(table, "\t"))
+		)
+	else:
+		source = UNCOMPRESSED_TEMPLATE.format(
+			name  = args.name,
+			short = ",\n".join(uint32_to_lines(short_table, "\t\t")),
+			long  = ",\n".join(uint32_to_lines(long_table,  "\t\t"))
+		)
+
+	with args.output as _file:
+		_file.write(source)
+
+if __name__ == "__main__":
+	main()