6
\$\begingroup\$

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 !

\$\endgroup\$
  • 1
    \$\begingroup\$ Can you please link to the input/output that is taking 63s? testblob_compress.bin decompresses in less than 1s. \$\endgroup\$ – mjolka Feb 21 '15 at 1:02
  • \$\begingroup\$ @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. \$\endgroup\$ – pepoluan Feb 21 '15 at 2:50
4
\$\begingroup\$

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.

\$\endgroup\$
  • 1
    \$\begingroup\$ 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). \$\endgroup\$ – pepoluan Feb 21 '15 at 2:56
  • \$\begingroup\$ 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. \$\endgroup\$ – pepoluan Feb 21 '15 at 2:59
  • \$\begingroup\$ Anyways, you do have very good points there. I'll try implementing your suggestions and get back to you. \$\endgroup\$ – pepoluan Feb 21 '15 at 3:05
  • \$\begingroup\$ 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 👍 \$\endgroup\$ – pepoluan Feb 21 '15 at 10:39
  • 2
    \$\begingroup\$ @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. \$\endgroup\$ – Veedrac Feb 21 '15 at 13:42

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.