mp_codec.py

"""Chunk-parallel wrapper around the context-mixing codec.

The input is split into fixed-size chunks, each compressed independently in a worker
process. Output layout is deterministic regardless of worker count:

    b'M4' | uvarint(chunk_count) | uvarint(blob_len)*count | blobs...

Chunk size is a fixed constant, so compressed bytes are identical on any machine.
"""
import multiprocessing as mp
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

CHUNK = 1 << 20  # 1 MiB
MAGIC = b'M4'


def _uvarint(value: int) -> bytes:
    out = bytearray()
    while True:
        b = value & 0x7F
        value >>= 7
        if value:
            out.append(b | 0x80)
        else:
            out.append(b)
            return bytes(out)


def _read_uvarint(buf: bytes, pos: int) -> tuple[int, int]:
    shift = 0
    value = 0
    while True:
        b = buf[pos]
        pos += 1
        value |= (b & 0x7F) << shift
        if not b & 0x80:
            return value, pos
        shift += 7


def _compress_chunk(chunk: bytes) -> bytes:
    from codec import compress
    return compress(chunk)


def _decompress_chunk(blob: bytes) -> bytes:
    from codec import decompress
    return decompress(blob)


def _pool_map(fn, items):
    if len(items) <= 1:
        return [fn(item) for item in items]
    workers = min(len(items), os.cpu_count() or 1)
    with mp.get_context('fork').Pool(workers) as pool:
        return pool.map(fn, items, chunksize=1)


def compress(data: bytes) -> bytes:
    chunks = [data[i:i + CHUNK] for i in range(0, len(data), CHUNK)] or [b'']
    blobs = _pool_map(_compress_chunk, chunks)
    head = bytearray(MAGIC)
    head += _uvarint(len(blobs))
    for blob in blobs:
        head += _uvarint(len(blob))
    return bytes(head) + b''.join(blobs)


def decompress(blob: bytes) -> bytes:
    assert blob[:2] == MAGIC, 'bad magic'
    count, pos = _read_uvarint(blob, 2)
    lengths = []
    for _ in range(count):
        n, pos = _read_uvarint(blob, pos)
        lengths.append(n)
    blobs = []
    for n in lengths:
        blobs.append(blob[pos:pos + n])
        pos += n
    return b''.join(_pool_map(_decompress_chunk, blobs))