/*
 *  Copyright (C) 2017 - This file is part of libecc project
 *
 *  Authors:
 *      Ryad BENADJILA <ryadbenadjila@gmail.com>
 *      Arnaud EBALARD <arnaud.ebalard@ssi.gouv.fr>
 *      Jean-Pierre FLORI <jean-pierre.flori@ssi.gouv.fr>
 *
 *  Contributors:
 *      Nicolas VIVET <nicolas.vivet@ssi.gouv.fr>
 *      Karim KHALFALLAH <karim.khalfallah@ssi.gouv.fr>
 *
 *  This software is licensed under a dual BSD and GPL v2 license.
 *  See LICENSE file at the root folder of the project.
 */
#include <libecc/nn/nn_mul_redc1.h>
#include <libecc/nn/nn_mul_public.h>
#include <libecc/nn/nn_add.h>
#include <libecc/nn/nn_logical.h>
#include <libecc/nn/nn_div_public.h>
#include <libecc/nn/nn_modinv.h>
#include <libecc/nn/nn.h>

/*
 * Given an odd number p, compute Montgomery coefficients r, r_square
 * as well as mpinv so that:
 *
 *      - r = 2^p_rounded_bitlen mod (p), where
 *        p_rounded_bitlen = BIT_LEN_WORDS(p) (i.e. bit length of
 *        minimum number of words required to store p)
 *      - r_square = r^2 mod (p)
 *      - mpinv = -p^-1 mod (2^WORDSIZE).
 *
 * Aliasing of outputs with the input is possible since p_in is
 * copied in local p at the beginning of the function.
 *
 * The function returns 0 on success, -1 on error. out parameters 'r',
 * 'r_square' and 'mpinv' are only meaningful on success.
 */
int nn_compute_redc1_coefs(nn_t r, nn_t r_square, nn_src_t p_in, word_t *mpinv)
{
        bitcnt_t p_rounded_bitlen;
        nn p, tmp_nn1, tmp_nn2;
        word_t _mpinv;
        int ret, isodd;
        p.magic = tmp_nn1.magic = tmp_nn2.magic = WORD(0);

        ret = nn_check_initialized(p_in); EG(ret, err);
        ret = nn_init(&p, 0); EG(ret, err);
        ret = nn_copy(&p, p_in); EG(ret, err);
        MUST_HAVE((mpinv != NULL), ret, err);

        /*
         * In order for our reciprocal division routines to work, it is
         * expected that the bit length (including leading zeroes) of
         * input prime p is >= 2 * wlen where wlen is the number of bits
         * of a word size.
         */
        if (p.wlen < 2) {
                ret = nn_set_wlen(&p, 2); EG(ret, err);
        }

        ret = nn_init(r, 0); EG(ret, err);
        ret = nn_init(r_square, 0); EG(ret, err);
        ret = nn_init(&tmp_nn1, 0); EG(ret, err);
        ret = nn_init(&tmp_nn2, 0); EG(ret, err);

        /* p_rounded_bitlen = bitlen of p rounded to word size */
        p_rounded_bitlen = (bitcnt_t)(WORD_BITS * p.wlen);

        /* _mpinv = 2^wlen - (modinv(prime, 2^wlen)) */
        ret = nn_set_wlen(&tmp_nn1, 2); EG(ret, err);
        tmp_nn1.val[1] = WORD(1);
        ret = nn_copy(&tmp_nn2, &tmp_nn1); EG(ret, err);
        ret = nn_modinv_2exp(&tmp_nn1, &p, WORD_BITS, &isodd); EG(ret, err);
        ret = nn_sub(&tmp_nn1, &tmp_nn2, &tmp_nn1); EG(ret, err);
        _mpinv = tmp_nn1.val[0];

        /* r = (0x1 << p_rounded_bitlen) (p) */
        ret = nn_one(r); EG(ret, err);
        ret = nn_lshift(r, r, p_rounded_bitlen); EG(ret, err);
        ret = nn_mod(r, r, &p); EG(ret, err);

        /*
         * r_square = (0x1 << (2*p_rounded_bitlen)) (p)
         * We are supposed to handle NN numbers of size  at least two times
         * the biggest prime we use. Thus, we should be able to compute r_square
         * with a multiplication followed by a reduction. (NB: we cannot use our
         * Montgomery primitives at this point since we are computing its
         * constants!)
         */
        /* Check we have indeed enough space for our r_square computation */
        MUST_HAVE(!(NN_MAX_BIT_LEN < (2 * p_rounded_bitlen)), ret, err);

        ret = nn_sqr(r_square, r); EG(ret, err);
        ret = nn_mod(r_square, r_square, &p); EG(ret, err);

        (*mpinv) = _mpinv;

err:
        nn_uninit(&p);
        nn_uninit(&tmp_nn1);
        nn_uninit(&tmp_nn2);

        return ret;
}

/*
 * Perform Montgomery multiplication, that is usual multiplication
 * followed by reduction modulo p.
 *
 * Inputs are supposed to be < p (i.e. taken modulo p).
 *
 * This uses the CIOS algorithm from Koc et al.
 *
 * The p input is the modulo number of the Montgomery multiplication,
 * and mpinv is -p^(-1) mod (2^WORDSIZE).
 *
 * The function does not check input parameters are initialized. This
 * MUST be done by the caller.
 *
 * The function returns 0 on success, -1 on error.
 */
ATTRIBUTE_WARN_UNUSED_RET static int _nn_mul_redc1(nn_t out, nn_src_t in1, nn_src_t in2, nn_src_t p,
                         word_t mpinv)
{
        word_t prod_high, prod_low, carry, acc, m;
        unsigned int i, j, len, len_mul;
        /* a and b inputs such that len(b) <= len(a) */
        nn_src_t a, b;
        int ret, cmp;
        u8 old_wlen;

        /*
         * These comparisons are input hypothesis and does not "break"
         * the following computation. However performance loss exists
         * when this check is always done, this is why we use our
         * SHOULD_HAVE primitive.
         */
        SHOULD_HAVE((!nn_cmp(in1, p, &cmp)) && (cmp < 0), ret, err);
        SHOULD_HAVE((!nn_cmp(in2, p, &cmp)) && (cmp < 0), ret, err);

        ret = nn_init(out, 0); EG(ret, err);

        /* Check which one of in1 or in2 is the biggest */
        a = (in1->wlen <= in2->wlen) ? in2 : in1;
        b = (in1->wlen <= in2->wlen) ? in1 : in2;

        /*
         * The inputs might have been reduced due to trimming
         * because of leading zeroes. It is important for our
         * Montgomery algorithm to work on sizes consistent with
         * out prime p real bit size. Thus, we expand the output
         * size to the size of p.
         */
        ret = nn_set_wlen(out, p->wlen); EG(ret, err);

        len = out->wlen;
        len_mul = b->wlen;
        /*
         * We extend out to store carries. We first check that we
         * do not have an overflow on the NN size.
         */
        MUST_HAVE(((WORD_BITS * (out->wlen + 1)) <= NN_MAX_BIT_LEN), ret, err);
        old_wlen = out->wlen;
        out->wlen = (u8)(out->wlen + 1);

        /*
         * This can be skipped if the first iteration of the for loop
         * is separated.
         */
        for (i = 0; i < out->wlen; i++) {
                out->val[i] = 0;
        }
        for (i = 0; i < len; i++) {
                carry = WORD(0);
                for (j = 0; j < len_mul; j++) {
                        WORD_MUL(prod_high, prod_low, a->val[i], b->val[j]);
                        prod_low  = (word_t)(prod_low + carry);
                        prod_high = (word_t)(prod_high + (prod_low < carry));
                        out->val[j] = (word_t)(out->val[j] + prod_low);
                        carry = (word_t)(prod_high + (out->val[j] < prod_low));
                }
                for (; j < len; j++) {
                        out->val[j] = (word_t)(out->val[j] + carry);
                        carry = (word_t)(out->val[j] < carry);
                }
                out->val[j] = (word_t)(out->val[j] + carry);
                acc = (word_t)(out->val[j] < carry);

                m = (word_t)(out->val[0] * mpinv);
                WORD_MUL(prod_high, prod_low, m, p->val[0]);
                prod_low = (word_t)(prod_low + out->val[0]);
                carry = (word_t)(prod_high + (prod_low < out->val[0]));
                for (j = 1; j < len; j++) {
                        WORD_MUL(prod_high, prod_low, m, p->val[j]);
                        prod_low  = (word_t)(prod_low + carry);
                        prod_high = (word_t)(prod_high + (prod_low < carry));
                        out->val[j - 1] = (word_t)(prod_low + out->val[j]);
                        carry = (word_t)(prod_high + (out->val[j - 1] < prod_low));
                }
                out->val[j - 1] = (word_t)(carry + out->val[j]);
                carry = (word_t)(out->val[j - 1] < out->val[j]);
                out->val[j] = (word_t)(acc + carry);
        }
        /*
         * Note that at this stage the msw of out is either 0 or 1.
         * If out > p we need to subtract p from out.
         */
        ret = nn_cmp(out, p, &cmp); EG(ret, err);
        ret = nn_cnd_sub(cmp >= 0, out, out, p); EG(ret, err);
        MUST_HAVE((!nn_cmp(out, p, &cmp)) && (cmp < 0), ret, err);
        /* We restore out wlen. */
        out->wlen = old_wlen;

err:
        return ret;
}

/*
 * Wrapper for previous function handling aliasing of one of the input
 * paramter with out, through a copy. The function does not check
 * input parameters are initialized. This MUST be done by the caller.
 */
ATTRIBUTE_WARN_UNUSED_RET static int _nn_mul_redc1_aliased(nn_t out, nn_src_t in1, nn_src_t in2,
                                 nn_src_t p, word_t mpinv)
{
        nn out_cpy;
        int ret;
        out_cpy.magic = WORD(0);

        ret = _nn_mul_redc1(&out_cpy, in1, in2, p, mpinv); EG(ret, err);
        ret = nn_init(out, out_cpy.wlen); EG(ret, err);
        ret = nn_copy(out, &out_cpy);

err:
        nn_uninit(&out_cpy);

        return ret;
}

/*
 * Public version, handling possible aliasing of out parameter with
 * one of the input parameters.
 */
int nn_mul_redc1(nn_t out, nn_src_t in1, nn_src_t in2, nn_src_t p,
                 word_t mpinv)
{
        int ret;

        ret = nn_check_initialized(in1); EG(ret, err);
        ret = nn_check_initialized(in2); EG(ret, err);
        ret = nn_check_initialized(p); EG(ret, err);

        /* Handle possible output aliasing */
        if ((out == in1) || (out == in2) || (out == p)) {
                ret = _nn_mul_redc1_aliased(out, in1, in2, p, mpinv);
        } else {
                ret = _nn_mul_redc1(out, in1, in2, p, mpinv);
        }

err:
        return ret;
}

/*
 * Compute in1 * in2 mod p where in1 and in2 are numbers < p.
 * When p is an odd number, the function redcifies in1 and in2
 * parameters, does the computation and then unredcifies the
 * result. When p is an even number, we use an unoptimized mul
 * then mod operations sequence.
 *
 * From a mathematical standpoint, the computation is equivalent
 * to performing:
 *
 *   nn_mul(&tmp2, in1, in2);
 *   nn_mod(&out, &tmp2, q);
 *
 * but the modular reduction is done progressively during
 * Montgomery reduction when p is odd (which brings more efficiency).
 *
 * Inputs are supposed to be < p (i.e. taken modulo p).
 *
 * The function returns 0 on success, -1 on error.
 */
int nn_mod_mul(nn_t out, nn_src_t in1, nn_src_t in2, nn_src_t p_in)
{
        nn r_square, p;
        nn in1_tmp, in2_tmp;
        word_t mpinv;
        int ret, isodd;
        r_square.magic = in1_tmp.magic = in2_tmp.magic = p.magic = WORD(0);

        /* When p_in is even, we cannot work with Montgomery multiplication */
        ret = nn_isodd(p_in, &isodd); EG(ret, err);
        if(!isodd){
                /* When p_in is even, we fallback to less efficient mul then mod */
                ret = nn_mul(out, in1, in2); EG(ret, err);
                ret = nn_mod(out, out, p_in); EG(ret, err);
        }
        else{
                /* Here, p_in is odd and we can use redcification */
                ret = nn_copy(&p, p_in); EG(ret, err);

                /*
                 * In order for our reciprocal division routines to work, it is
                 * expected that the bit length (including leading zeroes) of
                 * input prime p is >= 2 * wlen where wlen is the number of bits
                 * of a word size.
                 */
                if (p.wlen < 2) {
                        ret = nn_set_wlen(&p, 2); EG(ret, err);
                }

                /* Compute Mongtomery coefs.
                 * NOTE: in1_tmp holds a dummy value here after the operation.
                 */
                ret = nn_compute_redc1_coefs(&in1_tmp, &r_square, &p, &mpinv); EG(ret, err);

                /* redcify in1 and in2 */
                ret = nn_mul_redc1(&in1_tmp, in1, &r_square, &p, mpinv); EG(ret, err);
                ret = nn_mul_redc1(&in2_tmp, in2, &r_square, &p, mpinv); EG(ret, err);

                /* Compute in1 * in2 mod p in montgomery world.
                 * NOTE: r_square holds the result after the operation.
                 */
                ret = nn_mul_redc1(&r_square, &in1_tmp, &in2_tmp, &p, mpinv); EG(ret, err);

                /* Come back to real world by unredcifying result */
                ret = nn_init(&in1_tmp, 0); EG(ret, err);
                ret = nn_one(&in1_tmp); EG(ret, err);
                ret = nn_mul_redc1(out, &r_square, &in1_tmp, &p, mpinv); EG(ret, err);
        }

err:
        nn_uninit(&p);
        nn_uninit(&r_square);
        nn_uninit(&in1_tmp);
        nn_uninit(&in2_tmp);

        return ret;
}