/*************************************************************************
* Copyright (C) 2020 Intel Corporation
*
* Licensed under the Apache License,  Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* 	http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law  or agreed  to  in  writing,  software
* distributed under  the License  is  distributed  on  an  "AS IS"  BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the  specific  language  governing  permissions  and
* limitations under the License.
*************************************************************************/

#include "owndefs.h"
#include "owncp.h"
#include "pcpaesm.h"
#include "pcptool.h"
#include "aes_cfb_vaes_mb.h"

#if (_IPP32E >= _IPP32E_K1)

#define AES_ENCRYPT_VAES_MB16(b0, b1, b2, b3, pRkey, num_rounds)    \
    {                                                               \
        __m512i(*tkeys)[4] = &(pRkey)[(num_rounds)-9];              \
        (b0)               = _mm512_xor_si512((b0), (pRkey)[0][0]); \
        (b1)               = _mm512_xor_si512((b1), (pRkey)[0][1]); \
        (b2)               = _mm512_xor_si512((b2), (pRkey)[0][2]); \
        (b3)               = _mm512_xor_si512((b3), (pRkey)[0][3]); \
        switch (num_rounds) {                                       \
        case 14:                                                    \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[-4][0]);        \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[-4][1]);        \
            (b2) = _mm512_aesenc_epi128((b2), tkeys[-4][2]);        \
            (b3) = _mm512_aesenc_epi128((b3), tkeys[-4][3]);        \
                                                                    \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[-3][0]);        \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[-3][1]);        \
            (b2) = _mm512_aesenc_epi128((b2), tkeys[-3][2]);        \
            (b3) = _mm512_aesenc_epi128((b3), tkeys[-3][3]);        \
            IPPCP_FALLTHROUGH;                                      \
        case 12:                                                    \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[-2][0]);        \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[-2][1]);        \
            (b2) = _mm512_aesenc_epi128((b2), tkeys[-2][2]);        \
            (b3) = _mm512_aesenc_epi128((b3), tkeys[-2][3]);        \
                                                                    \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[-1][0]);        \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[-1][1]);        \
            (b2) = _mm512_aesenc_epi128((b2), tkeys[-1][2]);        \
            (b3) = _mm512_aesenc_epi128((b3), tkeys[-1][3]);        \
            IPPCP_FALLTHROUGH;                                      \
        default:                                                    \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[0][0]);         \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[0][1]);         \
            (b2) = _mm512_aesenc_epi128((b2), tkeys[0][2]);         \
            (b3) = _mm512_aesenc_epi128((b3), tkeys[0][3]);         \
                                                                    \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[1][0]);         \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[1][1]);         \
            (b2) = _mm512_aesenc_epi128((b2), tkeys[1][2]);         \
            (b3) = _mm512_aesenc_epi128((b3), tkeys[1][3]);         \
                                                                    \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[2][0]);         \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[2][1]);         \
            (b2) = _mm512_aesenc_epi128((b2), tkeys[2][2]);         \
            (b3) = _mm512_aesenc_epi128((b3), tkeys[2][3]);         \
                                                                    \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[3][0]);         \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[3][1]);         \
            (b2) = _mm512_aesenc_epi128((b2), tkeys[3][2]);         \
            (b3) = _mm512_aesenc_epi128((b3), tkeys[3][3]);         \
                                                                    \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[4][0]);         \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[4][1]);         \
            (b2) = _mm512_aesenc_epi128((b2), tkeys[4][2]);         \
            (b3) = _mm512_aesenc_epi128((b3), tkeys[4][3]);         \
                                                                    \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[5][0]);         \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[5][1]);         \
            (b2) = _mm512_aesenc_epi128((b2), tkeys[5][2]);         \
            (b3) = _mm512_aesenc_epi128((b3), tkeys[5][3]);         \
                                                                    \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[6][0]);         \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[6][1]);         \
            (b2) = _mm512_aesenc_epi128((b2), tkeys[6][2]);         \
            (b3) = _mm512_aesenc_epi128((b3), tkeys[6][3]);         \
                                                                    \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[7][0]);         \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[7][1]);         \
            (b2) = _mm512_aesenc_epi128((b2), tkeys[7][2]);         \
            (b3) = _mm512_aesenc_epi128((b3), tkeys[7][3]);         \
                                                                    \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[8][0]);         \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[8][1]);         \
            (b2) = _mm512_aesenc_epi128((b2), tkeys[8][2]);         \
            (b3) = _mm512_aesenc_epi128((b3), tkeys[8][3]);         \
                                                                    \
            (b0) = _mm512_aesenclast_epi128((b0), tkeys[9][0]);     \
            (b1) = _mm512_aesenclast_epi128((b1), tkeys[9][1]);     \
            (b2) = _mm512_aesenclast_epi128((b2), tkeys[9][2]);     \
            (b3) = _mm512_aesenclast_epi128((b3), tkeys[9][3]);     \
        }                                                           \
    }

// Disable optimization for MSVC
#if defined(_MSC_VER) && !defined(__INTEL_COMPILER)
#pragma optimize("", off)
#endif

/* clang-format off */
IPP_OWN_DEFN(void, aes_cfb16_enc_vaes_mb16, (const Ipp8u* const source_pa[16],
                                             Ipp8u* const dst_pa[16],
                                             const int arr_len[16],
                                             const int num_rounds,
                                             const Ipp32u* enc_keys[16],
                                             const Ipp8u* iv_pa[16]))
/* clang-format oт */
{
    int i, j, k;
    int maxLen = 0;
    int loc_len64[16];
    Ipp8u* loc_src[16];
    Ipp8u* loc_dst[16];
    __m512i iv512[4];

    __mmask8 mbMask128[16] = { 0x03, 0x0C, 0x30, 0xC0, 0x03, 0x0C, 0x30, 0xC0,
                               0x03, 0x0C, 0x30, 0xC0, 0x03, 0x0C, 0x30, 0xC0 };
    __mmask8 mbMask[16]    = { 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF,
                               0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF };

    // - Local copy of length, source and target pointers, maxLen calculation
    for (i = 0; i < 16; i++) {
        // The case of the empty input buffer
        if (arr_len[i] == 0) {
            mbMask128[i] = 0;
            mbMask[i]    = 0;
            loc_len64[i] = 0;
            continue;
        }

        loc_src[i]   = (Ipp8u*)source_pa[i];
        loc_dst[i]   = (Ipp8u*)dst_pa[i];
        int len64    = arr_len[i] / (Ipp32s)sizeof(Ipp64u); // length in 64-bit chunks
        loc_len64[i] = len64;

        if (len64 < 8)
            mbMask[i] = (__mmask8)(((1 << len64) - 1) & 0xFF);
        if (len64 > maxLen)
            maxLen = len64;
    }

    // Load the necessary number of 128-bit IV
    j = 0;
    for (i = 0; i < 16; i += 4) {
        iv512[j] = _mm512_setzero_si512();

        iv512[j] = _mm512_mask_expandloadu_epi64(iv512[j], mbMask128[i], iv_pa[i]);
        iv512[j] = _mm512_mask_expandloadu_epi64(iv512[j], mbMask128[i + 1], iv_pa[i + 1]);
        iv512[j] = _mm512_mask_expandloadu_epi64(iv512[j], mbMask128[i + 2], iv_pa[i + 2]);
        iv512[j] = _mm512_mask_expandloadu_epi64(iv512[j], mbMask128[i + 3], iv_pa[i + 3]);
        j += 1;
    }

    // Temporary block to left IV unchanged to use it on the next round
    __m512i chip0 = iv512[0];
    __m512i chip1 = iv512[1];
    __m512i chip2 = iv512[2];
    __m512i chip3 = iv512[3];

    // Prepare array with key schedule
    __m512i keySchedule[15][4];
    __m512i tmpKeyMb = _mm512_setzero_si512();
    for (i = 0; i <= num_rounds; i++) {
        k = 0;
        for (j = 0; j < 16; j += 4) {
            tmpKeyMb = _mm512_setzero_si512();
            tmpKeyMb = _mm512_mask_expandloadu_epi64(
                tmpKeyMb,
                mbMask128[j],
                (const void*)(enc_keys[j] + (Ipp32u)i * sizeof(Ipp32u)));
            tmpKeyMb = _mm512_mask_expandloadu_epi64(
                tmpKeyMb,
                mbMask128[j + 1],
                (const void*)(enc_keys[j + 1] + (Ipp32u)i * sizeof(Ipp32u)));
            tmpKeyMb = _mm512_mask_expandloadu_epi64(
                tmpKeyMb,
                mbMask128[j + 2],
                (const void*)(enc_keys[j + 2] + (Ipp32u)i * sizeof(Ipp32u)));
            tmpKeyMb = _mm512_mask_expandloadu_epi64(
                tmpKeyMb,
                mbMask128[j + 3],
                (const void*)(enc_keys[j + 3] + (Ipp32u)i * sizeof(Ipp32u)));

            keySchedule[i][k] = _mm512_loadu_si512(&tmpKeyMb);
            k += 1;
        }
    }

    for (; maxLen >= 0; maxLen -= 8) {
        __m512i b0 = _mm512_maskz_loadu_epi64(mbMask[0], loc_src[0]);
        loc_src[0] += CFB16_BLOCK_SIZE * 4;
        __m512i b1 = _mm512_maskz_loadu_epi64(mbMask[1], loc_src[1]);
        loc_src[1] += CFB16_BLOCK_SIZE * 4;
        __m512i b2 = _mm512_maskz_loadu_epi64(mbMask[2], loc_src[2]);
        loc_src[2] += CFB16_BLOCK_SIZE * 4;
        __m512i b3 = _mm512_maskz_loadu_epi64(mbMask[3], loc_src[3]);
        loc_src[3] += CFB16_BLOCK_SIZE * 4;
        __m512i b4 = _mm512_maskz_loadu_epi64(mbMask[4], loc_src[4]);
        loc_src[4] += CFB16_BLOCK_SIZE * 4;
        __m512i b5 = _mm512_maskz_loadu_epi64(mbMask[5], loc_src[5]);
        loc_src[5] += CFB16_BLOCK_SIZE * 4;
        __m512i b6 = _mm512_maskz_loadu_epi64(mbMask[6], loc_src[6]);
        loc_src[6] += CFB16_BLOCK_SIZE * 4;
        __m512i b7 = _mm512_maskz_loadu_epi64(mbMask[7], loc_src[7]);
        loc_src[7] += CFB16_BLOCK_SIZE * 4;
        __m512i b8 = _mm512_maskz_loadu_epi64(mbMask[8], loc_src[8]);
        loc_src[8] += CFB16_BLOCK_SIZE * 4;
        __m512i b9 = _mm512_maskz_loadu_epi64(mbMask[9], loc_src[9]);
        loc_src[9] += CFB16_BLOCK_SIZE * 4;
        __m512i b10 = _mm512_maskz_loadu_epi64(mbMask[10], loc_src[10]);
        loc_src[10] += CFB16_BLOCK_SIZE * 4;
        __m512i b11 = _mm512_maskz_loadu_epi64(mbMask[11], loc_src[11]);
        loc_src[11] += CFB16_BLOCK_SIZE * 4;
        __m512i b12 = _mm512_maskz_loadu_epi64(mbMask[12], loc_src[12]);
        loc_src[12] += CFB16_BLOCK_SIZE * 4;
        __m512i b13 = _mm512_maskz_loadu_epi64(mbMask[13], loc_src[13]);
        loc_src[13] += CFB16_BLOCK_SIZE * 4;
        __m512i b14 = _mm512_maskz_loadu_epi64(mbMask[14], loc_src[14]);
        loc_src[14] += CFB16_BLOCK_SIZE * 4;
        __m512i b15 = _mm512_maskz_loadu_epi64(mbMask[15], loc_src[15]);
        loc_src[15] += CFB16_BLOCK_SIZE * 4;

        TRANSPOSE_4x4_I128(b0, b1, b2, b3);     // {0,0,0,0}, {1,1,1,1}, {2,2,2,2}, {3,3,3,3}
        TRANSPOSE_4x4_I128(b4, b5, b6, b7);     // {0,0,0,0}, {1,1,1,1}, {2,2,2,2}, {3,3,3,3}
        TRANSPOSE_4x4_I128(b8, b9, b10, b11);   // {0,0,0,0}, {1,1,1,1}, {2,2,2,2}, {3,3,3,3}
        TRANSPOSE_4x4_I128(b12, b13, b14, b15); // {0,0,0,0}, {1,1,1,1}, {2,2,2,2}, {3,3,3,3}

        AES_ENCRYPT_VAES_MB16(chip0, chip1, chip2, chip3, keySchedule, num_rounds);
        chip0 = b0 = _mm512_xor_si512(b0, chip0);
        chip1 = b4 = _mm512_xor_si512(b4, chip1);
        chip2 = b8 = _mm512_xor_si512(b8, chip2);
        chip3 = b12 = _mm512_xor_si512(b12, chip3);

        AES_ENCRYPT_VAES_MB16(chip0, chip1, chip2, chip3, keySchedule, num_rounds);
        chip0 = b1 = _mm512_xor_si512(b1, chip0);
        chip1 = b5 = _mm512_xor_si512(b5, chip1);
        chip2 = b9 = _mm512_xor_si512(b9, chip2);
        chip3 = b13 = _mm512_xor_si512(b13, chip3);

        AES_ENCRYPT_VAES_MB16(chip0, chip1, chip2, chip3, keySchedule, num_rounds);
        chip0 = b2 = _mm512_xor_si512(b2, chip0);
        chip1 = b6 = _mm512_xor_si512(b6, chip1);
        chip2 = b10 = _mm512_xor_si512(b10, chip2);
        chip3 = b14 = _mm512_xor_si512(b14, chip3);

        AES_ENCRYPT_VAES_MB16(chip0, chip1, chip2, chip3, keySchedule, num_rounds);
        chip0 = b3 = _mm512_xor_si512(b3, chip0);
        chip1 = b7 = _mm512_xor_si512(b7, chip1);
        chip2 = b11 = _mm512_xor_si512(b11, chip2);
        chip3 = b15 = _mm512_xor_si512(b15, chip3);

        TRANSPOSE_4x4_I128(b0, b1, b2, b3);
        TRANSPOSE_4x4_I128(b4, b5, b6, b7);
        TRANSPOSE_4x4_I128(b8, b9, b10, b11);
        TRANSPOSE_4x4_I128(b12, b13, b14, b15);

        _mm512_mask_storeu_epi64(loc_dst[0], mbMask[0], b0);
        loc_dst[0] += CFB16_BLOCK_SIZE * 4;
        loc_len64[0] -= 2 * 4;
        UPDATE_MASK(loc_len64[0], mbMask[0]);
        _mm512_mask_storeu_epi64(loc_dst[1], mbMask[1], b1);
        loc_dst[1] += CFB16_BLOCK_SIZE * 4;
        loc_len64[1] -= 2 * 4;
        UPDATE_MASK(loc_len64[1], mbMask[1]);
        _mm512_mask_storeu_epi64(loc_dst[2], mbMask[2], b2);
        loc_dst[2] += CFB16_BLOCK_SIZE * 4;
        loc_len64[2] -= 2 * 4;
        UPDATE_MASK(loc_len64[2], mbMask[2]);
        _mm512_mask_storeu_epi64(loc_dst[3], mbMask[3], b3);
        loc_dst[3] += CFB16_BLOCK_SIZE * 4;
        loc_len64[3] -= 2 * 4;
        UPDATE_MASK(loc_len64[3], mbMask[3]);
        _mm512_mask_storeu_epi64(loc_dst[4], mbMask[4], b4);
        loc_dst[4] += CFB16_BLOCK_SIZE * 4;
        loc_len64[4] -= 2 * 4;
        UPDATE_MASK(loc_len64[4], mbMask[4]);
        _mm512_mask_storeu_epi64(loc_dst[5], mbMask[5], b5);
        loc_dst[5] += CFB16_BLOCK_SIZE * 4;
        loc_len64[5] -= 2 * 4;
        UPDATE_MASK(loc_len64[5], mbMask[5]);
        _mm512_mask_storeu_epi64(loc_dst[6], mbMask[6], b6);
        loc_dst[6] += CFB16_BLOCK_SIZE * 4;
        loc_len64[6] -= 2 * 4;
        UPDATE_MASK(loc_len64[6], mbMask[6]);
        _mm512_mask_storeu_epi64(loc_dst[7], mbMask[7], b7);
        loc_dst[7] += CFB16_BLOCK_SIZE * 4;
        loc_len64[7] -= 2 * 4;
        UPDATE_MASK(loc_len64[7], mbMask[7]);
        _mm512_mask_storeu_epi64(loc_dst[8], mbMask[8], b8);
        loc_dst[8] += CFB16_BLOCK_SIZE * 4;
        loc_len64[8] -= 2 * 4;
        UPDATE_MASK(loc_len64[8], mbMask[8]);
        _mm512_mask_storeu_epi64(loc_dst[9], mbMask[9], b9);
        loc_dst[9] += CFB16_BLOCK_SIZE * 4;
        loc_len64[9] -= 2 * 4;
        UPDATE_MASK(loc_len64[9], mbMask[9]);
        _mm512_mask_storeu_epi64(loc_dst[10], mbMask[10], b10);
        loc_dst[10] += CFB16_BLOCK_SIZE * 4;
        loc_len64[10] -= 2 * 4;
        UPDATE_MASK(loc_len64[10], mbMask[10]);
        _mm512_mask_storeu_epi64(loc_dst[11], mbMask[11], b11);
        loc_dst[11] += CFB16_BLOCK_SIZE * 4;
        loc_len64[11] -= 2 * 4;
        UPDATE_MASK(loc_len64[11], mbMask[11]);
        _mm512_mask_storeu_epi64(loc_dst[12], mbMask[12], b12);
        loc_dst[12] += CFB16_BLOCK_SIZE * 4;
        loc_len64[12] -= 2 * 4;
        UPDATE_MASK(loc_len64[12], mbMask[12]);
        _mm512_mask_storeu_epi64(loc_dst[13], mbMask[13], b13);
        loc_dst[13] += CFB16_BLOCK_SIZE * 4;
        loc_len64[13] -= 2 * 4;
        UPDATE_MASK(loc_len64[13], mbMask[13]);
        _mm512_mask_storeu_epi64(loc_dst[14], mbMask[14], b14);
        loc_dst[14] += CFB16_BLOCK_SIZE * 4;
        loc_len64[14] -= 2 * 4;
        UPDATE_MASK(loc_len64[14], mbMask[14]);
        _mm512_mask_storeu_epi64(loc_dst[15], mbMask[15], b15);
        loc_dst[15] += CFB16_BLOCK_SIZE * 4;
        loc_len64[15] -= 2 * 4;
        UPDATE_MASK(loc_len64[15], mbMask[15]);
    }
}

#if defined(_MSC_VER) && !defined(__INTEL_COMPILER)
#pragma optimize("", on)
#endif

#endif
