#include <stdint.h>
#include <string.h>

#include "constants.h"
#include "parameters.h"
#include "tweakey.h"


#define LANE_BITS  64
#define LANE_BYTES (LANE_BITS/8)
#define LANES_NB (TWEAKEY_BYTES/LANE_BYTES)


void tweakey_state_init(
    uint8_t TK[TWEAKEY_BYTES],
    const uint8_t key[KEY_BYTES],
    const uint8_t tweak[TWEAK_BYTES]
)
{
    memcpy(TK,             tweak, TWEAK_BYTES);
    memcpy(TK+TWEAK_BYTES, key,   KEY_BYTES);
}


void tweakey_state_extract(
    const uint8_t TK[TWEAKEY_BYTES],
    uint8_t round_constant,
    uint8_t round_tweakey[ROUND_TWEAKEY_BYTES]
)
{
    memset(round_tweakey, 0, ROUND_TWEAKEY_BYTES);

    for (const uint8_t *lane=TK; lane<TK+TWEAKEY_BYTES; lane+=LANE_BYTES)
    {
        for (size_t j=0; j<LANE_BYTES; j++)
        {
            round_tweakey[j] ^= lane[j];
        }
    }

    round_tweakey[0] ^= round_constant;
}


static void _permute_state(uint8_t TK[TWEAKEY_BYTES])
{
    uint8_t TK_old[TWEAKEY_BYTES];
    memcpy(TK_old, TK, TWEAKEY_BYTES);

    for (size_t j=0; j<TWEAKEY_BYTES; j+=LANE_BYTES)
    {
        for (size_t k=0; k<LANE_BYTES; k++)
        {
            TK[j+h[k]] = TK_old[j+k];
        }
    }
}

static void _multiply_state(uint8_t TK[TWEAKEY_BYTES])
{
    /* Each byte in lane 0 is multiplied by alpha_0 = 1, i.e. it
     * remains unchanged.
     *
     * Each byte b in lanes j = { 1, ..., p-1 } is multiplied by
     * alpha_j; the result of b*alpha_j is stored in P_j[b].
     *
     * In this implementation, P_j sequences are stored in array P;
     * P_j = P[j-1].
     */

    for (size_t j=1; j<LANES_NB; j++)
    {
        const uint8_t *P_j = P[j-1];

        for (size_t k=0; k<LANE_BYTES; k++)
        {
            size_t offset = j*LANE_BYTES + k;
            TK[offset] = P_j[TK[offset]];
        }
    }
}

void tweakey_state_update(uint8_t TK[TWEAKEY_BYTES])
{
    _permute_state(TK);
    _multiply_state(TK);
}