/*************************************************************************
* Copyright (C) 2019 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 "owncp.h"

#if (_IPP32E >= _IPP32E_K1)

#include "pcptool.h"

#include "pcpngmontexpstuff_avx512.h"
#include "ifma_rsa_arith.h"

IPP_OWN_DEFN(cpSize, gsMontDualExpWinBuffer_avx512, (int modulusBits))
{
    cpSize redNum = numofVariable_avx512(modulusBits);
    /* Reg (ymm) capacity = 4 qwords */
    cpSize redBufferNum = numofVariableBuff_avx512(redNum, 4);
    return redBufferNum * 2 * 8;
}

/* clang-format off */
IPP_OWN_FUNPTR(void, AMM52, (Ipp64u* out,
                             const Ipp64u* a,
                             const Ipp64u* b,
                             const Ipp64u* m,
                             Ipp64u k0))
IPP_OWN_FUNPTR(void, DEXP52, (Ipp64u* out,
                              const Ipp64u* base,
                              const Ipp64u* exp[2],
                              const Ipp64u* modulus,
                              const Ipp64u* pMont,
                              const Ipp64u k0[2]))


IPP_OWN_DEFN(cpSize, gsMontDualExpWin_BNU_sscm_avx512, (BNU_CHUNK_T* dataY[2],
                                                        const BNU_CHUNK_T* dataX[2],
                                                        cpSize nsX[2],
                                                        const BNU_CHUNK_T* dataE[2],
                                                        gsModEngine* pMont[2],
                                                        BNU_CHUNK_T* pBuffer))
/* clang-format on */
{
    const BNU_CHUNK_T* modulus[2] = { 0 };
    const BNU_CHUNK_T* rr[2]      = { 0 };
    cpSize modulusSize[2]         = { 0 };
    BNU_CHUNK_T k0[2]             = { 0 };

    modulus[0]     = MOD_MODULUS(pMont[0]);
    modulus[1]     = MOD_MODULUS(pMont[1]);
    modulusSize[0] = MOD_LEN(pMont[0]);
    modulusSize[1] = MOD_LEN(pMont[1]);
    rr[0]          = MOD_MNT_R2(pMont[0]);
    rr[1]          = MOD_MNT_R2(pMont[1]);
    k0[0]          = MOD_MNT_FACTOR(pMont[0]);
    k0[1]          = MOD_MNT_FACTOR(pMont[1]);

    /*
    * Dual expo implemenentation assumes that bit sizes of P and Q are
    * the same, so query size from the first.
    */
    int modulusBitSize = BITSIZE_BNU(modulus[0], modulusSize[0]);
    int redLen         = NUMBER_OF_DIGITS(modulusBitSize + 2, EXP_DIGIT_SIZE_AVX512);
    /* For ymm-based implementation reg capacity = 4 qwords */
    int redBufferLen = numofVariableBuff_avx512(redLen, 4);

    /* Allocate buffers */
    BNU_CHUNK_T* redX      = pBuffer;
    BNU_CHUNK_T* redM      = redX + 2 * redBufferLen;
    BNU_CHUNK_T* redRR     = redM + 2 * redBufferLen;
    BNU_CHUNK_T* redCoeff  = redRR + 2 * redBufferLen;
    BNU_CHUNK_T* redBuffer = redCoeff + redBufferLen;

    AMM52 ammFunc   = NULL;
    DEXP52 dexpFunc = NULL;
    switch (modulusBitSize) {
    case 1024:
        ammFunc  = ifma256_amm52x20;
        dexpFunc = (DEXP52)ifma256_exp52x20_dual;
        break;
    case 1536:
        ammFunc  = ifma256_amm52x30;
        dexpFunc = (DEXP52)ifma256_exp52x30_dual;
        break;
    case 2048:
        ammFunc  = ifma256_amm52x40;
        dexpFunc = (DEXP52)ifma256_exp52x40_dual;
        break;
    default:
        /* Other modulus sizes not supported. This function shall not be called for them. */
        return 0;
    }

    ZEXPAND_BNU(redCoeff, 0, redBufferLen);
    int conv_coeff = 4 * (EXP_DIGIT_SIZE_AVX512 * redLen - modulusBitSize);
    /* Set corresponding bit in reduced domain */
    SET_BIT(redCoeff, 64 * (int)(conv_coeff / 52) + conv_coeff % 52);

    for (int i = 0; i < 2; i++) {
        /* Convert base into redundant domain */
        ZEXPAND_COPY_BNU(redBuffer, redBufferLen, dataX[i], nsX[i]);
        regular_dig52(redX + i * redBufferLen, redBufferLen, redBuffer, modulusBitSize);

        /* Convert modulus into redundant domain */
        ZEXPAND_COPY_BNU(redBuffer, redBufferLen, modulus[i], modulusSize[i]);
        regular_dig52(redM + i * redBufferLen, redBufferLen, redBuffer, modulusBitSize);

        /*
       * Compute target domain Montgomery converter RR' based on original domain RR.
       *
       * Example: modlen = 1024: RR = 2^2048 mod m, RR' = 2^2080 mod m
       *     conv_coeff = 2^64
       *     (1st amm): 2^2048*2^2048/2^1040= 2^3056 mod m
       *     (2nd amm): 2^3056*2^64/2^1040 = 2^2080 mod m
       *
       */
        ZEXPAND_COPY_BNU(redBuffer, redBufferLen, rr[i], modulusSize[i]);
        regular_dig52(redRR + i * redBufferLen, redBufferLen, redBuffer, modulusBitSize);
        ammFunc(redRR + i * redBufferLen,
                redRR + i * redBufferLen,
                redRR + i * redBufferLen,
                redM + i * redBufferLen,
                k0[i]);
        ammFunc(redRR + i * redBufferLen,
                redRR + i * redBufferLen,
                redCoeff,
                redM + i * redBufferLen,
                k0[i]);
    }

    dexpFunc(redRR, redX, dataE, redM, redRR, k0);

    /* Convert result back to regular domain */
    for (int i = 0; i < 2; i++)
        dig52_regular(dataY[i], redRR + i * redBufferLen, modulusBitSize);

    /* Clear redundant exponents buffer */
    PurgeBlock(redRR, 2 * redBufferLen * (int)sizeof(BNU_CHUNK_T));

    return (modulusSize[0] + modulusSize[1]);
}

#endif
