From 31ad2856fd877b1d1d898cec31ad63554cc92a4e Mon Sep 17 00:00:00 2001 From: Jennifer Taylor Date: Sat, 28 Nov 2020 19:05:27 +0000 Subject: [PATCH] Add a much faster C++ implementation of Lz77 decompress. We can achieve approximately a 10x speedup decompressing with this enabled. --- bemani/protocol/lz77.cpp | 100 ++++++++++++++++++++++++++++++++++++++ bemani/protocol/lz77.py | 37 +++++++++++++- bemani/tests/test_lz77.py | 8 +-- setup.py | 56 +++++++++++++++++---- 4 files changed, 186 insertions(+), 15 deletions(-) create mode 100644 bemani/protocol/lz77.cpp diff --git a/bemani/protocol/lz77.cpp b/bemani/protocol/lz77.cpp new file mode 100644 index 0000000..d4e08c5 --- /dev/null +++ b/bemani/protocol/lz77.cpp @@ -0,0 +1,100 @@ +#include +#include + +#define FLAG_COPY 1 +#define FLAG_BACKREF 0 + +extern "C" +{ + int decompress(uint8_t *indata, unsigned int inlen, uint8_t *outdata, unsigned int outlen) + { + // First, let's assume a worst case compression which in theory is just a copy. + // The math is basically 9 bytes used for every 8 bytes. So, the minimum output + // buffer we need is (inlen * 8/9). If we have an outlen smaller than that, we + // are hosed. + if (outlen < ((inlen * 8) / 9)) + { + // We cannot decompress, we will run out of room! + return -1; + } + + // Now, let's enter a loop where we read control bytes and act on them. + unsigned int inloc = 0; + unsigned int outloc = 0; + bool eof = false; + while (inloc < inlen && !eof) + { + if (inloc >= inlen) + { + // We failed to decompress, we overran the input buffer. + return -2; + } + + uint8_t flags = indata[inloc++]; + for (unsigned int flagpos = 0; flagpos < 8; flagpos++) + { + if (((flags >> flagpos) & 1) == FLAG_COPY) + { + // Copy a byte, move on + if (inloc >= inlen) + { + // We failed to decompress, we overran the input buffer. + return -2; + } + if (outloc >= outlen) + { + // We overwrote our output buffer, we probably corrupted memory somewhere. + return -3; + } + outdata[outloc++] = indata[inloc++]; + } + else + { + // Backref copy + if (inloc >= (inlen - 1)) + { + // We failed to decompress, we overran the input buffer. + return -2; + } + uint8_t hi = indata[inloc++]; + uint8_t lo = indata[inloc++]; + + unsigned int copy_len = (lo & 0xF) + 3; + unsigned int copy_pos = (hi << 4) | (lo >> 4); + + if (copy_pos == 0) + { + // This is the end of a file. + eof = true; + break; + } + + // Copy backref a byte at a time. This is because a backref can stick + // out into as-of-yet uncopied data in order to reference what we're + // about to write. + for (unsigned int backref_copy_amt = 0; backref_copy_amt < copy_len; backref_copy_amt++) + { + if (outloc >= outlen) + { + // We overwrote our output buffer, we probably corrupted memory somewhere. + return -3; + } + + if (((int)outloc - (int)copy_pos) < 0) + { + outdata[outloc++] = 0; + } + else + { + outdata[outloc] = outdata[outloc - copy_pos]; + outloc++; + } + } + } + } + } + + // Update the outlen with the actual data length. + return outloc; + } +} diff --git a/bemani/protocol/lz77.py b/bemani/protocol/lz77.py index c36f0b7..e491497 100644 --- a/bemani/protocol/lz77.py +++ b/bemani/protocol/lz77.py @@ -1,7 +1,22 @@ +import ctypes +import os from collections import defaultdict from typing import Generator, List, MutableMapping, Optional, Set, Tuple +# Attempt to use the faster C++ libraries if they're available +try: + clib = None + clib_path = os.path.dirname(os.path.abspath(__file__)) + files = [f for f in os.listdir(clib_path) if f.startswith("lz77alt") and f.endswith(".so")] + if len(files) > 0: + clib = ctypes.cdll.LoadLibrary(os.path.join(clib_path, files[0])) + clib.decompress.argtypes = (ctypes.c_char_p, ctypes.c_int, ctypes.c_char_p, ctypes.c_int) + clib.decompress.restype = ctypes.c_int +except Exception: + clib = None + + class LzException(Exception): """ An exception thrown when we encounter an error with Lz77 encoding/decoding. @@ -414,8 +429,26 @@ class Lz77: Returns: Raw binary data. """ - lz = Lz77Decompress(data, backref=self.backref) - return b''.join(lz.decompress_bytes()) + if clib is not None: + # Given a maximum backref length of 18, if we had a file that was + # only backrefs of maximum size. We would get a compression of around + # (18 * 8) / (2 * 8 + 1), or 8.47. So, allocate 9 times in the output + # buffer. + outbuf = b"\0" * (len(data) * 9) + result = clib.decompress(data, len(data), outbuf, len(outbuf)) + if result >= 0: + return outbuf[:result] + elif result == -1: + raise LzException("Not enough room in output buffer!") + elif result == -2: + raise LzException("Unexpected EOF during decompression!") + elif result == -3: + raise LzException("Not enough room to write output byte!") + else: + raise LzException("Unknown exception in C++ code!") + else: + lz = Lz77Decompress(data, backref=self.backref) + return b''.join(lz.decompress_bytes()) def compress(self, data: bytes) -> bytes: """ diff --git a/bemani/tests/test_lz77.py b/bemani/tests/test_lz77.py index 3e3b2b8..e00139f 100644 --- a/bemani/tests/test_lz77.py +++ b/bemani/tests/test_lz77.py @@ -12,7 +12,7 @@ class TestLZ77Decompressor(unittest.TestCase): for _ in range(100): amount = random.randint(1, Lz77Decompress.RING_LENGTH) - data = bytes([random.randint(0, 255) for _ in range(amount)]) + data = os.urandom(amount) # Save our ring position, write a chunk of data readpos = dec.write_pos @@ -35,7 +35,7 @@ def get_fixture(name: str) -> bytes: class TestLz77RealCompressor(unittest.TestCase): def test_small_data_random(self) -> None: lz77 = Lz77() - data = bytes([random.randint(0, 255) for _ in range(1 * 1024)]) + data = os.urandom(1 * 1024) compresseddata = lz77.compress(data) self.assertNotEqual(data, compresseddata) @@ -45,7 +45,7 @@ class TestLz77RealCompressor(unittest.TestCase): def test_large_data_random(self) -> None: lz77 = Lz77() - data = bytes([random.randint(0, 255) for _ in range(100 * 1024)]) + data = os.urandom(100 * 1024) compresseddata = lz77.compress(data) self.assertNotEqual(data, compresseddata) @@ -55,7 +55,7 @@ class TestLz77RealCompressor(unittest.TestCase): def test_huge_data_random(self) -> None: lz77 = Lz77() - data = bytes([random.randint(0, 255) for _ in range(1 * 1024 * 1024)]) + data = os.urandom(1 * 1024 * 1024) compresseddata = lz77.compress(data) self.assertNotEqual(data, compresseddata) diff --git a/setup.py b/setup.py index ac01823..3f4a126 100644 --- a/setup.py +++ b/setup.py @@ -1,4 +1,4 @@ -from setuptools import setup +from setuptools import Extension, setup from Cython.Build import cythonize setup( @@ -69,16 +69,54 @@ setup( install_requires=[ req for req in open('requirements.txt').read().split('\n') if len(req) > 0 ], + # None of these are required for operating any of the utilities found in this repo. + # They are all present for speed. If you cannot compile arbitrary code or cython, + # remove the ext_modules line and run setuptools again. Everything should work, but + # it will run slower. ext_modules=cythonize( - # Compile various low-level portions of the protocol to get a massive - # per-packet speed boost. [ - "bemani/protocol/binary.py", - "bemani/protocol/lz77.py", - "bemani/protocol/node.py", - "bemani/protocol/protocol.py", - "bemani/protocol/stream.py", - "bemani/protocol/xml.py", + Extension( + "bemani.protocol.binary", + [ + "bemani/protocol/binary.py", + ] + ), + Extension( + "bemani.protocol.lz77", + [ + "bemani/protocol/lz77.py", + ] + ), + Extension( + "bemani.protocol.lz77alt", + [ + "bemani/protocol/lz77.cpp", + ] + ), + Extension( + "bemani.protocol.node", + [ + "bemani/protocol/node.py", + ] + ), + Extension( + "bemani.protocol.protocol", + [ + "bemani/protocol/protocol.py", + ] + ), + Extension( + "bemani.protocol.stream", + [ + "bemani/protocol/stream.py", + ] + ), + Extension( + "bemani.protocol.xml", + [ + "bemani/protocol/xml.py", + ] + ), ], language_level=3, ),