/*************************************************************************
* 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, pRkey, num_rounds)          \
    {                                                             \
        __m512i(*tkeys)[2] = &(pRkey)[(num_rounds)-9];            \
        (b0)               = _mm512_xor_si512(b0, (pRkey)[0][0]); \
        (b1)               = _mm512_xor_si512(b1, (pRkey)[0][1]); \
        switch (num_rounds) {                                     \
        case 14:                                                  \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[-4][0]);      \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[-4][1]);      \
                                                                  \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[-3][0]);      \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[-3][1]);      \
            IPPCP_FALLTHROUGH;                                    \
        case 12:                                                  \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[-2][0]);      \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[-2][1]);      \
                                                                  \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[-1][0]);      \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[-1][1]);      \
            IPPCP_FALLTHROUGH;                                    \
        default:                                                  \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[0][0]);       \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[0][1]);       \
                                                                  \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[1][0]);       \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[1][1]);       \
                                                                  \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[2][0]);       \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[2][1]);       \
                                                                  \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[3][0]);       \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[3][1]);       \
                                                                  \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[4][0]);       \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[4][1]);       \
                                                                  \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[5][0]);       \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[5][1]);       \
                                                                  \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[6][0]);       \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[6][1]);       \
                                                                  \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[7][0]);       \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[7][1]);       \
                                                                  \
            (b0) = _mm512_aesenc_epi128((b0), tkeys[8][0]);       \
            (b1) = _mm512_aesenc_epi128((b1), tkeys[8][1]);       \
                                                                  \
            (b0) = _mm512_aesenclast_epi128((b0), tkeys[9][0]);   \
            (b1) = _mm512_aesenclast_epi128((b1), tkeys[9][1]);   \
        }                                                         \
    }

// 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_mb8, (const Ipp8u* const source_pa[8],
                                            Ipp8u* const dst_pa[8],
                                            const int arr_len[8],
                                            const int num_rounds,
                                            const Ipp32u* enc_keys[8],
                                            const Ipp8u* iv_pa[8]))
/* clang-format on */
{
    int i;
    int maxLen = 0;
    int loc_len64[8];
    Ipp8u* loc_src[8];
    Ipp8u* loc_dst[8];
    __m512i iv512[2];

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

    for (i = 0; i < 8; i++) {
        // The case of the empty input buffer
        if (arr_len[i] == 0) {
            mbMask128[i] = 0;
            mbMask[i]    = 0;
            loc_len64[i] = 0;
            loc_src[i]   = NULL;
            loc_dst[i]   = NULL;
            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
    iv512[0] = _mm512_setzero_si512();
    iv512[0] = _mm512_mask_expandloadu_epi64(iv512[0], mbMask128[0], iv_pa[0]);
    iv512[0] = _mm512_mask_expandloadu_epi64(iv512[0], mbMask128[1], iv_pa[1]);
    iv512[0] = _mm512_mask_expandloadu_epi64(iv512[0], mbMask128[2], iv_pa[2]);
    iv512[0] = _mm512_mask_expandloadu_epi64(iv512[0], mbMask128[3], iv_pa[3]);

    iv512[1] = _mm512_setzero_si512();
    iv512[1] = _mm512_mask_expandloadu_epi64(iv512[1], mbMask128[4], iv_pa[4]);
    iv512[1] = _mm512_mask_expandloadu_epi64(iv512[1], mbMask128[5], iv_pa[5]);
    iv512[1] = _mm512_mask_expandloadu_epi64(iv512[1], mbMask128[6], iv_pa[6]);
    iv512[1] = _mm512_mask_expandloadu_epi64(iv512[1], mbMask128[7], iv_pa[7]);


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

    // Prepare array with key schedule
    __m512i keySchedule[15][2];
    __m512i tmpKeyMb = _mm512_setzero_si512();

    for (i = 0; i <= num_rounds; i++) {
        tmpKeyMb =
            _mm512_mask_expandloadu_epi64(tmpKeyMb,
                                          mbMask128[0],
                                          (const void*)(enc_keys[0] + (Ipp32u)i * sizeof(Ipp32u)));
        tmpKeyMb =
            _mm512_mask_expandloadu_epi64(tmpKeyMb,
                                          mbMask128[1],
                                          (const void*)(enc_keys[1] + (Ipp32u)i * sizeof(Ipp32u)));
        tmpKeyMb =
            _mm512_mask_expandloadu_epi64(tmpKeyMb,
                                          mbMask128[2],
                                          (const void*)(enc_keys[2] + (Ipp32u)i * sizeof(Ipp32u)));
        tmpKeyMb =
            _mm512_mask_expandloadu_epi64(tmpKeyMb,
                                          mbMask128[3],
                                          (const void*)(enc_keys[3] + (Ipp32u)i * sizeof(Ipp32u)));

        keySchedule[i][0] = _mm512_loadu_si512(&tmpKeyMb);

        tmpKeyMb =
            _mm512_mask_expandloadu_epi64(tmpKeyMb,
                                          mbMask128[4],
                                          (const void*)(enc_keys[4] + (Ipp32u)i * sizeof(Ipp32u)));
        tmpKeyMb =
            _mm512_mask_expandloadu_epi64(tmpKeyMb,
                                          mbMask128[5],
                                          (const void*)(enc_keys[5] + (Ipp32u)i * sizeof(Ipp32u)));
        tmpKeyMb =
            _mm512_mask_expandloadu_epi64(tmpKeyMb,
                                          mbMask128[6],
                                          (const void*)(enc_keys[6] + (Ipp32u)i * sizeof(Ipp32u)));
        tmpKeyMb =
            _mm512_mask_expandloadu_epi64(tmpKeyMb,
                                          mbMask128[7],
                                          (const void*)(enc_keys[7] + (Ipp32u)i * sizeof(Ipp32u)));

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

    for (; maxLen >= 0; maxLen -= 8) {
        // Load plain text from the different buffers
        __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;

        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}

        AES_ENCRYPT_VAES_MB16(chip0, chip1, keySchedule, num_rounds);
        chip0 = b0 = _mm512_xor_si512(b0, chip0);
        chip1 = b4 = _mm512_xor_si512(b4, chip1);

        AES_ENCRYPT_VAES_MB16(chip0, chip1, keySchedule, num_rounds);
        chip0 = b1 = _mm512_xor_si512(b1, chip0);
        chip1 = b5 = _mm512_xor_si512(b5, chip1);

        AES_ENCRYPT_VAES_MB16(chip0, chip1, keySchedule, num_rounds);
        chip0 = b2 = _mm512_xor_si512(b2, chip0);
        chip1 = b6 = _mm512_xor_si512(b6, chip1);

        AES_ENCRYPT_VAES_MB16(chip0, chip1, keySchedule, num_rounds);
        chip0 = b3 = _mm512_xor_si512(b3, chip0);
        chip1 = b7 = _mm512_xor_si512(b7, chip1);

        TRANSPOSE_4x4_I128(b0, b1, b2, b3);
        TRANSPOSE_4x4_I128(b4, b5, b6, b7);

        _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]);
    }
}

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

#endif
