entropy/rans.py

import struct
from collections import Counter

RANS_L = 1 << 31


def _build_freq(tokens):
    freq = Counter(tokens)
    max_sym = max(tokens) if tokens else 0
    cum = [0]
    for i in range(max_sym + 1):
        cum.append(cum[-1] + freq.get(i, 0))
    total = cum[-1]
    if total == 0:
        return {}, 0
    table = {}
    for sym, f in freq.items():
        table[sym] = (f, cum[sym])
    return table, total


def rans_encode(tokens):
    if not tokens:
        return b""
    freq, total = _build_freq(tokens)
    if total == 0:
        return struct.pack("<I", 0)

    state = RANS_L
    out = []

    for sym in reversed(tokens):
        f, start = freq[sym]
        slot = state % f
        state = (state // f) * total + start + slot
        while state >= RANS_L:
            out.append(state & 0xFFFF)
            state >>= 16

    out.append(state & 0xFFFF)
    state >>= 16
    if state:
        out.append(state)

    fbytes = struct.pack(f"<{len(out)}H", *out)

    meta = struct.pack("<III", len(tokens), total, len(freq))
    for sym, (f, start) in sorted(freq.items()):
        meta += struct.pack("<II", sym, f)
    return struct.pack("<I", len(meta)) + meta + fbytes


def rans_decode(payload):
    if not payload:
        return []

    hlen = struct.unpack_from("<I", payload, 0)[0]
    meta = payload[4:4 + hlen]
    num_syms, total, nfreq = struct.unpack_from("<III", meta, 0)
    if total == 0:
        return []

    freq = {}
    cum = {}
    offset = 12
    c = 0
    for _ in range(nfreq):
        sym, f = struct.unpack_from("<II", meta, offset)
        freq[sym] = f
        cum[sym] = c
        c += f
        offset += 8

    data = payload[4 + hlen:]
    words = list(struct.unpack(f"<{len(data)//2}H", data))

    state = 0
    for _ in range(min(2, len(words))):
        state = (state << 16) | words.pop()
    if not words and state == 0:
        state = 0

    result = []
    cum_list = sorted(cum.items())
    syms = [s for s, _ in cum_list]
    cum_vals = [c for _, c in cum_list]

    for _ in range(num_syms):
        slot = state % total
        state //= total

        lo, hi = 0, len(cum_vals)
        while lo < hi:
            mid = (lo + hi) // 2
            if cum_vals[mid] <= slot:
                lo = mid + 1
            else:
                hi = mid
        idx = max(0, lo - 1)
        sym = syms[idx]
        f = freq[sym]
        start = cum[sym]
        state = state * f + (slot - start)

        result.append(sym)

        while state < RANS_L and words:
            state = (state << 16) | words.pop()

    return result