#include <stdio.h>
#include <stdint.h>
#include <immintrin.h>

/*
Author: Ondrej Hladuvka, hladuond@fit.cvut.cz
Template: Jiri Bucek 2017
AES specification:
http://csrc.nist.gov/publications/fips/fips197/fips-197.pdf
*/

namespace aes128_6 {
static inline __m128i expandRoundKey(__m128i key, __m128i keygened){
    keygened = _mm_shuffle_epi32(keygened, _MM_SHUFFLE(3,3,3,3));
    key = _mm_xor_si128(key, _mm_slli_si128(key, 4));
    key = _mm_xor_si128(key, _mm_slli_si128(key, 4));
    key = _mm_xor_si128(key, _mm_slli_si128(key, 4));
    return _mm_xor_si128(key, keygened);
}

void expandKey(uint8_t *key, __m128i *ex){
    ex[0] = _mm_loadu_si128((const __m128i*) key);
    ex[1]  = expandRoundKey(ex[0], _mm_aeskeygenassist_si128(ex[0], 0x01));
    ex[2]  = expandRoundKey(ex[1], _mm_aeskeygenassist_si128(ex[1], 0x02));
    ex[3]  = expandRoundKey(ex[2], _mm_aeskeygenassist_si128(ex[2], 0x04));
    ex[4]  = expandRoundKey(ex[3], _mm_aeskeygenassist_si128(ex[3], 0x08));
    ex[5]  = expandRoundKey(ex[4], _mm_aeskeygenassist_si128(ex[4], 0x10));
    ex[6]  = expandRoundKey(ex[5], _mm_aeskeygenassist_si128(ex[5], 0x20));
    ex[7]  = expandRoundKey(ex[6], _mm_aeskeygenassist_si128(ex[6], 0x40));
    ex[8]  = expandRoundKey(ex[7], _mm_aeskeygenassist_si128(ex[7], 0x80));
    ex[9]  = expandRoundKey(ex[8], _mm_aeskeygenassist_si128(ex[8], 0x1B));
    ex[10] = expandRoundKey(ex[9], _mm_aeskeygenassist_si128(ex[9], 0x36));
}

void aes128_6(uint8_t *in, uint8_t *out, uint32_t* key) {
    __m128i *k = (__m128i *)key;
    __m128i state = _mm_loadu_si128((__m128i *) in);

    state = _mm_xor_si128       (state, k[ 0]);
    state = _mm_aesenc_si128    (state, k[ 1]);
    state = _mm_aesenc_si128    (state, k[ 2]);
    state = _mm_aesenc_si128    (state, k[ 3]);
    state = _mm_aesenc_si128    (state, k[ 4]);
    state = _mm_aesenc_si128    (state, k[ 5]);
    state = _mm_aesenc_si128    (state, k[ 6]);
    state = _mm_aesenc_si128    (state, k[ 7]);
    state = _mm_aesenc_si128    (state, k[ 8]);
    state = _mm_aesenc_si128    (state, k[ 9]);
    state = _mm_aesenclast_si128(state, k[10]);

    _mm_storeu_si128((__m128i *) out, state);
}

}