Add a much faster C++ implementation of Lz77 decompress. We can achieve approximately a 10x speedup decompressing with this enabled.

bemani/protocol/lz77.cpp

@@ -0,0 +1,100 @@
+#include <stdio.h>
+#include <stdint.h>
+
+#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;
+    }
+}

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:
         """

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)

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,
     ),