ref: d4a6a1b2bcf94915020f2a6c96b8004479185325
dir: /tools/pokemontools/lz.py/
# -*- coding: utf-8 -*-
"""
Pokemon Crystal data de/compression.
"""
"""
A rundown of Pokemon Crystal's compression scheme:
Control commands occupy bits 5-7.
Bits 0-4 serve as the first parameter <n> for each command.
"""
lz_commands = {
'literal': 0, # n values for n bytes
'iterate': 1, # one value for n bytes
'alternate': 2, # alternate two values for n bytes
'blank': 3, # zero for n bytes
}
"""
Repeater commands repeat any data that was just decompressed.
They take an additional signed parameter <s> to mark a relative starting point.
These wrap around (positive from the start, negative from the current position).
"""
lz_commands.update({
'repeat': 4, # n bytes starting from s
'flip': 5, # n bytes in reverse bit order starting from s
'reverse': 6, # n bytes backwards starting from s
})
"""
The long command is used when 5 bits aren't enough. Bits 2-4 contain a new control code.
Bits 0-1 are appended to a new byte as 8-9, allowing a 10-bit parameter.
"""
lz_commands.update({
'long': 7, # n is now 10 bits for a new control code
})
max_length = 1 << 10 # can't go higher than 10 bits
lowmax = 1 << 5 # standard 5-bit param
"""
If 0xff is encountered instead of a command, decompression ends.
"""
lz_end = 0xff
bit_flipped = [
sum(((byte >> i) & 1) << (7 - i) for i in xrange(8))
for byte in xrange(0x100)
]
class Compressed:
"""
Usage:
lz = Compressed(data).output
or
lz = Compressed().compress(data)
or
c = Compressed()
c.data = data
lz = c.compress()
There are some issues with reproducing the target compressor.
Some notes are listed here:
- the criteria for detecting a lookback is inconsistent
- sometimes lookbacks that are mostly 0s are pruned, sometimes not
- target appears to skip ahead if it can use a lookback soon, stopping the current command short or in some cases truncating it with literals.
- this has been implemented, but the specifics are unknown
- self.min_scores: It's unknown if blank's minimum score should be 1 or 2. Most likely it's 1, with some other hack to account for edge cases.
- may be related to the above
- target does not appear to compress backwards
"""
def __init__(self, *args, **kwargs):
self.min_scores = {
'blank': 1,
'iterate': 2,
'alternate': 3,
'repeat': 3,
'reverse': 3,
'flip': 3,
}
self.preference = [
'repeat',
'blank',
'flip',
'reverse',
'iterate',
'alternate',
#'literal',
]
self.lookback_methods = 'repeat', 'reverse', 'flip'
self.__dict__.update({
'data': None,
'commands': lz_commands,
'debug': False,
'literal_only': False,
})
self.arg_names = 'data', 'commands', 'debug', 'literal_only'
self.__dict__.update(kwargs)
self.__dict__.update(dict(zip(self.arg_names, args)))
if self.data is not None:
self.compress()
def compress(self, data=None):
if data is not None:
self.data = data
self.data = list(bytearray(self.data))
self.indexes = {}
self.lookbacks = {}
for method in self.lookback_methods:
self.lookbacks[method] = {}
self.address = 0
self.end = len(self.data)
self.output = []
self.literal = None
while self.address < self.end:
if self.score():
self.do_literal()
self.do_winner()
else:
if self.literal == None:
self.literal = self.address
self.address += 1
self.do_literal()
self.output += [lz_end]
return self.output
def reset_scores(self):
self.scores = {}
self.offsets = {}
self.helpers = {}
for method in self.min_scores.iterkeys():
self.scores[method] = 0
def bit_flip(self, byte):
return bit_flipped[byte]
def do_literal(self):
if self.literal != None:
length = abs(self.address - self.literal)
start = min(self.literal, self.address + 1)
self.helpers['literal'] = self.data[start:start+length]
self.do_cmd('literal', length)
self.literal = None
def score(self):
self.reset_scores()
map(self.score_literal, ['iterate', 'alternate', 'blank'])
for method in self.lookback_methods:
self.scores[method], self.offsets[method] = self.find_lookback(method, self.address)
self.stop_short()
return any(
score
> self.min_scores[method] + int(score > lowmax)
for method, score in self.scores.iteritems()
)
def stop_short(self):
"""
If a lookback is close, reduce the scores of other commands.
"""
best_method, best_score = max(
self.scores.items(),
key = lambda x: (
x[1],
-self.preference.index(x[0])
)
)
for method in self.lookback_methods:
min_score = self.min_scores[method]
for address in xrange(self.address+1, self.address+best_score):
length, index = self.find_lookback(method, address)
if length > max(min_score, best_score):
# BUG: lookbacks can reduce themselves. This appears to be a bug in the target also.
for m, score in self.scores.items():
self.scores[m] = min(score, address - self.address)
def read(self, address=None):
if address is None:
address = self.address
if 0 <= address < len(self.data):
return self.data[address]
return None
def find_all_lookbacks(self):
for method in self.lookback_methods:
for address, byte in enumerate(self.data):
self.find_lookback(method, address)
def find_lookback(self, method, address=None):
"""Temporarily stubbed, because the real function doesn't run in polynomial time."""
return 0, None
def broken_find_lookback(self, method, address=None):
if address is None:
address = self.address
existing = self.lookbacks.get(method, {}).get(address)
if existing != None:
return existing
lookback = 0, None
# Better to not carelessly optimize at the moment.
"""
if address < 2:
return lookback
"""
byte = self.read(address)
if byte is None:
return lookback
direction, mutate = {
'repeat': ( 1, int),
'reverse': (-1, int),
'flip': ( 1, self.bit_flip),
}[method]
# Doesn't seem to help
"""
if mutate == self.bit_flip:
if byte == 0:
self.lookbacks[method][address] = lookback
return lookback
"""
data_len = len(self.data)
is_two_byte_index = lambda index: int(index < address - 0x7f)
for index in self.get_indexes(mutate(byte)):
if index >= address:
break
old_length, old_index = lookback
if direction == 1:
if old_length > data_len - index: break
else:
if old_length > index: continue
if self.read(index) in [None]: continue
length = 1 # we know there's at least one match, or we wouldn't be checking this index
while 1:
this_byte = self.read(address + length)
that_byte = self.read(index + length * direction)
if that_byte == None or this_byte != mutate(that_byte):
break
length += 1
score = length - is_two_byte_index(index)
old_score = old_length - is_two_byte_index(old_index)
if score >= old_score or (score == old_score and length > old_length):
# XXX maybe avoid two-byte indexes when possible
if score >= lookback[0] - is_two_byte_index(lookback[1]):
lookback = length, index
self.lookbacks[method][address] = lookback
return lookback
def get_indexes(self, byte):
if not self.indexes.has_key(byte):
self.indexes[byte] = []
index = -1
while 1:
try:
index = self.data.index(byte, index + 1)
except ValueError:
break
self.indexes[byte].append(index)
return self.indexes[byte]
def score_literal(self, method):
address = self.address
compare = {
'blank': [0],
'iterate': [self.read(address)],
'alternate': [self.read(address), self.read(address + 1)],
}[method]
# XXX may or may not be correct
if method == 'alternate' and compare[0] == 0:
return
length = 0
while self.read(address + length) == compare[length % len(compare)]:
length += 1
self.scores[method] = length
self.helpers[method] = compare
def do_winner(self):
winners = filter(
lambda (method, score):
score
> self.min_scores[method] + int(score > lowmax),
self.scores.iteritems()
)
winners.sort(
key = lambda (method, score): (
-(score - self.min_scores[method] - int(score > lowmax)),
self.preference.index(method)
)
)
winner, score = winners[0]
length = min(score, max_length)
self.do_cmd(winner, length)
self.address += length
def do_cmd(self, cmd, length):
start_address = self.address
cmd_length = length - 1
output = []
if length > lowmax:
output.append(
(self.commands['long'] << 5)
+ (self.commands[cmd] << 2)
+ (cmd_length >> 8)
)
output.append(
cmd_length & 0xff
)
else:
output.append(
(self.commands[cmd] << 5)
+ cmd_length
)
self.helpers['blank'] = [] # quick hack
output += self.helpers.get(cmd, [])
if cmd in self.lookback_methods:
offset = self.offsets[cmd]
# Negative offsets are one byte.
# Positive offsets are two.
if 0 < start_address - offset - 1 <= 0x7f:
offset = (start_address - offset - 1) | 0x80
output += [offset]
else:
output += [offset / 0x100, offset % 0x100] # big endian
if self.debug:
print ' '.join(map(str, [
cmd, length, '\t',
' '.join(map('{:02x}'.format, output)),
self.data[start_address:start_address+length] if cmd in self.lookback_methods else '',
]))
self.output += output
class Decompressed:
"""
Interpret and decompress lz-compressed data, usually 2bpp.
"""
"""
Usage:
data = Decompressed(lz).output
or
data = Decompressed().decompress(lz)
or
d = Decompressed()
d.lz = lz
data = d.decompress()
To decompress from offset 0x80000 in a rom:
data = Decompressed(rom, start=0x80000).output
"""
lz = None
start = 0
commands = lz_commands
debug = False
arg_names = 'lz', 'start', 'commands', 'debug'
def __init__(self, *args, **kwargs):
self.__dict__.update(dict(zip(self.arg_names, args)))
self.__dict__.update(kwargs)
self.command_names = dict(map(reversed, self.commands.items()))
self.address = self.start
if self.lz is not None:
self.decompress()
if self.debug: print self.command_list()
def command_list(self):
"""
Print a list of commands that were used. Useful for debugging.
"""
text = ''
output_address = 0
for name, attrs in self.used_commands:
length = attrs['length']
address = attrs['address']
offset = attrs['offset']
direction = attrs['direction']
text += '{2:03x} {0}: {1}'.format(name, length, output_address)
text += '\t' + ' '.join(
'{:02x}'.format(int(byte))
for byte in self.lz[ address : address + attrs['cmd_length'] ]
)
if offset is not None:
repeated_data = self.output[ offset : offset + length * direction : direction ]
if name == 'flip':
repeated_data = map(bit_flipped.__getitem__, repeated_data)
text += ' [' + ' '.join(map('{:02x}'.format, repeated_data)) + ']'
text += '\n'
output_address += length
return text
def decompress(self, lz=None):
if lz is not None:
self.lz = lz
self.lz = bytearray(self.lz)
self.used_commands = []
self.output = []
while 1:
cmd_address = self.address
self.offset = None
self.direction = None
if (self.byte == lz_end):
self.next()
break
self.cmd = (self.byte & 0b11100000) >> 5
if self.cmd_name == 'long':
# 10-bit length
self.cmd = (self.byte & 0b00011100) >> 2
self.length = (self.next() & 0b00000011) * 0x100
self.length += self.next() + 1
else:
# 5-bit length
self.length = (self.next() & 0b00011111) + 1
self.__class__.__dict__[self.cmd_name](self)
self.used_commands += [(
self.cmd_name,
{
'length': self.length,
'address': cmd_address,
'offset': self.offset,
'cmd_length': self.address - cmd_address,
'direction': self.direction,
}
)]
# Keep track of the data we just decompressed.
self.compressed_data = self.lz[self.start : self.address]
@property
def byte(self):
return self.lz[ self.address ]
def next(self):
byte = self.byte
self.address += 1
return byte
@property
def cmd_name(self):
return self.command_names.get(self.cmd)
def get_offset(self):
if self.byte >= 0x80: # negative
# negative
offset = self.next() & 0x7f
offset = len(self.output) - offset - 1
else:
# positive
offset = self.next() * 0x100
offset += self.next()
self.offset = offset
def literal(self):
"""
Copy data directly.
"""
self.output += self.lz[ self.address : self.address + self.length ]
self.address += self.length
def iterate(self):
"""
Write one byte repeatedly.
"""
self.output += [self.next()] * self.length
def alternate(self):
"""
Write alternating bytes.
"""
alts = [self.next(), self.next()]
self.output += [ alts[x & 1] for x in xrange(self.length) ]
def blank(self):
"""
Write zeros.
"""
self.output += [0] * self.length
def flip(self):
"""
Repeat flipped bytes from output.
Example: 11100100 -> 00100111
"""
self._repeat(table=bit_flipped)
def reverse(self):
"""
Repeat reversed bytes from output.
"""
self._repeat(direction=-1)
def repeat(self):
"""
Repeat bytes from output.
"""
self._repeat()
def _repeat(self, direction=1, table=None):
self.get_offset()
self.direction = direction
# Note: appends must be one at a time (this way, repeats can draw from themselves if required)
for i in xrange(self.length):
byte = self.output[ self.offset + i * direction ]
self.output.append( table[byte] if table else byte )