/*
 *  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_modinv.h>
#include <libecc/nn/nn_div_public.h>
#include <libecc/nn/nn_logical.h>
#include <libecc/nn/nn_add.h>
#include <libecc/nn/nn_mod_pow.h>
#include <libecc/nn/nn.h>
/* Include the "internal" header as we use non public API here */
#include "../nn/nn_mul.h"

/*
 * Compute out = x^-1 mod m, i.e. out such that (out * x) = 1 mod m
 * out is initialized by the function, i.e. caller needs not initialize
 * it; only provide the associated storage space. Done in *constant
 * time* if underlying routines are.
 *
 * Asserts that m is odd and that x is smaller than m.
 * This second condition is not strictly necessary,
 * but it allows to perform all operations on nn's of the same length,
 * namely the length of m.
 *
 * Uses a binary xgcd algorithm,
 * only keeps track of coefficient for inverting x,
 * and performs reduction modulo m at each step.
 *
 * This does not normalize out on return.
 *
 * 0 is returned on success (everything went ok and x has reciprocal), -1
 * on error or if x has no reciprocal. On error, out is not meaningful.
 *
 * The function is an internal helper: caller MUST check params have been
 * initialized, i.e. this is not done by the function.
 *
 */
ATTRIBUTE_WARN_UNUSED_RET static int _nn_modinv_odd(nn_t out, nn_src_t x, nn_src_t m)
{
	int isodd, swap, smaller, ret, cmp, iszero, tmp_isodd;
	nn a, b, u, tmp, mp1d2;
	nn_t uu = out;
	bitcnt_t cnt;
	a.magic = b.magic = u.magic = tmp.magic = mp1d2.magic = WORD(0);

	ret = nn_init(out, 0); EG(ret, err);
	ret = nn_init(&a, (u16)(m->wlen * WORD_BYTES)); EG(ret, err);
	ret = nn_init(&b, (u16)(m->wlen * WORD_BYTES)); EG(ret, err);
	ret = nn_init(&u, (u16)(m->wlen * WORD_BYTES)); EG(ret, err);
	ret = nn_init(&mp1d2, (u16)(m->wlen * WORD_BYTES)); EG(ret, err);
	/*
	 * Temporary space needed to only deal with positive stuff.
	 */
	ret = nn_init(&tmp, (u16)(m->wlen * WORD_BYTES)); EG(ret, err);

	MUST_HAVE((!nn_isodd(m, &isodd)) && isodd, ret, err);
	MUST_HAVE((!nn_cmp(x, m, &cmp)) && (cmp < 0), ret, err);
	MUST_HAVE((!nn_iszero(x, &iszero)) && (!iszero), ret, err);

	/*
	 * Maintain:
	 *
	 * a = u * x (mod m)
	 * b = uu * x (mod m)
	 *
	 * and b odd at all times. Initially,
	 *
	 * a = x, u = 1
	 * b = m, uu = 0
	 */
	ret = nn_copy(&a, x); EG(ret, err);
	ret = nn_set_wlen(&a, m->wlen); EG(ret, err);
	ret = nn_copy(&b, m); EG(ret, err);
	ret = nn_one(&u); EG(ret, err);
	ret = nn_zero(uu); EG(ret, err);

	/*
	 * The lengths of u and uu should not affect constant timeness but it
	 * does not hurt to set them already.
	 * They will always be strictly smaller than m.
	 */
	ret = nn_set_wlen(&u, m->wlen); EG(ret, err);
	ret = nn_set_wlen(uu, m->wlen); EG(ret, err);

	/*
	 * Precompute inverse of 2 mod m:
	 *	2^-1 = (m+1)/2
	 * computed as (m >> 1) + 1.
	 */
	ret = nn_rshift_fixedlen(&mp1d2, m, 1); EG(ret, err);

	ret = nn_inc(&mp1d2, &mp1d2); EG(ret, err); /* no carry can occur here
						       because of prev. shift */

	cnt = (bitcnt_t)((a.wlen + b.wlen) * WORD_BITS);
	while (cnt > 0) {
		cnt = (bitcnt_t)(cnt - 1);
		/*
		 * Always maintain b odd. The logic of the iteration is as
		 * follows.
		 */

		/*
		 * For a, b:
		 *
		 * odd = a & 1
		 * swap = odd & (a < b)
		 * if (swap)
		 *      swap(a, b)
		 * if (odd)
		 *      a -= b
		 * a /= 2
		 */

		MUST_HAVE((!nn_isodd(&b, &tmp_isodd)) && tmp_isodd, ret, err);

		ret = nn_isodd(&a, &isodd); EG(ret, err);
		ret = nn_cmp(&a, &b, &cmp); EG(ret, err);
		swap = isodd & (cmp == -1);

		ret = nn_cnd_swap(swap, &a, &b); EG(ret, err);
		ret = nn_cnd_sub(isodd, &a, &a, &b); EG(ret, err);

		MUST_HAVE((!nn_isodd(&a, &tmp_isodd)) && (!tmp_isodd), ret, err); /* a is now even */

		ret = nn_rshift_fixedlen(&a, &a, 1);  EG(ret, err);/* division by 2 */

		/*
		 * For u, uu:
		 *
		 * if (swap)
		 *      swap u, uu
		 * smaller = (u < uu)
		 * if (odd)
		 *      if (smaller)
		 *              u += m - uu
		 *      else
		 *              u -= uu
		 * u >>= 1
		 * if (u was odd)
		 *      u += (m+1)/2
		 */
		ret = nn_cnd_swap(swap, &u, uu); EG(ret, err);

		/* This parameter is used to avoid handling negative numbers. */
		ret = nn_cmp(&u, uu, &cmp); EG(ret, err);
		smaller = (cmp == -1);

		/* Computation of 'm - uu' can always be performed. */
		ret = nn_sub(&tmp, m, uu); EG(ret, err);

		/* Selection btw 'm-uu' and '-uu' is made by the following function calls. */
		ret = nn_cnd_add(isodd & smaller, &u, &u, &tmp); EG(ret, err); /* no carry can occur as 'u+(m-uu) = m-(uu-u) < m' */
		ret = nn_cnd_sub(isodd & (!smaller), &u, &u, uu); EG(ret, err);

		/* Divide u by 2 */
		ret = nn_isodd(&u, &isodd); EG(ret, err);
		ret = nn_rshift_fixedlen(&u, &u, 1); EG(ret, err);
		ret = nn_cnd_add(isodd, &u, &u, &mp1d2); EG(ret, err); /* no carry can occur as u=1+u' with u'<m-1 and u' even so u'/2+(m+1)/2<(m-1)/2+(m+1)/2=m */

		MUST_HAVE((!nn_cmp(&u, m, &cmp)) && (cmp < 0), ret, err);
		MUST_HAVE((!nn_cmp(uu, m, &cmp)) && (cmp < 0), ret, err);

		/*
		 * As long as a > 0, the quantity
		 * (bitsize of a) + (bitsize of b)
		 * is reduced by at least one bit per iteration,
		 * hence after (bitsize of x) + (bitsize of m) - 1
		 * iterations we surely have a = 0. Then b = gcd(x, m)
		 * and if b = 1 then also uu = x^{-1} (mod m).
		 */
	}

	MUST_HAVE((!nn_iszero(&a, &iszero)) && iszero, ret, err);

	/* Check that gcd is one. */
	ret = nn_cmp_word(&b, WORD(1), &cmp); EG(ret, err);

	/* If not, set "inverse" to zero. */
	ret = nn_cnd_sub(cmp != 0, uu, uu, uu); EG(ret, err);

	ret = cmp ? -1 : 0;

err:
	nn_uninit(&a);
	nn_uninit(&b);
	nn_uninit(&u);
	nn_uninit(&mp1d2);
	nn_uninit(&tmp);

	PTR_NULLIFY(uu);

	return ret;
}

/*
 * Same as above without restriction on m.
 * No attempt to make it constant time.
 * Uses the above constant-time binary xgcd when m is odd
 * and a not constant time plain Euclidean xgcd when m is even.
 *
 * _out parameter need not be initialized; this will be done by the function.
 * x and m must be initialized nn.
 *
 * Return -1 on error or if if x has no reciprocal modulo m. out is zeroed.
 * Return  0 if x has reciprocal modulo m.
 *
 * The function supports aliasing.
 */
int nn_modinv(nn_t _out, nn_src_t x, nn_src_t m)
{
	int sign, ret, cmp, isodd, isone;
	nn_t x_mod_m;
	nn u, v, out; /* Out to support aliasing */
	out.magic = u.magic = v.magic = WORD(0);

	ret = nn_check_initialized(x); EG(ret, err);
	ret = nn_check_initialized(m); EG(ret, err);

	/* Initialize out */
	ret = nn_init(&out, 0); EG(ret, err);
	ret = nn_isodd(m, &isodd); EG(ret, err);
	if (isodd) {
		ret = nn_cmp(x, m, &cmp); EG(ret, err);
		if (cmp >= 0) {
			/*
			 * If x >= m, (x^-1) mod m = ((x mod m)^-1) mod m
			 * Hence, compute x mod m. In order to avoid
			 * additional stack usage, we use 'u' (not
			 * already useful at that point in the function).
			 */
			x_mod_m = &u;
			ret = nn_mod(x_mod_m, x, m); EG(ret, err);
			ret = _nn_modinv_odd(&out, x_mod_m, m); EG(ret, err);
		} else {
			ret = _nn_modinv_odd(&out, x, m); EG(ret, err);
		}
		ret = nn_copy(_out, &out);
		goto err;
	}

	/* Now m is even */
	ret = nn_isodd(x, &isodd); EG(ret, err);
	MUST_HAVE(isodd, ret, err);

	ret = nn_init(&u, 0); EG(ret, err);
	ret = nn_init(&v, 0); EG(ret, err);
	ret = nn_xgcd(&out, &u, &v, x, m, &sign); EG(ret, err);
	ret = nn_isone(&out, &isone); EG(ret, err);
	MUST_HAVE(isone, ret, err);

	ret = nn_mod(&out, &u, m); EG(ret, err);
	if (sign == -1) {
		ret = nn_sub(&out, m, &out); EG(ret, err);
	}
	ret = nn_copy(_out, &out);

err:
	nn_uninit(&out);
	nn_uninit(&u);
	nn_uninit(&v);

	PTR_NULLIFY(x_mod_m);

	return ret;
}

/*
 * Compute (A - B) % 2^(storagebitsizeof(B) + 1). A and B must be initialized nn.
 * the function is an internal helper and does not verify params have been
 * initialized; this must be done by the caller. No assumption on A and B values
 * such as A >= B. Done in *constant time. Returns 0 on success, -1 on error.
 */
ATTRIBUTE_WARN_UNUSED_RET static inline int _nn_sub_mod_2exp(nn_t A, nn_src_t B)
{
	u8 Awlen = A->wlen;
	int ret;

	ret = nn_set_wlen(A, (u8)(Awlen + 1)); EG(ret, err);

	/* Make sure A > B */
	/* NOTE: A->wlen - 1 is not an issue here thant to the nn_set_wlen above */
	A->val[A->wlen - 1] = WORD(1);
	ret = nn_sub(A, A, B); EG(ret, err);

	/* The artificial word will be cleared in the following function call */
	ret = nn_set_wlen(A, Awlen);

err:
	return ret;
}

/*
 * Invert x modulo 2^exp using Hensel lifting. Returns 0 on success, -1 on
 * error. On success, x_isodd is 1 if x is odd, 0 if it is even.
 * Please note that the result is correct (inverse of x) only when x is prime
 * to 2^exp, i.e. x is odd (x_odd is 1).
 *
 * Operations are done in *constant time*.
 *
 * Aliasing is supported.
 */
int nn_modinv_2exp(nn_t _out, nn_src_t x, bitcnt_t exp, int *x_isodd)
{
	bitcnt_t cnt;
	u8 exp_wlen = (u8)BIT_LEN_WORDS(exp);
	bitcnt_t exp_cnt = exp % WORD_BITS;
	word_t mask = (word_t)((exp_cnt == 0) ? WORD_MASK : (word_t)((WORD(1) << exp_cnt) - WORD(1)));
	nn tmp_sqr, tmp_mul;
	/* for aliasing */
	int isodd, ret;
	nn out;
	out.magic = tmp_sqr.magic = tmp_mul.magic = WORD(0);

	MUST_HAVE((x_isodd != NULL), ret, err);
	ret = nn_check_initialized(x); EG(ret, err);
	ret = nn_check_initialized(_out); EG(ret, err);

	ret = nn_init(&out, 0); EG(ret, err);
	ret = nn_init(&tmp_sqr, 0); EG(ret, err);
	ret = nn_init(&tmp_mul, 0); EG(ret, err);
	ret = nn_isodd(x, &isodd); EG(ret, err);
	if (exp == (bitcnt_t)0){
		/* Specific case of zero exponent, output 0 */
		(*x_isodd) = isodd;
		goto err;
	}
	if (!isodd) {
		ret = nn_zero(_out); EG(ret, err);
		(*x_isodd) = 0;
		goto err;
	}

	/*
	 * Inverse modulo 2.
	 */
	cnt = 1;
	ret = nn_one(&out); EG(ret, err);

	/*
	 * Inverse modulo 2^(2^i) <= 2^WORD_BITS.
	 * Assumes WORD_BITS is a power of two.
	 */
	for (; cnt < WORD_MIN(WORD_BITS, exp); cnt = (bitcnt_t)(cnt << 1)) {
		ret = nn_sqr_low(&tmp_sqr, &out, out.wlen); EG(ret, err);
		ret = nn_mul_low(&tmp_mul, &tmp_sqr, x, out.wlen); EG(ret, err);
		ret = nn_lshift_fixedlen(&out, &out, 1); EG(ret, err);

		/*
		 * Allowing "negative" results for a subtraction modulo
		 * a power of two would allow to use directly:
		 * nn_sub(out, out, tmp_mul)
		 * which is always negative in ZZ except when x is one.
		 *
		 * Another solution is to add the opposite of tmp_mul.
		 * nn_modopp_2exp(tmp_mul, tmp_mul);
		 * nn_add(out, out, tmp_mul);
		 *
		 * The current solution is to add a sufficiently large power
		 * of two to out unconditionally to absorb the potential
		 * borrow. The result modulo 2^(2^i) is correct whether the
		 * borrow occurs or not.
		 */
		ret = _nn_sub_mod_2exp(&out, &tmp_mul); EG(ret, err);
	}

	/*
	 * Inverse modulo 2^WORD_BITS < 2^(2^i) < 2^exp.
	 */
	for (; cnt < ((exp + 1) >> 1); cnt = (bitcnt_t)(cnt << 1)) {
		ret = nn_set_wlen(&out, (u8)(2 * out.wlen)); EG(ret, err);
		ret = nn_sqr_low(&tmp_sqr, &out, out.wlen); EG(ret, err);
		ret = nn_mul_low(&tmp_mul, &tmp_sqr, x, out.wlen); EG(ret, err);
		ret = nn_lshift_fixedlen(&out, &out, 1); EG(ret, err);
		ret = _nn_sub_mod_2exp(&out, &tmp_mul); EG(ret, err);
	}

	/*
	 * Inverse modulo 2^(2^i + j) >= 2^exp.
	 */
	if (exp > WORD_BITS) {
		ret = nn_set_wlen(&out, exp_wlen); EG(ret, err);
		ret = nn_sqr_low(&tmp_sqr, &out, out.wlen); EG(ret, err);
		ret = nn_mul_low(&tmp_mul, &tmp_sqr, x, out.wlen); EG(ret, err);
		ret = nn_lshift_fixedlen(&out, &out, 1); EG(ret, err);
		ret = _nn_sub_mod_2exp(&out, &tmp_mul); EG(ret, err);
	}

	/*
	 * Inverse modulo 2^exp.
	 */
	out.val[exp_wlen - 1] &= mask;

	ret = nn_copy(_out, &out); EG(ret, err);

	(*x_isodd) = 1;

err:
	nn_uninit(&out);
	nn_uninit(&tmp_sqr);
	nn_uninit(&tmp_mul);

	return ret;
}

/*
 * Invert word w modulo m.
 *
 * The function supports aliasing.
 */
int nn_modinv_word(nn_t out, word_t w, nn_src_t m)
{
	nn nn_tmp;
	int ret;
	nn_tmp.magic = WORD(0);

	ret = nn_init(&nn_tmp, 0); EG(ret, err);
	ret = nn_set_word_value(&nn_tmp, w); EG(ret, err);
	ret = nn_modinv(out, &nn_tmp, m);

err:
	nn_uninit(&nn_tmp);

	return ret;
}


/*
 * Internal function for nn_modinv_fermat and nn_modinv_fermat_redc used
 * hereafter.
 */
ATTRIBUTE_WARN_UNUSED_RET static int _nn_modinv_fermat_common(nn_t out, nn_src_t x, nn_src_t p, nn_t p_minus_two, int *lesstwo)
{
	int ret, cmp, isodd;
	nn two;
	two.magic = WORD(0);

	/* Sanity checks on inputs */
	ret = nn_check_initialized(x); EG(ret, err);
	ret = nn_check_initialized(p); EG(ret, err);
	/* NOTE: since this is an internal function, we are ensured that p_minus_two,
	 * two and regular are OK.
	 */

	/* 0 is not invertible in any case */
	ret = nn_iszero(x, &cmp); EG(ret, err);
	if(cmp){
		/* Zero the output and return an error */
		ret = nn_init(out, 0); EG(ret, err);
		ret = nn_zero(out); EG(ret, err);
		ret = -1;
		goto err;
	}

	/* For p <= 2, p being prime either p = 1 or p = 2.
	 * When p = 2, only 1 has an inverse, if p = 1 no one has an inverse.
	 */
	(*lesstwo) = 0;
	ret = nn_cmp_word(p, WORD(2), &cmp); EG(ret, err);
        if(cmp == 0){
		/* This is the p = 2 case, parity of x provides the result */
		ret = nn_isodd(x, &isodd); EG(ret, err);
		if(isodd){
			/* x is odd, 1 is its inverse */
			ret = nn_init(out, 0); EG(ret, err);
			ret = nn_one(out); EG(ret, err);
			ret = 0;
		}
		else{
			/* x is even, no inverse. Zero the output */
			ret = nn_init(out, 0); EG(ret, err);
			ret = nn_zero(out); EG(ret, err);
			ret = -1;
		}
		(*lesstwo) = 1;
		goto err;
        } else if (cmp < 0){
		/* This is the p = 1 case, no inverse here: hence return an error */
		/* Zero the output */
		ret = nn_init(out, 0); EG(ret, err);
		ret = nn_zero(out); EG(ret, err);
		ret = -1;
		(*lesstwo) = 1;
		goto err;
	}

	/* Else we compute (p-2) for the upper layer */
	if(p != p_minus_two){
		/* Handle aliasing of p and p_minus_two */
		ret = nn_init(p_minus_two, 0); EG(ret, err);
	}

	ret = nn_init(&two, 0); EG(ret, err);
	ret = nn_set_word_value(&two, WORD(2)); EG(ret, err);
	ret = nn_sub(p_minus_two, p, &two);

err:
	nn_uninit(&two);

	return ret;
}

/*
 * Invert NN x modulo p using Fermat's little theorem for our inversion:
 *
 *    p prime means that:
 *    x^(p-1) = 1 mod (p)
 *    which means that x^(p-2) mod(p) is the modular inverse of x mod (p)
 *    for x != 0
 *
 * NOTE: the input hypothesis is that p is prime.
 * XXX WARNING: using this function with p not prime will produce wrong
 * results without triggering an error!
 *
 * The function returns 0 on success, -1 on error
 * (e.g. if x has no inverse modulo p, i.e. x = 0).
 *
 * Aliasing is supported.
 */
int nn_modinv_fermat(nn_t out, nn_src_t x, nn_src_t p)
{
	int ret, lesstwo;
	nn p_minus_two;
	p_minus_two.magic = WORD(0);

	/* Call our helper.
	 * NOTE: "marginal" cases where x = 0 and p <= 2 should be caught in this helper.
	 */
	ret = _nn_modinv_fermat_common(out, x, p, &p_minus_two, &lesstwo); EG(ret, err);

	if(!lesstwo){
		/* Compute x^(p-2) mod (p) */
		ret = nn_mod_pow(out, x, &p_minus_two, p);
	}

err:
	nn_uninit(&p_minus_two);

	return ret;
}

/*
 * Invert NN x modulo m using Fermat's little theorem for our inversion.
 *
 * This is a version with already (pre)computed Montgomery coefficients.
 *
 * NOTE: the input hypothesis is that p is prime.
 * XXX WARNING: using this function with p not prime will produce wrong
 * results without triggering an error!
 *
 * The function returns 0 on success, -1 on error
 * (e.g. if x has no inverse modulo p, i.e. x = 0).
 *
 * Aliasing is supported.
 */
int nn_modinv_fermat_redc(nn_t out, nn_src_t x, nn_src_t p, nn_src_t r, nn_src_t r_square, word_t mpinv)
{
	int ret, lesstwo;
	nn p_minus_two;
	p_minus_two.magic = WORD(0);

	/* Call our helper.
	 * NOTE: "marginal" cases where x = 0 and p <= 2 should be caught in this helper.
	 */
	ret = _nn_modinv_fermat_common(out, x, p, &p_minus_two, &lesstwo); EG(ret, err);

	if(!lesstwo){
		/* Compute x^(p-2) mod (p) using precomputed Montgomery coefficients as input */
		ret = nn_mod_pow_redc(out, x, &p_minus_two, p, r, r_square, mpinv);
	}

err:
	nn_uninit(&p_minus_two);

	return ret;
}