Encode Integers with LEB128 Varints

from __future__ import annotations

LIMIT = 1 << 64


def encode_varint(value: int) -> bytes:
    if not 0 <= value < LIMIT:
        raise ValueError("value must fit in 64 bits")
    output = bytearray()
    while True:
        byte = value & 0x7F
        value >>= 7
        if value:
            output.append(byte | 0x80)
        else:
            output.append(byte)
            return bytes(output)


def decode_varint(data: bytes, offset: int = 0) -> tuple[int, int]:
    value = 0
    shift = 0
    while True:
        if offset >= len(data):
            raise ValueError("truncated varint")
        byte = data[offset]
        offset += 1
        value |= (byte & 0x7F) << shift
        if not byte & 0x80:
            if value >= LIMIT:
                raise ValueError("varint exceeds 64 bits")
            return value, offset
        shift += 7
        if shift > 63:
            raise ValueError("varint too long")


def zigzag_encode(value: int) -> int:
    return (value << 1) ^ (value >> 63)


def zigzag_decode(value: int) -> int:
    return (value >> 1) ^ -(value & 1)


def main() -> None:
    for number in [0, 1, 127, 128, 300, 2**32]:
        encoded = encode_varint(number)
        decoded, next_offset = decode_varint(encoded)
        print(number, encoded.hex(), decoded)
        assert decoded == number
        assert next_offset == len(encoded)

    combined = encode_varint(0) + encode_varint(300)
    first, offset = decode_varint(combined)
    second, offset = decode_varint(combined, offset)
    assert (first, second, offset) == (0, 300, len(combined))

    for value in [-500, -1, 0, 1, 500]:
        assert zigzag_decode(zigzag_encode(value)) == value


if __name__ == "__main__":
    main()