limbs.py

"""Limb codec shared by the scaled CoT encodings (divcot2, karatsuba).

A limb is one digit in base B (10 or 100), encoded as a SINGLE byte
0x80+value (128..227) so sequence length scales with limb count, not decimal
digits. Base-100 halves W and roughly quarters quadratic CoT lengths.
Strings are latin1-safe: encode("latin1") round-trips; never use ascii.
"""
from __future__ import annotations

LIMB_OFFSET = 0x80  # limb v -> chr(0x80+v); supports base <= 128


def limb_char(v: int) -> str:
    return chr(LIMB_OFFSET + v)


def char_limb(c: str) -> int:
    """Inverse of limb_char; out-of-range chars decode to 0 (eval leniency)."""
    v = ord(c) - LIMB_OFFSET
    return v if v >= 0 else 0


def to_limbs(n: int, width: int, base: int, msb_first: bool = False) -> list[int]:
    out = []
    for _ in range(width):
        out.append(n % base)
        n //= base
    assert n == 0, "width too small for value"
    return out[::-1] if msb_first else out


def from_limbs(limbs, base: int, msb_first: bool = False) -> int:
    seq = limbs[::-1] if msb_first else limbs
    val = 0
    for k, v in enumerate(seq):
        val += v * base ** k
    return val


def limb_str(n: int, width: int, base: int, msb_first: bool = False) -> str:
    return "".join(limb_char(v) for v in to_limbs(n, width, base, msb_first))


def parse_limbs(s: str, base: int, msb_first: bool = False) -> int:
    vals = [min(char_limb(c), base - 1) for c in s]
    return from_limbs(vals, base, msb_first)


def width_for(p_max: int, base: int) -> int:
    w = 1
    while base ** w <= p_max - 1:
        w += 1
    return w


if __name__ == "__main__":
    import random
    rng = random.Random(0)
    for base in (10, 100):
        for _ in range(2000):
            w = rng.randint(1, 12)
            n = rng.randrange(base ** w)
            for msb in (False, True):
                s = limb_str(n, w, base, msb)
                assert len(s) == w and parse_limbs(s, base, msb) == n
    print("limbs.py: all round-trip tests pass")