Optimizing critical loop for consuming a byte-buffer

Question

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 !

Answer 1

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)
    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.')

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
        raise InvalidCompressionException(message='NoMagic, the Magic byte 0xFB not found!')

    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.