/*************************************************************************
* 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:
//     EncryptStreamCTR32_VAES_NI
//
*/

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

#if (_IPP32E >= _IPP32E_K1)

/* Mask to convert low 32-bit parts of four 128-bit Big-Endian numbers
 * stored in 512-bit register. The low 32-bit of each number converted
 * from LE to BE */
static __ALIGN64 Ipp8u swapBytes[] = { 0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10, 11, 15,
                                       14, 13, 12, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25,
                                       26, 27, 31, 30, 29, 28, 32, 33, 34, 35, 36, 37, 38,
                                       39, 40, 41, 42, 43, 47, 46, 45, 44, 48, 49, 50, 51,
                                       52, 53, 54, 55, 56, 57, 58, 59, 63, 62, 61, 60 };

////////////////////////////////////////////////////////////////////////////////

/* clang-format off */
IPP_OWN_DEFN(void, EncryptStreamCTR32_VAES_NI, (const Ipp8u* pSrc,
                                                Ipp8u* pDst,
                                                int nr,
                                                const Ipp8u* pRKey,
                                                int length,  /* message length in bytes   */
                                                Ipp8u* pIV)) /* BE counter representation */
/* clang-format on */
{
    int cipherRounds = nr - 1;

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

    Ipp32u* pIV32 = (Ipp32u*)pIV;
    Ipp64u* pIV64 = (Ipp64u*)pIV;

    // start increment mask for IV
    __m512i startIncMask = _mm512_set_epi32(0x3,
                                            0x0,
                                            0x0,
                                            0x0,
                                            0x2,
                                            0x0,
                                            0x0,
                                            0x0,
                                            0x1,
                                            0x0,
                                            0x0,
                                            0x0,
                                            0x0,
                                            0x0,
                                            0x0,
                                            0x0);

    // continuous increment mask for IV
    __m512i incMask = _mm512_set_epi32(0x4,
                                       0x0,
                                       0x0,
                                       0x0,
                                       0x4,
                                       0x0,
                                       0x0,
                                       0x0,
                                       0x4,
                                       0x0,
                                       0x0,
                                       0x0,
                                       0x4,
                                       0x0,
                                       0x0,
                                       0x0);

    // initial BE counter with 32-bit low part converted to LE:
    // pIV: f0 f1 f2 f3 f4 f5 f6 f7 f8 f9 fa fb | fc fd fe ff
    // IV:  f0 f1 f2 f3 f4 f5 f6 f7 f8 f9 fa fb | ff fe fd fc
    __m512i IV = _mm512_set4_epi32((Ipp32s)ENDIANNESS32(pIV32[3]),
                                   (Ipp32s)pIV32[2],
                                   (Ipp32s)pIV32[1],
                                   (Ipp32s)pIV32[0]);


    // Update IV counter for next function calls
    int blocks = length / MBS_RIJ128;

    Ipp64u ctr64_h = ENDIANNESS64(pIV64[0]); // high 64-bit of BE IV converted to LE
    Ipp64u ctr64_l = ENDIANNESS64(pIV64[1]); // low 32-bit of BE pIV converted to LE

    ctr64_l += (Ipp64u)blocks;
    if (ctr64_l < (Ipp64u)blocks) {          // overflow of low part
        ctr64_h += 1;
    }

    // update IV
    pIV64[0] = ENDIANNESS64(ctr64_h);
    pIV64[1] = ENDIANNESS64(ctr64_l);

    //--------------------------------------

    for (blocks = length / MBS_RIJ128; blocks >= (4 * 4); blocks -= (4 * 4)) {
        // counter0:  f0 f1 f2 f3 f4 f5 f6 f7 f8 f9 fa fb | ff fe fd fc
        // counter1:  f0 f1 f2 f3 f4 f5 f6 f7 f8 f9 fa fb | 00 ff fd fc
        // counter2:  f0 f1 f2 f3 f4 f5 f6 f7 f8 f9 fa fb | 01 ff fd fc
        // counter3:  f0 f1 f2 f3 f4 f5 f6 f7 f8 f9 fa fb | 02 ff fd fc
        __m512i counter0 = _mm512_add_epi32(startIncMask, IV);
        __m512i counter1 = _mm512_add_epi32(incMask, counter0);
        __m512i counter2 = _mm512_add_epi32(incMask, counter1);
        __m512i counter3 = _mm512_add_epi32(incMask, counter2);

        startIncMask = incMask;
        IV           = counter3;

        // convert last 32-bit (LE->BE):
        // counter0: f0 f1 f2 f3 f4 f5 f6 f7 f8 f9 fa fb | fc fd fe ff
        // counter1: f0 f1 f2 f3 f4 f5 f6 f7 f8 f9 fa fb | fc fd ff 00
        // counter2: f0 f1 f2 f3 f4 f5 f6 f7 f8 f9 fa fb | fc fd ff 01
        // counter3: f0 f1 f2 f3 f4 f5 f6 f7 f8 f9 fa fb | fc fd ff 02
        counter0 = _mm512_shuffle_epi8(counter0, *((__m512i*)swapBytes));
        counter1 = _mm512_shuffle_epi8(counter1, *((__m512i*)swapBytes));
        counter2 = _mm512_shuffle_epi8(counter2, *((__m512i*)swapBytes));
        counter3 = _mm512_shuffle_epi8(counter3, *((__m512i*)swapBytes));

        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) <= blocks) {
        __m512i counter0 = _mm512_add_epi32(startIncMask, IV);
        __m512i counter1 = _mm512_add_epi32(incMask, counter0);
        __m512i counter2 = _mm512_add_epi32(incMask, counter1);

        startIncMask = incMask;
        IV           = counter2;

        // convert last 32-bit (LE->BE)
        counter0 = _mm512_shuffle_epi8(counter0, *((__m512i*)swapBytes));
        counter1 = _mm512_shuffle_epi8(counter1, *((__m512i*)swapBytes));
        counter2 = _mm512_shuffle_epi8(counter2, *((__m512i*)swapBytes));

        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;
        blocks -= (3 * 4);
    }

    if ((4 * 2) <= blocks) {
        __m512i counter0 = _mm512_add_epi32(startIncMask, IV);
        __m512i counter1 = _mm512_add_epi32(incMask, counter0);

        startIncMask = incMask;
        IV           = counter1;

        // convert last 32-bit (LE->BE)
        counter0 = _mm512_shuffle_epi8(counter0, *((__m512i*)swapBytes));
        counter1 = _mm512_shuffle_epi8(counter1, *((__m512i*)swapBytes));

        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;
        blocks -= (2 * 4);
    }

    for (; blocks >= 4; blocks -= 4) {
        __m512i counter0 = _mm512_add_epi32(startIncMask, IV);

        startIncMask = incMask;
        IV           = counter0;

        // convert last 32-bit (LE->BE)
        counter0 = _mm512_shuffle_epi8(counter0, *((__m512i*)swapBytes));

        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 (blocks) {
        __mmask8 k8   = (__mmask8)((1 << (blocks + blocks)) - 1); // 64-bit chunks
        __mmask16 k16 = (__mmask16)((1 << (blocks << 2)) - 1);    // 32-bit chunks
        __mmask64 k64 = (__mmask64)((1LL << (blocks << 4)) - 1);  // 8-bit chunks

        __m512i counter0 = _mm512_maskz_add_epi32(k16, startIncMask, IV);

        // swap last 32-bit (LE->BE)
        counter0 = _mm512_maskz_shuffle_epi8(k64, counter0, *((__m512i*)swapBytes));

        cpAESEncrypt1_VAES_NI(&counter0, pKeys, cipherRounds);

        __m512i blk0 = _mm512_maskz_loadu_epi64(k8, pSrc512);
        blk0         = _mm512_maskz_xor_epi64(k8, blk0, counter0);
        _mm512_mask_storeu_epi64(pDst512, k8, blk0);
    }
}

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