# Optimizing critical loop for consuming a byte-buffer

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

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)
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:
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 !

• Can you please link to the input/output that is taking 63s? testblob_compress.bin decompresses in less than 1s. Feb 21 '15 at 1:02
• @mjolka that was a test file containing 8013 chunks of which 6000-ish are compressed. There's another class that parses this file, and whenever encountering a compressed chunk, hands the chunk to decompress(). I can provide the file, but then you'll need to run the whole program instead of just focusing on the function. Feb 21 '15 at 2:50

The lazy thing to note is that pypy runs this in about 20% of the time, so PyPy should be preferred if possible.

Your with io.BytesIO() as output: can safely cover a smaller fraction of the code, so I suggest moving as much as possible (within reason) out of its context.

You have

    except IndexError:
... # stuff
pass

finally:
pass


This should just be

    except IndexError:
... # stuff


I don't get your justification for running sys.exc_clear. I suggest you make sure this is really the right thing to do, because it looks wrong.

You have

        while True:
if output_len >= fullsize and ignore_extra:
break


This looks like it would better be written

        while not (output_len >= fullsize and ignore_extra):


It seems to me that your while loop:

while num_copy:
output.seek(copy_pos)
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.')


is not needed:

If the argument is positive, and the underlying raw stream is not interactive, multiple raw reads may be issued to satisfy the byte count (unless EOF is reached first).

A simple assert should be fine.

Your output.seek(0, 2) should be output.seek(0, io.SEEK_END).

You might find things easier if you use a bytearray over a memoryview as you can avoid all of the ord calls.

output can also be a bytearray. I find this gives a significant speed improvement.

You spend a lot of upkeep on output_len; there's no real harm in using len(output), so I suggest you do so.

There seems to be no good reason for this line:

control_0, control_1, control_2, control_3 = 0, 0, 0, 0


so remove it.

It looks to me like the if control_0 is not None: in the except IndexError can be replaced with a pos < len(buf) in the while and the try can be moved in. This does come at a slight speed cost so I avoided the change.

Some of your bit twiddling can be simplified. The ifs:

# 0x00 ~ 0x7F
if control_0 < 0x80:
...

# 0x80 ~ 0xBF
elif control_0 < 0xX0:
...

# 0xC0 ~ 0xDF
elif control_0 < 0xE0:
...

# 0xE0 ~ 0xFB
elif control_0 < 0xFC:
...

# 0xFC ~ 0xFF
else:
...


Your

num_plain = ((control_1 & 0b11000000) >> 6) & 0b11


can be just

num_plain = (control_1 >> 6) & 0b11


I don't see how to speed it up further, but this should be a 2x improvement when staying on CPython and a 10x improvement if moving to PyPy (an extra 5x from the better interpreter).

def decompress(byte_buffer, strict_size=True, ignore_extra=True):
"""
Performs a decompression on a bytearray(byte_buffer).

:param byte_buffer: byte_buffer containing the compressed resource
:type byte_buffer: bytearray
: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, bytearray)

buf = byte_buffer

comptype, magic = buf[0:2]
if magic != 0xFB:
# Not a valid compression format

s1 = 0
if comptype & 0x80:
s1, s2, s3, s4 = buf[2:6]
pos = 6
else:
s2, s3, s4 = buf[2:5]
pos = 5
fullsize = (s1 << 24) | (s2 << 16) | (s3 << 8) | s4

output = bytearray()

# If the compression structure has been decompressed successfully although without 'end of compression'
# control (0xFC~0xFF), AND before encoded fullsize is reached, the pos < len(buf) condition will
# break the loop.
# 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)
try:
while not (len(output) >= fullsize and ignore_extra):
# The while ensures that buf[pos] always valid
control_0 = None
control_0 = buf[pos]

pos += 1

# 0x00 ~ 0x7F
if control_0 < 0x80:
control_1 = buf[pos]
pos += 1

num_plain = control_0 & 0b11
num_copy = ((control_0 >> 2) & 0b111) + 3
copy_offset = ((control_0 & 0b1100000) << 3) + control_1 + 1

# 0x80 ~ 0xBF
elif control_0 < 0xC0:
control_1 = buf[pos]
control_2 = buf[pos + 1]
pos += 2

num_plain = (control_1 >> 6) & 0b11
num_copy = (control_0 & 0b111111) + 4
copy_offset = ((control_1 & 0b111111) << 8) + control_2 + 1

# 0xC0 ~ 0xDF
elif control_0 < 0xE0:
control_1 = buf[pos]
control_2 = buf[pos + 1]
control_3 = buf[pos + 2]
pos += 3

num_plain = control_0 & 0b11
num_copy = ((control_0 & 0b1100) << 6) + control_3 + 5
copy_offset = ((control_0 & 0b10000) << 12) + (control_1 << 8) + control_2 + 1

# 0xE0 ~ 0xFB
elif control_0 < 0xFC:
num_plain = ((control_0 & 0b11111) << 2) + 4

output += buf[pos:pos + num_plain]
pos += num_plain
continue

# 0xFC ~ 0xFF
else:
num_plain = control_0 & 0b11

output += buf[pos:pos + num_plain]
pos += num_plain
break

output += buf[pos:pos + num_plain]
pos += num_plain

# Don't use negative indices lest the addition makes the end point 0
copy_pos = len(output) - copy_offset

if copy_pos < 0:
raise IndexError('There should be {} values in buffer, got {}.'.format(num_copy, len(output)))

to_copy = output[copy_pos:copy_pos + num_copy]
output += to_copy

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)

# 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 != len(output):
raise InvalidCompressionException(message='Size mismatch, want {} got {}'.format(fullsize, len(output)))

return output


One last thing that seems to help is removing the control_1, control_2 and control_3 intermediates:

# 0x00 ~ 0x7F
if control_0 < 0x80:
num_plain = control_0 & 0b11
num_copy = ((control_0 >> 2) & 0b111) + 3
copy_offset = ((control_0 & 0b1100000) << 3) + buf[pos] + 1
pos += 1

# 0x80 ~ 0xBF
elif control_0 < 0xC0:
num_plain = (buf[pos] >> 6) & 0b11
num_copy = (control_0 & 0b111111) + 4
copy_offset = ((buf[pos] & 0b111111) << 8) + buf[pos + 1] + 1
pos += 2

# 0xC0 ~ 0xDF
elif control_0 < 0xE0:
num_plain = control_0 & 0b11
num_copy = ((control_0 & 0b1100) << 6) + buf[pos + 2] + 5
copy_offset = ((control_0 & 0b10000) << 12) + (buf[pos] << 8) + buf[pos + 1] + 1
pos += 3

# 0xE0 ~ 0xFB
elif control_0 < 0xFC:
num_plain = ((control_0 & 0b11111) << 2) + 4

output += buf[pos:pos + num_plain]
pos += num_plain
continue

# 0xFC ~ 0xFF
else:
num_plain = control_0 & 0b11

output += buf[pos:pos + num_plain]
pos += num_plain
break


This actually helps PyPy a lot more than CPython, getting PyPy to a full 20x the speed of the original CPython here.

• The while num_copy is necessary because, according to the specification, it is possible that copy_pos seeks back less than num_copy. E.g., the output is 100 bytes, copy_pos is 96, but num_copy is 10. In this case, I must read 4 bytes to append to output, lengthening output to 104, then another 4 bytes starting position (96+4) lengthening output to 108, then another 2 bytes starting position (96+4+4). Feb 21 '15 at 2:56
• About PyPy, IMO requiring the eventual user to use PyPy would be an unnecessary burden on their part. There's nothing wrong with total execution time of tens of seconds, actually. I'm just looking for code-based ways to shave the time further. Feb 21 '15 at 2:59
• Anyways, you do have very good points there. I'll try implementing your suggestions and get back to you. Feb 21 '15 at 3:05
• I've implemented most of your suggestions... and it's amazing! The time taken by the decompress() gets cut down from 63 to 41!! Great job 👍 Feb 21 '15 at 10:39
• @pepoluan I've added another suggestion. I would also consider at least mentioning PyPy; a 10x speed improvement isn't to be taken lightly ;P. Feb 21 '15 at 13:42