mirror of
https://github.com/RetroPortingToolKit/psxrecomp
synced 2026-10-26 18:25:35 -04:00
Matthew Stanley
142 lines
6.2 KiB
Python
142 lines
6.2 KiB
Python
"""BPE container contract on invented bytes; no game data."""
|
|
import struct
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
import bpe_container
|
|
from bpe_container import BpeError, decode
|
|
|
|
MAGIC = 0x1A455042
|
|
|
|
|
|
def encode_table(left, right):
|
|
"""Run-length table coding with identity stretches skipped."""
|
|
out, c = bytearray(), 0
|
|
while c < 256:
|
|
stretch = 0
|
|
while c + stretch < 256 and stretch < 128 and left[c + stretch] == c + stretch:
|
|
stretch += 1
|
|
if stretch >= 3:
|
|
out.append(0x7E + stretch) # skips stretch - 1 entries, then one entry
|
|
out.append(c + stretch - 1) # the last identity entry of the stretch
|
|
c += stretch
|
|
continue
|
|
run = 0
|
|
while c + run < 256 and run < 128 and not (
|
|
left[c + run] == c + run and
|
|
all(left[c + run + k] == c + run + k for k in range(1, 3) if c + run + k < 256) and
|
|
c + run + 3 <= 256):
|
|
run += 1
|
|
run = max(run, 1)
|
|
out.append(run - 1)
|
|
for k in range(c, c + run):
|
|
out.append(left[k])
|
|
if left[k] != k:
|
|
out.append(right[k])
|
|
c += run
|
|
return bytes(out)
|
|
|
|
|
|
def encode_block(symbols, decoded_size, left=None, right=None):
|
|
"""A block of input `symbols` that decodes to `decoded_size` bytes."""
|
|
left = list(range(256)) if left is None else left
|
|
right = [0] * 256 if right is None else right
|
|
body = encode_table(left, right) + bytes(symbols)
|
|
body += bytes(-len(body) % 4)
|
|
return struct.pack('<I', decoded_size) + body
|
|
|
|
|
|
def container(blocks, declared, magic=MAGIC):
|
|
return struct.pack('<II', declared, magic) + b''.join(blocks) + struct.pack('<I', 0)
|
|
|
|
|
|
class BpeContainerTests(unittest.TestCase):
|
|
def pairs(self):
|
|
left, right = list(range(256)), [0] * 256
|
|
left[0xF0], right[0xF0] = ord('A'), ord('B') # 0xF0 -> "AB"
|
|
left[0xF1], right[0xF1] = 0xF0, ord('C') # 0xF1 -> "ABC" (chained)
|
|
left[0x10], right[0x10] = 0x20, 0x30 # 0x10 -> 0x20 0x30
|
|
return left, right
|
|
|
|
def one(self, symbols, size, left=None, right=None):
|
|
return container([encode_block(symbols, size, left, right)], size)
|
|
|
|
def test_literals_pairs_chains_and_multiple_blocks(self):
|
|
left, right = self.pairs()
|
|
first = encode_block([0xF1, ord('x'), 0xF0, 0x10], 8, left, right)
|
|
second = encode_block(list(b'plain'), 5)
|
|
data = container([first, second], 13)
|
|
body, consumed = decode(data, magic=MAGIC)
|
|
self.assertEqual(body, b'ABCxAB' + bytes([0x20, 0x30]) + b'plain')
|
|
self.assertEqual(consumed, len(data))
|
|
|
|
def test_consumed_counts_words_and_offset_selects_the_container(self):
|
|
data = self.one(list(b'abc'), 3)
|
|
for lead in (b'', bytes(8)):
|
|
with self.subTest(lead=len(lead)):
|
|
body, consumed = decode(lead + data + b'trailing', magic=MAGIC, offset=len(lead))
|
|
self.assertEqual((body, consumed), (b'abc', len(data)))
|
|
|
|
def test_skip_coded_identity_stretches_keep_tables_short_and_decode_correctly(self):
|
|
left, right = self.pairs()
|
|
table = encode_table(left, right)
|
|
self.assertLess(len(table), 40, 'identity stretches must be skip-coded')
|
|
body, _ = decode(self.one([0xF1, 0x10, 0x41, 0xF0], 8, left, right), magic=MAGIC)
|
|
self.assertEqual(body, b'ABC' + bytes([0x20, 0x30]) + b'A' + b'AB')
|
|
|
|
def test_a_skip_that_reaches_the_last_entry_ends_the_table(self):
|
|
# 128 explicit entries (entry 0 is a pair), then a 128-entry skip: index 256.
|
|
table = bytes([127, ord('Z'), ord('Y')]) + bytes(range(1, 128)) + bytes([0xFF])
|
|
block = struct.pack('<I', 2) + table + bytes([0])
|
|
block += bytes(-len(block) % 4)
|
|
body, _ = decode(container([block], 2), magic=MAGIC)
|
|
self.assertEqual(body, b'ZY')
|
|
# After a skip that stays inside the table, one entry still follows.
|
|
table = bytes([127]) + bytes(range(128)) + bytes([0xFE, 0x41, 0x42]) # skip 127 -> 255, entry 255
|
|
block = struct.pack('<I', 1) + table + bytes([5])
|
|
block += bytes(-len(block) % 4)
|
|
self.assertEqual(decode(container([block], 1), magic=MAGIC)[0], b'')
|
|
|
|
def test_rejects_wrong_magic_size_truncation_and_misalignment(self):
|
|
good = self.one(list(b'abc'), 3)
|
|
cases = [
|
|
(self.one(list(b'abc'), 3).replace(struct.pack('<I', MAGIC), struct.pack('<I', 0x12345678), 1),
|
|
dict(), 'magic'),
|
|
(container([encode_block(list(b'abc'), 3)], 4), dict(), 'declares'),
|
|
(container([encode_block(list(b'abc'), 3)], 2), dict(), 'declared size'),
|
|
(good[:-4], dict(), 'truncated'),
|
|
(good[:12], dict(), 'truncated'),
|
|
(good[:9], dict(), 'truncated'),
|
|
(good, dict(offset=2), 'misaligned'),
|
|
(good, dict(max_output=2), 'output limit'),
|
|
]
|
|
for data, options, error in cases:
|
|
with self.subTest(error=error, options=options), self.assertRaisesRegex(BpeError, error):
|
|
decode(data, magic=MAGIC, **options)
|
|
|
|
def test_rejects_cyclic_tables_and_table_runs_past_the_last_entry(self):
|
|
left, right = list(range(256)), [0] * 256
|
|
left[0xF0], right[0xF0] = 0xF1, ord('x')
|
|
left[0xF1], right[0xF1] = 0xF0, ord('y')
|
|
with self.assertRaisesRegex(BpeError, 'does not terminate'):
|
|
decode(self.one([0xF0], 1, left, right), magic=MAGIC)
|
|
# 128 + 101 + 41 entries = 270 > 256.
|
|
table = (bytes([127]) + bytes(range(128)) + bytes([100]) + bytes(range(128, 229)) +
|
|
bytes([40]) + bytes(range(229, 256)) + bytes(16))
|
|
block = struct.pack('<I', 1) + table + bytes([5])
|
|
block += bytes(-len(block) % 4)
|
|
with self.assertRaisesRegex(BpeError, 'passes its last entry'):
|
|
decode(container([block], 1), magic=MAGIC)
|
|
|
|
def test_magic_is_required_evidence_not_a_default(self):
|
|
with self.assertRaises(TypeError):
|
|
decode(self.one(list(b'a'), 1)) # no magic given
|
|
self.assertTrue(issubclass(BpeError, ValueError))
|
|
self.assertIs(bpe_container.BpeError, BpeError)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|