ref: 8cc79c0891c7a4baa046cc0efa90a0244aab4788
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 )