/*************************************************************************
* Copyright (C) 2013 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.
*************************************************************************/

/*
//
//  Purpose:
//     Cryptography Primitive.
//     AES encryption/decryption (CTR mode)
//
//  Contents:
//     EncryptCTR_RIJ128pipe_VAES_NI
//
*/

#include "owncp.h"
#include "pcpaes_encrypt_vaes512.h"
#include "pcpaesm.h"

#if (_IPP32E >= _IPP32E_K1)
#if defined(_MSC_VER) && !defined(__INTEL_COMPILER)
#pragma warning(disable : 4310) // cast truncates constant value in MSVC
#endif

#define M512(mem) (*((__m512i*)((Ipp8u*)(mem))))

/* Mask to convert 64-bit parts of four 128-bit numbers stored in one 512-bit register
 * from Little-Endian to Big-Endian */

/* clang-format off */
static __ALIGN64 Ipp8u swapBytes[] = { 7,  6,  5,  4,  3,  2,  1,  0,  15, 14, 13, 12, 11, 10, 9,  8,
                                       23, 22, 21, 20, 19, 18, 17, 16, 31, 30, 29, 28, 27, 26, 25, 24,
                                       39, 38, 37, 36, 35, 34, 33, 32, 47, 46, 45, 44, 43, 42, 41, 40,
                                       55, 54, 53, 52, 51, 50, 49, 48, 63, 62, 61, 60, 59, 58, 57, 56 };
/* clang-format on */

/* Increment masks for Hi-Lo 64-bit parts of 128-bit numbers in 512-bit register */
static __ALIGN64 Ipp64u startIncLoMask[] = { 0x0, 0x0, 0x0, 0x1, 0x0, 0x2, 0x0, 0x3 };
static __ALIGN64 Ipp64u nextIncLoMask[]  = { 0x0, 0x4, 0x0, 0x4, 0x0, 0x4, 0x0, 0x4 };
static __ALIGN64 Ipp64u incLoByOneMask[] = { 0x0, 0x1, 0x0, 0x1, 0x0, 0x1, 0x0, 0x1 };
static __ALIGN64 Ipp64u incHiByOneMask[] = { 0x1, 0x0, 0x1, 0x0, 0x1, 0x0, 0x1, 0x0 };

IPPCP_INLINE __m512i adcLo_epi64(__m512i a, __m512i b)
{
    a = _mm512_add_epi64(a, b);
    // check overflow in each low 64-bit of 128-bit numbers
    __mmask8 overMsk = _mm512_cmplt_epu64_mask(a, b);
    // get mask of each high 64-bit that need to be increased
    overMsk <<= 1;
    a = _mm512_mask_add_epi64(a, overMsk, a, M512(incHiByOneMask));
    return a;
}

IPPCP_INLINE __m512i applyNonce(__m512i a, __m512i ctrBitMask, __m512i templateCtr)
{
    a = _mm512_shuffle_epi8(a, M512(swapBytes));
    a = _mm512_and_epi64(a, ctrBitMask);
    a = _mm512_or_epi64(a, templateCtr);

    return a;
}

////////////////////////////////////////////////////////////////////////////////
/* clang-format off */
IPP_OWN_DEFN(void, EncryptCTR_RIJ128pipe_VAES_NI, (const Ipp8u* pSrc,
                                                   Ipp8u* pDst,
                                                   int nr,
                                                   const Ipp8u* pRKey,
                                                   int length, /* message length in bytes   */
                                                   Ipp8u* pCtrValue,
                                                   const Ipp8u* pCtrBitMask))
/* clang-format on */
{
    int cipherRounds = nr - 1;

    __m128i* pKeys   = (__m128i*)pRKey;
    __m512i* pSrc512 = (__m512i*)pSrc;
    __m512i* pDst512 = (__m512i*)pDst;

    Ipp64u* pCtr64 = (Ipp64u*)pCtrValue;

    // Initial counter
    __m128i initialCtr128 = _mm_maskz_loadu_epi64(0x03, pCtrValue);
    __m128i ctrBitMask128 = _mm_maskz_loadu_epi64(0x03, pCtrBitMask);
    // Unchanged counter part
    __m128i templateCtr128 =
        _mm_maskz_andnot_epi64(0x03 /* all 128-bits */, ctrBitMask128, initialCtr128);

    __m512i ctrBitMask512  = _mm512_broadcast_i64x2(ctrBitMask128);
    __m512i templateCtr512 = _mm512_broadcast_i64x2(templateCtr128);

    Ipp64u ctr64_h = ENDIANNESS64(pCtr64[0]); // high 64-bit of BE counter converted to LE
    Ipp64u ctr64_l = ENDIANNESS64(pCtr64[1]); // low 64-bit of BE counter converted to LE

    __m512i ctr512 =
        _mm512_set4_epi64((Ipp64s)ctr64_l, (Ipp64s)ctr64_h, (Ipp64s)ctr64_l, (Ipp64s)ctr64_h);

    // int blocks;
    int remainded_length;
    __m512i incMsk = M512(startIncLoMask);
    for (remainded_length = length; remainded_length >= (4 * 4 * MBS_RIJ128);
         remainded_length -= (4 * 4 * MBS_RIJ128)) {
        __m512i counter0 = adcLo_epi64(incMsk, ctr512);
        __m512i counter1 = adcLo_epi64(M512(nextIncLoMask), counter0);
        __m512i counter2 = adcLo_epi64(M512(nextIncLoMask), counter1);
        __m512i counter3 = adcLo_epi64(M512(nextIncLoMask), counter2);

        incMsk = M512(nextIncLoMask);
        ctr512 = counter3;

        // convert back to BE and add nonce
        counter0 = applyNonce(counter0, ctrBitMask512, templateCtr512);
        counter1 = applyNonce(counter1, ctrBitMask512, templateCtr512);
        counter2 = applyNonce(counter2, ctrBitMask512, templateCtr512);
        counter3 = applyNonce(counter3, ctrBitMask512, templateCtr512);

        cpAESEncrypt4_VAES_NI(&counter0, &counter1, &counter2, &counter3, pKeys, cipherRounds);

        __m512i blk0 = _mm512_loadu_si512(pSrc512);
        __m512i blk1 = _mm512_loadu_si512(pSrc512 + 1);
        __m512i blk2 = _mm512_loadu_si512(pSrc512 + 2);
        __m512i blk3 = _mm512_loadu_si512(pSrc512 + 3);

        blk0 = _mm512_xor_si512(blk0, counter0);
        blk1 = _mm512_xor_si512(blk1, counter1);
        blk2 = _mm512_xor_si512(blk2, counter2);
        blk3 = _mm512_xor_si512(blk3, counter3);

        _mm512_storeu_si512(pDst512, blk0);
        _mm512_storeu_si512(pDst512 + 1, blk1);
        _mm512_storeu_si512(pDst512 + 2, blk2);
        _mm512_storeu_si512(pDst512 + 3, blk3);

        pSrc512 += 4;
        pDst512 += 4;
    }

    if ((3 * 4 * MBS_RIJ128) <= remainded_length) {
        __m512i counter0 = adcLo_epi64(incMsk, ctr512);
        __m512i counter1 = adcLo_epi64(M512(nextIncLoMask), counter0);
        __m512i counter2 = adcLo_epi64(M512(nextIncLoMask), counter1);

        incMsk = M512(nextIncLoMask);
        ctr512 = counter2;

        // convert back to BE and add nonce
        counter0 = applyNonce(counter0, ctrBitMask512, templateCtr512);
        counter1 = applyNonce(counter1, ctrBitMask512, templateCtr512);
        counter2 = applyNonce(counter2, ctrBitMask512, templateCtr512);

        cpAESEncrypt3_VAES_NI(&counter0, &counter1, &counter2, pKeys, cipherRounds);

        __m512i blk0 = _mm512_loadu_si512(pSrc512);
        __m512i blk1 = _mm512_loadu_si512(pSrc512 + 1);
        __m512i blk2 = _mm512_loadu_si512(pSrc512 + 2);

        blk0 = _mm512_xor_si512(blk0, counter0);
        blk1 = _mm512_xor_si512(blk1, counter1);
        blk2 = _mm512_xor_si512(blk2, counter2);

        _mm512_storeu_si512(pDst512, blk0);
        _mm512_storeu_si512(pDst512 + 1, blk1);
        _mm512_storeu_si512(pDst512 + 2, blk2);

        pSrc512 += 3;
        pDst512 += 3;
        remainded_length -= (3 * 4 * MBS_RIJ128);
    }

    if ((4 * 2 * MBS_RIJ128) <= remainded_length) {
        __m512i counter0 = adcLo_epi64(incMsk, ctr512);
        __m512i counter1 = adcLo_epi64(M512(nextIncLoMask), counter0);

        incMsk = M512(nextIncLoMask);
        ctr512 = counter1;

        // convert back to BE and add nonce
        counter0 = applyNonce(counter0, ctrBitMask512, templateCtr512);
        counter1 = applyNonce(counter1, ctrBitMask512, templateCtr512);

        cpAESEncrypt2_VAES_NI(&counter0, &counter1, pKeys, cipherRounds);

        __m512i blk0 = _mm512_loadu_si512(pSrc512);
        __m512i blk1 = _mm512_loadu_si512(pSrc512 + 1);

        blk0 = _mm512_xor_si512(blk0, counter0);
        blk1 = _mm512_xor_si512(blk1, counter1);

        _mm512_storeu_si512(pDst512, blk0);
        _mm512_storeu_si512(pDst512 + 1, blk1);

        pSrc512 += 2;
        pDst512 += 2;
        remainded_length -= (2 * 4 * MBS_RIJ128);
    }

    for (; remainded_length >= 4 * MBS_RIJ128; remainded_length -= 4 * MBS_RIJ128) {
        __m512i counter0 = adcLo_epi64(incMsk, ctr512);

        incMsk = M512(nextIncLoMask);
        ctr512 = counter0;

        // convert back to BE and add nonce
        counter0 = applyNonce(counter0, ctrBitMask512, templateCtr512);

        cpAESEncrypt1_VAES_NI(&counter0, pKeys, cipherRounds);

        __m512i blk0 = _mm512_loadu_si512(pSrc512);
        blk0         = _mm512_xor_si512(blk0, counter0);
        _mm512_storeu_si512(pDst512, blk0);

        pSrc512 += 1;
        pDst512 += 1;
    }

    if (remainded_length) {
        __m512i counter0 = _mm512_add_epi64(incMsk, ctr512);
        __mmask8 overMsk = _mm512_cmplt_epu64_mask(counter0, ctr512);
        overMsk <<= 1;
        counter0 = _mm512_mask_add_epi64(counter0, overMsk, counter0, M512(incHiByOneMask));

        incMsk = M512(nextIncLoMask);
        ctr512 = counter0;

        // convert back to BE and add nonce
        counter0 = _mm512_shuffle_epi8(counter0, M512(swapBytes));
        counter0 = _mm512_and_epi64(counter0, ctrBitMask512);
        counter0 = _mm512_or_epi64(counter0, templateCtr512);

        cpAESEncrypt1_VAES_NI(&counter0, pKeys, cipherRounds);

        __mmask64 rw_mask = (__mmask64)((1LL << remainded_length) - 1);

        __m512i blk0 = _mm512_maskz_loadu_epi8(rw_mask, pSrc512);
        blk0         = _mm512_xor_si512(blk0, counter0);
        _mm512_mask_storeu_epi8(pDst512, rw_mask, blk0);
    }

    // return last counter
    int blocks         = remainded_length % MBS_RIJ128 == 0 ? remainded_length / MBS_RIJ128
                                                            : (remainded_length / MBS_RIJ128) + 1;
    __mmask8 lastCtrK8 = blocks == 0 ? 0xC0 : (__mmask8)((Ipp8u)0x03 << ((blocks - 1) << 1));
    __mmask64 lastCtrK64 =
        blocks == 0 ? 0xFFFF000000000000 : (__mmask64)((Ipp64u)0xFFFF << ((blocks - 1) << 4));

    ctr512 = adcLo_epi64(M512(incLoByOneMask), ctr512);

    ctr512 = _mm512_maskz_shuffle_epi8(lastCtrK64, ctr512, M512(swapBytes));
    ctr512 = _mm512_maskz_and_epi64(lastCtrK8, ctr512, ctrBitMask512);
    ctr512 = _mm512_maskz_or_epi64(lastCtrK8, ctr512, templateCtr512);

    _mm512_mask_compressstoreu_epi64(pCtrValue, lastCtrK8, ctr512);
}

#endif /* #if (_IPP32E>=_IPP32E_K1) */
