/*
 *  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/fp/fp_sqrt.h>
#include <libecc/nn/nn_add.h>
#include <libecc/nn/nn_logical.h>

/*
 * Compute the legendre symbol of an element of Fp:
 *
 *   Legendre(a) = a^((p-1)/2) (p) = { -1, 0, 1 }
 *
 */
ATTRIBUTE_WARN_UNUSED_RET static int legendre(fp_src_t a)
{
        int ret, iszero, cmp;
        fp pow; /* The result if the exponentiation is in Fp */
        fp one; /* The element 1 in the field */
        nn exp; /* The power exponent is in NN */
        pow.magic = one.magic = WORD(0);
        exp.magic = WORD(0);

        /* Initialize elements */
        ret = fp_check_initialized(a); EG(ret, err);
        ret = fp_init(&pow, a->ctx); EG(ret, err);
        ret = fp_init(&one, a->ctx); EG(ret, err);
        ret = nn_init(&exp, 0); EG(ret, err);

        /* Initialize our variables from the Fp context of the
         * input a.
         */
        ret = fp_init(&pow, a->ctx); EG(ret, err);
        ret = fp_init(&one, a->ctx); EG(ret, err);
        ret = nn_init(&exp, 0); EG(ret, err);

        /* one = 1 in Fp */
        ret = fp_one(&one); EG(ret, err);

        /* Compute the exponent (p-1)/2
         * The computation is done in NN, and the division by 2
         * is performed using a right shift by one
         */
        ret = nn_dec(&exp, &(a->ctx->p)); EG(ret, err);
        ret = nn_rshift(&exp, &exp, 1); EG(ret, err);

        /* Compute a^((p-1)/2) in Fp using our fp_pow
         * API.
         */
        ret = fp_pow(&pow, a, &exp); EG(ret, err);

        ret = fp_iszero(&pow, &iszero); EG(ret, err);
        ret = fp_cmp(&pow, &one, &cmp); EG(ret, err);
        if (iszero) {
                ret = 0;
        } else if (cmp == 0) {
                ret = 1;
        } else {
                ret = -1;
        }

err:
        /* Cleaning */
        fp_uninit(&pow);
        fp_uninit(&one);
        nn_uninit(&exp);

        return ret;
}

/*
 * We implement the Tonelli-Shanks algorithm for finding
 * square roots (quadratic residues) modulo a prime number,
 * i.e. solving the equation:
 *     x^2 = n (p)
 * where p is a prime number. This can be seen as an equation
 * over the finite field Fp where a and x are elements of
 * this finite field.
 *   Source: https://en.wikipedia.org/wiki/Tonelli%E2%80%93Shanks_algorithm
 *   All   ≡   are taken to mean   (mod p)   unless stated otherwise.
 *   Input : p an odd prime, and an integer n .
 *       Step 0. Check that n is indeed a square  : (n | p) must be ≡ 1
 *       Step 1. [Factors out powers of 2 from p-1] Define q -odd- and s such as p-1 = q * 2^s
 *           - if s = 1 , i.e p ≡ 3 (mod 4) , output the two solutions r ≡ +/- n^((p+1)/4) .
 *       Step 2. Select a non-square z such as (z | p) = -1 , and set c ≡ z^q .
 *       Step 3. Set r ≡ n ^((q+1)/2) , t ≡ n^q, m = s .
 *       Step 4. Loop.
 *           - if t ≡ 1 output r, p-r .
 *           - Otherwise find, by repeated squaring, the lowest i , 0 < i < m , such as t^(2^i) ≡ 1
 *           - Let b ≡ c^(2^(m-i-1)), and set r ≡ r*b, t ≡ t*b^2 , c ≡ b^2 and m = i.
 *
 * NOTE: the algorithm is NOT constant time.
 *
 * The outputs, sqrt1 and sqrt2 ARE initialized by the function.
 * The function returns 0 on success, -1 on error (in which case values of sqrt1 and sqrt2
 * must not be considered).
 *
 * Aliasing is supported.
 * 
 */
int fp_sqrt(fp_t sqrt1, fp_t sqrt2, fp_src_t n)
{
        int ret, iszero, cmp, isodd;
        nn q, s, one_nn, two_nn, m, i, tmp_nn;
        fp z, t, b, r, c, one_fp, tmp_fp, __n;
        fp_t _n = &__n;
        q.magic = s.magic = one_nn.magic = two_nn.magic = m.magic = WORD(0);
        i.magic = tmp_nn.magic = z.magic = t.magic = b.magic = WORD(0);
        r.magic = c.magic = one_fp.magic = tmp_fp.magic = __n.magic = WORD(0);

        ret = nn_init(&q, 0); EG(ret, err);
        ret = nn_init(&s, 0); EG(ret, err);
        ret = nn_init(&tmp_nn, 0); EG(ret, err);
        ret = nn_init(&one_nn, 0); EG(ret, err);
        ret = nn_init(&two_nn, 0); EG(ret, err);
        ret = nn_init(&m, 0); EG(ret, err);
        ret = nn_init(&i, 0); EG(ret, err);
        ret = fp_init(&z, n->ctx); EG(ret, err);
        ret = fp_init(&t, n->ctx); EG(ret, err);
        ret = fp_init(&b, n->ctx); EG(ret, err);
        ret = fp_init(&r, n->ctx); EG(ret, err);
        ret = fp_init(&c, n->ctx); EG(ret, err);
        ret = fp_init(&one_fp, n->ctx); EG(ret, err);
        ret = fp_init(&tmp_fp, n->ctx); EG(ret, err);

        /* Handle input aliasing */
        ret = fp_copy(_n, n); EG(ret, err);

        /* Initialize outputs */
        ret = fp_init(sqrt1, _n->ctx); EG(ret, err);
        ret = fp_init(sqrt2, _n->ctx); EG(ret, err);

        /* one_nn = 1 in NN */
        ret = nn_one(&one_nn); EG(ret, err);
        /* two_nn = 2 in NN */
        ret = nn_set_word_value(&two_nn, WORD(2)); EG(ret, err);

        /* If our p prime of Fp is 2, then return the input as square roots */
        ret = nn_cmp(&(_n->ctx->p), &two_nn, &cmp); EG(ret, err);
        if (cmp == 0) {
                ret = fp_copy(sqrt1, _n); EG(ret, err);
                ret = fp_copy(sqrt2, _n); EG(ret, err);
                ret = 0;
                goto err;
        }

        /* Square root of 0 is 0 */
        ret = fp_iszero(_n, &iszero); EG(ret, err);
        if (iszero) {
                ret = fp_zero(sqrt1); EG(ret, err);
                ret = fp_zero(sqrt2); EG(ret, err);
                ret = 0;
                goto err;
        }
        /* Step 0. Check that n is indeed a square  : (n | p) must be ≡ 1 */
        if (legendre(_n) != 1) {
                /* a is not a square */
                ret = -1;
                goto err;
        }
        /* Step 1. [Factors out powers of 2 from p-1] Define q -odd- and s such as p-1 = q * 2^s */
        /* s = 0 */
        ret = nn_zero(&s); EG(ret, err);
        /* q = p - 1 */
        ret = nn_copy(&q, &(_n->ctx->p)); EG(ret, err);
        ret = nn_dec(&q, &q); EG(ret, err);
        while (1) {
                /* i is used as a temporary unused variable here */
                ret = nn_divrem(&tmp_nn, &i, &q, &two_nn); EG(ret, err);
                ret = nn_inc(&s, &s); EG(ret, err);
                ret = nn_copy(&q, &tmp_nn); EG(ret, err);
                /* If r is odd, we have finished our division */
                ret = nn_isodd(&q, &isodd); EG(ret, err);
                if (isodd) {
                        break;
                }
        }
        /* - if s = 1 , i.e p ≡ 3 (mod 4) , output the two solutions r ≡ +/- n^((p+1)/4) . */
        ret = nn_cmp(&s, &one_nn, &cmp); EG(ret, err);
        if (cmp == 0) {
                ret = nn_inc(&tmp_nn, &(_n->ctx->p)); EG(ret, err);
                ret = nn_rshift(&tmp_nn, &tmp_nn, 2); EG(ret, err);
                ret = fp_pow(sqrt1, _n, &tmp_nn); EG(ret, err);
                ret = fp_neg(sqrt2, sqrt1); EG(ret, err);
                ret = 0;
                goto err;
        }
        /* Step 2. Select a non-square z such as (z | p) = -1 , and set c ≡ z^q . */
        ret = fp_zero(&z); EG(ret, err);
        while (legendre(&z) != -1) {
                ret = fp_inc(&z, &z); EG(ret, err);
        }
        ret = fp_pow(&c, &z, &q); EG(ret, err);
        /* Step 3. Set r ≡ n ^((q+1)/2) , t ≡ n^q, m = s . */
        ret = nn_inc(&tmp_nn, &q); EG(ret, err);
        ret = nn_rshift(&tmp_nn, &tmp_nn, 1); EG(ret, err);
        ret = fp_pow(&r, _n, &tmp_nn); EG(ret, err);
        ret = fp_pow(&t, _n, &q); EG(ret, err);
        ret = nn_copy(&m, &s); EG(ret, err);
        ret = fp_one(&one_fp); EG(ret, err);

        /* Step 4. Loop. */
        while (1) {
                /* - if t ≡ 1 output r, p-r . */
                ret = fp_cmp(&t, &one_fp, &cmp); EG(ret, err);
                if (cmp == 0) {
                        ret = fp_copy(sqrt1, &r); EG(ret, err);
                        ret = fp_neg(sqrt2, sqrt1); EG(ret, err);
                        ret = 0;
                        goto err;
                }
                /* - Otherwise find, by repeated squaring, the lowest i , 0 < i < m , such as t^(2^i) ≡ 1 */
                ret = nn_one(&i); EG(ret, err);
                ret = fp_copy(&tmp_fp, &t); EG(ret, err);
                while (1) {
                        ret = fp_sqr(&tmp_fp, &tmp_fp); EG(ret, err);
                        ret = fp_cmp(&tmp_fp, &one_fp, &cmp); EG(ret, err);
                        if (cmp == 0) {
                                break;
                        }
                        ret = nn_inc(&i, &i); EG(ret, err);
                        ret = nn_cmp(&i, &m, &cmp); EG(ret, err);
                        if (cmp == 0) {
                                /* i has reached m, that should not happen ... */
                                ret = -2;
                                goto err;
                        }
                }
                /* - Let b ≡ c^(2^(m-i-1)), and set r ≡ r*b, t ≡ t*b^2 , c ≡ b^2 and m = i. */
                ret = nn_sub(&tmp_nn, &m, &i); EG(ret, err);
                ret = nn_dec(&tmp_nn, &tmp_nn); EG(ret, err);
                ret = fp_copy(&b, &c); EG(ret, err);
                ret = nn_iszero(&tmp_nn, &iszero); EG(ret, err);
                while (!iszero) {
                        ret = fp_sqr(&b, &b); EG(ret, err);
                        ret = nn_dec(&tmp_nn, &tmp_nn); EG(ret, err);
                        ret = nn_iszero(&tmp_nn, &iszero); EG(ret, err);
                }
                /* r ≡ r*b */
                ret = fp_mul(&r, &r, &b); EG(ret, err);
                /* c ≡ b^2 */
                ret = fp_sqr(&c, &b); EG(ret, err);
                /* t ≡ t*b^2 */
                ret = fp_mul(&t, &t, &c); EG(ret, err);
                /* m = i */
                ret = nn_copy(&m, &i); EG(ret, err);
        }

 err:
        /* Uninitialize local variables */
        nn_uninit(&q);
        nn_uninit(&s);
        nn_uninit(&tmp_nn);
        nn_uninit(&one_nn);
        nn_uninit(&two_nn);
        nn_uninit(&m);
        nn_uninit(&i);
        fp_uninit(&z);
        fp_uninit(&t);
        fp_uninit(&b);
        fp_uninit(&r);
        fp_uninit(&c);
        fp_uninit(&one_fp);
        fp_uninit(&tmp_fp);
        fp_uninit(&__n);

        PTR_NULLIFY(_n);

        return ret;
}