entropy/__init__.py

import struct
import zlib

from mediatok.container.gtkv import ENTROPY_CODEC_ID_MAP
from mediatok.entropy.pack import pack_tokens, unpack_tokens


ZSTD_LEVEL = 3


def entropy_encode(tokens: list[int], codec_id: int, bits: int = 32) -> bytes:
    if codec_id == 0:
        return _raw_encode(tokens, bits)
    elif codec_id == 1:
        from .rans import rans_encode
        return rans_encode(tokens)
    elif codec_id == 2:
        raw = _raw_encode(tokens, bits)
        return zlib.compress(raw, level=ZSTD_LEVEL)
    elif codec_id == 3:
        raw = _raw_encode(tokens, bits)
        return zlib.compress(raw, level=ZSTD_LEVEL)
    elif codec_id == 4:
        from .ans import ans_encode
        return ans_encode(tokens)
    raise ValueError(f"unknown codec: {codec_id}")


def entropy_decode(payload: bytes, codec_id: int, expected_count: int, bits: int = 32) -> list[int]:
    if codec_id == 0:
        return _raw_decode(payload, expected_count, bits)
    elif codec_id == 1:
        from .rans import rans_decode
        return rans_decode(payload)
    elif codec_id == 2:
        raw = zlib.decompress(payload)
        return _raw_decode(raw, expected_count, bits)
    elif codec_id == 3:
        raw = zlib.decompress(payload)
        return _raw_decode(raw, expected_count, bits)
    elif codec_id == 4:
        from .ans import ans_decode
        return ans_decode(payload)
    raise ValueError(f"unknown codec: {codec_id}")


def _raw_encode(tokens: list[int], bits: int = 32) -> bytes:
    if bits == 18:
        return bytes(pack_tokens(tokens, 18))
    data = bytearray()
    for t in tokens:
        while t >= 128:
            data.append((t & 127) | 128)
            t >>= 7
        data.append(t)
    return bytes(data)


def _raw_decode(data: bytes, expected_count: int, bits: int = 32) -> list[int]:
    if bits == 18:
        return unpack_tokens(list(data), expected_count, 18)
    tokens = []
    i = 0
    while len(tokens) < expected_count and i < len(data):
        val = 0
        shift = 0
        while i < len(data):
            b = data[i]
            i += 1
            val |= (b & 127) << shift
            shift += 7
            if not (b & 128):
                break
        tokens.append(val)
    return tokens