rho.py

ROT = [
    [0, 36, 3, 41, 18],
    [1, 44, 10, 45, 2],
    [62, 6, 43, 15, 61],
    [28, 55, 25, 21, 56],
    [27, 20, 39, 8, 14],
]

# OFFSETs = [[64 - i for i in row] for row in ROT]
OFFSETs = [
    [64, 28, 61, 23, 46],
    [63, 20, 54, 19, 62],
    [2, 58, 21, 49, 3],
    [36, 9, 39, 43, 8],
    [37, 44, 25, 56, 50]
]


def keccak_u64_first_converter(n: int) -> int:
    """
    n is the sum of 12 different bits.
    The theta step has 12 xor operations
    If the sum is odd, then the 12 xor operations result 1
    If the sum is even, then the 12 xor operations result 0
    """
    assert n < 13
    return n & 1


def get_step_size(cur_offset: int, rot: int) -> int:
    """
    Is the `check_offset_helper` in rust with
    base_num_of_chunks = 4
    """
    if 60 - cur_offset < rot < 64 - cur_offset:
        return 64 - rot - cur_offset
    # near the end of the lane
    if cur_offset < 64 < cur_offset + 4:
        return 64 - cur_offset
    return 4


def query_table_accumulate(table, x, f_of_x, g_of_x, acc, coef, is_final=False):
    assert table.lookup(x, f_of_x, g_of_x)
    acc_next = acc - coef * x
    if is_final:
        # also means acc_next == 0
        assert acc == coef * x
    return f_of_x, g_of_x, acc_next


def cnst_rotate_and_convert(_input: int, rot: int) -> int:
    init_base = 9**rot
    special_chunk = 0
    _input = _input + 2**64
    acc = 0

    for i in range(65):
        remainder = _input % 13
        if i in (0, 64):
            special_chunk += remainder
        else:
            acc += keccak_u64_first_converter(remainder) * base
        _input /= 13
        base *= 9
        if i == 64 - rot:
            base = 1
    acc += keccak_u64_first_converter(special_chunk) * init_base
    return acc


def rho(state, first_iteration: bool):

    offset_map = {
        1: [],
        2: [],
        3: []
    }

    for i in range(5):
        for j in range(5):
            output_slices = []
            output_coefs = []
            cur_offset = 0
            cur_input_coef = 1
            cur_output_coef = 9**ROT[i][j]
            acc = state[i][j]

            state[i][j] = cnst_rotate_and_convert(state[i][j])

            last_chunk_low_value = 0
            if first_iteration:
                last_chunk_low_value = raw_value % 13
                raw_value /= 13

            last_coef = cur_output_coef
            cur_input_coef *= 13
            if ROT[i][j] == 63:
                cur_output_coef = 1
            else:
                cur_output_coef *= 9
            cur_offset += 1
            while cur_offset < 64:
                step = get_step_size(cur_offset, ROT[i][j])
                if first_iteration:
                    input_slice = raw_value % 13**step
                    raw_value /= 13**step

                g_chunk, output_slice, acc = query_table_accumulate(
                    of_first_to_second_base_converter_table,
                    input_slice,
                    acc,
                    cur_input_coef,
                    is_final=False
                )
                output_slices.append(output_slice)
                cur_offset += step
                cur_input_coef *= 13 ** step
                # modify output_chunk coefficient
                if cur_offset == 64 - ROT[i][j]:
                    cur_output_coef = 1
                else:
                    cur_output_coef *= 9 ** step

                if step < 4:
                    offset_map[step].append(g_chunk)
            last_chunk_high_value = 0
            if first_iteration:
                last_chunk_high_value = raw_value % 13
                raw_value /= 13

            last_chunk_value = last_chunk_high_value * (13**64) + last_chunk_low_value

            (_, output_slice, _) = query_table_accumulate(
                of_first_to_second_base_converter_table,
                last_chunk_value,
                acc,
                1,
                is_final=True
            )
            output_coefs.append(last_coef)
            output_slices.append(output_slice)
    offset_transformed = [0, 0, 1, 13, 170]
    for i in (1,2, 3):
        sum_of_inputs = sum(offset_map[i])
        total = offset_transformed[i] * len(offset_map[i])
        # Perform a range check here
        assert sum_of_inputs <= total