143 lines
4.7 KiB
Python
143 lines
4.7 KiB
Python
"""LZO1X, as used by the game's compressed UE3 packages (docs/formats/upk.md).
|
|
|
|
`decompress` is a pure-Python LZO1X decoder (same semantics as lzo1x_decompress_safe; the packages
|
|
use 128 KiB blocks, so speed is fine for block-level access). `compress_stored` writes a valid LZO1X
|
|
stream that holds the data as one literal run (no matches): bigger than real compression, but any
|
|
LZO1X decoder (the game's included) reads it back exactly, and nothing else has to change.
|
|
"""
|
|
M2_MAX_OFFSET = 0x0800
|
|
|
|
|
|
class LzoError(ValueError):
|
|
pass
|
|
|
|
|
|
def _copy_match(out, dist, n):
|
|
start = len(out) - dist
|
|
if start < 0:
|
|
raise LzoError("match distance beyond output start")
|
|
if dist >= n:
|
|
out += out[start:start + n]
|
|
else: # overlapping: repeat the last `dist` bytes
|
|
pat = bytes(out[start:])
|
|
out += (pat * (n // dist + 1))[:n]
|
|
|
|
|
|
def decompress(src, out_len=None):
|
|
"""LZO1X block -> bytes. `out_len` (if given) is checked."""
|
|
src = bytes(src)
|
|
n = len(src)
|
|
out = bytearray()
|
|
ip = 0
|
|
|
|
def ext(t, base):
|
|
nonlocal ip
|
|
while src[ip] == 0:
|
|
t += 255
|
|
ip += 1
|
|
t += base + src[ip]
|
|
ip += 1
|
|
return t
|
|
|
|
try:
|
|
state = "loop"
|
|
if src[0] > 17:
|
|
t = src[0] - 17
|
|
ip = 1
|
|
if t < 4:
|
|
state = "match_next"
|
|
else:
|
|
out += src[ip:ip + t]
|
|
ip += t
|
|
state = "first_literal_run"
|
|
while True:
|
|
if state == "loop":
|
|
t = src[ip]
|
|
ip += 1
|
|
if t >= 16:
|
|
state = "match"
|
|
continue
|
|
if t == 0:
|
|
t = ext(0, 15)
|
|
out += src[ip:ip + t + 3]
|
|
ip += t + 3
|
|
state = "first_literal_run"
|
|
continue
|
|
if state == "first_literal_run":
|
|
t = src[ip]
|
|
ip += 1
|
|
if t >= 16:
|
|
state = "match"
|
|
continue
|
|
dist = 1 + M2_MAX_OFFSET + (t >> 2) + (src[ip] << 2)
|
|
ip += 1
|
|
_copy_match(out, dist, 3)
|
|
state = "match_done"
|
|
continue
|
|
if state == "match_next":
|
|
out += src[ip:ip + t]
|
|
ip += t
|
|
t = src[ip]
|
|
ip += 1
|
|
state = "match"
|
|
continue
|
|
if state == "match":
|
|
if t >= 64: # M2
|
|
dist = 1 + ((t >> 2) & 7) + (src[ip] << 3)
|
|
ip += 1
|
|
_copy_match(out, dist, (t >> 5) + 1)
|
|
elif t >= 32: # M3
|
|
t &= 31
|
|
if t == 0:
|
|
t = ext(0, 31)
|
|
dist = 1 + ((src[ip] | (src[ip + 1] << 8)) >> 2)
|
|
ip += 2
|
|
_copy_match(out, dist, t + 2)
|
|
elif t >= 16: # M4 (distance 0 = end of stream)
|
|
hi = (t & 8) << 11
|
|
t &= 7
|
|
if t == 0:
|
|
t = ext(0, 7)
|
|
dist = hi + ((src[ip] | (src[ip + 1] << 8)) >> 2)
|
|
ip += 2
|
|
if dist == 0:
|
|
break
|
|
_copy_match(out, dist + 0x4000, t + 2)
|
|
else: # M1 after a match: 2 bytes
|
|
dist = 1 + (t >> 2) + (src[ip] << 2)
|
|
ip += 1
|
|
_copy_match(out, dist, 2)
|
|
state = "match_done"
|
|
continue
|
|
# match_done: the low 2 bits of the instruction's last byte = trailing literals
|
|
t = src[ip - 2] & 3
|
|
if t == 0:
|
|
state = "loop"
|
|
else:
|
|
state = "match_next"
|
|
except IndexError:
|
|
raise LzoError("truncated LZO stream") from None
|
|
if ip != n:
|
|
raise LzoError("%d bytes after the end of the LZO stream" % (n - ip))
|
|
if out_len is not None and len(out) != out_len:
|
|
raise LzoError("LZO block decoded to %d bytes, expected %d" % (len(out), out_len))
|
|
return bytes(out)
|
|
|
|
|
|
EOF = b"\x11\x00\x00"
|
|
|
|
|
|
def compress_stored(data):
|
|
"""bytes -> LZO1X stream made of one literal run (decodes back to `data`)."""
|
|
data = bytes(data)
|
|
n = len(data)
|
|
if n == 0:
|
|
return EOF
|
|
if n <= 238:
|
|
return bytes([n + 17]) + data + EOF
|
|
# long literal run: 0x00, k zero bytes (255 each), last byte; length = 18 + 255*k + last
|
|
r = n - 18
|
|
k = (r - 1) // 255
|
|
last = r - 255 * k
|
|
return b"\x00" + b"\x00" * k + bytes([last]) + data + EOF
|