I'm currently developing an Open Source project called sims3py using Python2.7. One overarching goal of the project is to ensure a Pure Python implementation of functions currently available only through a .Net-based library.
Because the project must be Pure Python, 'cheats' using C code is not allowed (except for the implicit C code invoked by standard Python libraries).
I have performed profiling on my code, and discovered that a certain function, decompress()
was very slow. This function is called thousands of time by many tools I created based on the sims3py library. So naturally this is a right place to perform optimization.
The decompression algorithm is documented here.
(Please note that I've performed profiling; the times listed below are time taken only for the decompression part. More specifically, time taken to perform chunk = decompress(buff)
from another place in the program. They are not the total time taken for the program to run. Just for the decompress()
function).
The first iteration uses a custom class called SlidingUnpacker
. It works, but it's darned slow. When this first iteration was used against a test file (containing about 6000 chunks to decompress), the decompression of elements in that file took 240+ seconds.
The second iteration, which is the current Public version, dispenses with SlidingUnpacker
and accesses the blob to decompress directly (via memoryview()
), and it's much faster; testing against the same test file shows that the second iteration took about 75 seconds.
After reviewing how the code works, I've optimized decompress()
further as follows:
def decompress(byte_buffer, strict_size=True, ignore_extra=True):
"""
Performs a decompression on a memoryview(byte_buffer).
:param byte_buffer: byte_buffer containing the compressed resource
:type byte_buffer: memoryview
:param strict_size: Whether to raise exception on decompressed size mismatch or not
:type strict_size: bool
:param ignore_extra: Whether to ignore compressed bytes beyond specified fullsize
:type ignore_extra: bool
:return: a bytes() containing the uncompressed resource
:rtype: bytes
"""
assert isinstance(byte_buffer, memoryview)
with io.BytesIO() as output:
buf = byte_buffer
comptype, magic = map(ord, buf[0:2])
if magic != 0xFB:
# Not a valid compression format
raise InvalidCompressionException(message='NoMagic, the Magic byte 0xFB not found!')
s1 = 0
if comptype & 0x80:
s1, s2, s3, s4 = map(ord, buf[2:6])
pos = 6
else:
s2, s3, s4 = map(ord, buf[2:5])
pos = 5
fullsize = (s1 << 24) | (s2 << 16) | (s3 << 8) | s4
output_len = 0
control_0, control_1, control_2, control_3 = 0, 0, 0, 0
try:
while True:
if output_len >= fullsize and ignore_extra:
break
# The following is a sentinel. If buf[pos] results in IndexError, control_0 stays None and we can
# detect end_of_buffer without 'hanging' control bytes
control_0 = None
control_0 = ord(buf[pos])
pos += 1
# 0x00 ~ 0x7F
if (control_0 & 0x80) == 0:
control_1 = ord(buf[pos])
pos += 1
num_plain = control_0 & 0x03
num_copy = ((control_0 & 0x1c) >> 2) + 3
copy_offset = ((control_0 & 0x60) << 3) + control_1 + 1
# 0x80 ~ 0xBF
elif (control_0 & 0xc0) == 0x80:
control_1 = ord(buf[pos])
control_2 = ord(buf[pos + 1])
pos += 2
num_plain = ((control_1 & 0xC0) >> 6) & 0x03
num_copy = (control_0 & 0x3F) + 4
copy_offset = ((control_1 & 0x3F) << 8) + control_2 + 1
# 0xC0 ~ 0xDF
elif (control_0 & 0xE0) == 0xC0:
control_1 = ord(buf[pos])
control_2 = ord(buf[pos + 1])
control_3 = ord(buf[pos + 2])
pos += 3
num_plain = control_0 & 0x03
num_copy = ((control_0 & 0x0C) << 6) + control_3 + 5
copy_offset = ((control_0 & 0x10) << 12) + (control_1 << 8) + control_2 + 1
# 0xE0 ~ 0xFB
elif 0xE0 <= control_0 <= 0xFB:
num_plain = ((control_0 & 0x1F) << 2) + 4
npos = pos + num_plain
output.write(buf[pos:npos])
pos = npos
output_len += num_plain
continue
# 0xFC ~ 0xFF
else:
num_plain = control_0 & 0x03
output.write(buf[pos:pos + num_plain])
pos += num_plain
output_len += num_plain
break
npos = pos + num_plain
output.write(buf[pos:npos])
pos = npos
output_len += num_plain
copy_pos = output_len - copy_offset
# Pre-add num_copy to output_len because we WILL get num_copy bytes anyways... or Error trying
output_len += num_copy
# We do not use for: loop, because for: loop handles bytes one-by-one.
# This construct tries to read as many bytes as possible per iteration
while num_copy:
output.seek(copy_pos)
to_copy = output.read(num_copy)
le = len(to_copy)
if le:
output.seek(0, 2) # Seek to end of stream
output.write(to_copy)
copy_pos += le
num_copy -= le
else:
raise IndexError('There should be at least 1 char but got none.')
except IndexError: # This will be raised by buf[] if we try reading beyond its bounds
if control_0 is not None:
# Exception raised if buffer is exhausted while algorithm still requires a control byte, or
# control bytes specified a number of plain data to consume but the buffer exhausted before the
# required number of bytes are received
raise InvalidCompressionException(
message='Truncated or corrupt resource, buffer exhausted after reading {0} bytes'.format(pos)
)
# If we reach this point, this means that the compression structure has been decompressed successfully
# although without 'end of compression' control (0xFC~0xFF), AND before encoded fullsize is reached.
# Because technically we don't find any errors in the compressed structure, we do not do anything, letting
# an external sanity check to decide.
# (This situation is situation (3) as described in the sanity check's comments)
sys.exc_clear()
pass
finally:
pass
# We reach this point only if one these are true:
# (1) len(output) >= fullsize (while ignore_extra == True)
# (2) End of compression control detected (0xFC ~ 0xFF)
# (3) byte_buffer has been exhausted before (1) or (2) reached
# In any case, all compression controls have been decoded properly (i.e., no incomplete control codes and/or
# truncated data needed by control codes). So, technically the compressed data was NOT corrupt.
# What we do depends on whether strict flag is set or not.
if strict_size and fullsize != output_len:
raise InvalidCompressionException(message='Size mismatch, want {0} got {1}'.format(fullsize, output_len))
return output.getvalue()
This final iteration indeed improves the decompression. Testing the same test file, decompression now takes only about 63 seconds.
However, I still feel the code can be optimized further. Can you provide suggestions as to how I can further optimize the code?
Note: I've made available test vectors for this function.
After unzipping the .7z archive, there should be a pair of files: testblob_compressed.bin
and testblob_uncompressed.xml
. The following code should be enough to test:
filename = 'testblob_compressed.bin'
ba = bytearray(os.path.getsize(filename))
with open(filename, 'rb') as fin:
fin.readinto(ba)
mv = memoryview(ba)
output = decompress(mv)
(output
should be byte-identical with testblob_uncompressed.xml
)
UPDATE
Just in case you're wondering, here's the final version of the decompress()
function:
https://bitbucket.org/pepoluan/sims3py/src/6f97b77fd4b12a4d294cd4a904742072e09a2747/sims3py/init.py
Thanks to everyone pitching in, especially @Veedrac !
testblob_compress.bin
decompresses in less than 1s. \$\endgroup\$