/* fe25519.c -- see fe25519.h for the representation and why. */
#include "starkernel/fe25519.h"
#include <string.h>

typedef unsigned __int128 u128;

#define WBITS 51
#define MASK51 (((int64_t)1 << WBITS) - 1)

void fe25519_0(fe25519 *r) { memset(r->v, 0, sizeof(r->v)); }

void fe25519_1(fe25519 *r) { memset(r->v, 0, sizeof(r->v)); r->v[0] = 1; }

void fe25519_copy(fe25519 *r, const fe25519 *a) { *r = *a; }

/* Carry-propagate v[] in place, base 2^51, 5 limbs; the carry out of
 * limb 4 wraps back into limb 0 multiplied by 19 (2^255 = 19 mod p, and
 * 51*5 = 255 exactly, so this is the plain reduction constant -- no
 * extra scaling needed, unlike a limb count/width that doesn't divide
 * 255 evenly). Two passes: the wraparound carry can itself need to
 * ripple through limb 0 again on rare inputs. Works for negative limbs
 * too (arithmetic shift preserves sign) -- needed after fe25519_sub. */
static void fe25519_carry(fe25519 *r)
{
    for (int pass = 0; pass < 2; pass++) {
        for (int i = 0; i < 4; i++) {
            int64_t c = r->v[i] >> WBITS;
            r->v[i] -= c << WBITS;
            r->v[i+1] += c;
        }
        int64_t c = r->v[4] >> WBITS;
        r->v[4] -= c << WBITS;
        r->v[0] += c * 19;
    }
}

void fe25519_add(fe25519 *r, const fe25519 *a, const fe25519 *b)
{
    for (int i = 0; i < 5; i++) r->v[i] = a->v[i] + b->v[i];
    fe25519_carry(r);
}

void fe25519_sub(fe25519 *r, const fe25519 *a, const fe25519 *b)
{
    for (int i = 0; i < 5; i++) r->v[i] = a->v[i] - b->v[i];
    fe25519_carry(r);
}

void fe25519_neg(fe25519 *r, const fe25519 *a)
{
    fe25519 zero; fe25519_0(&zero);
    fe25519_sub(r, &zero, a);
}

/* Schoolbook multiply into a 9-element __int128 wide accumulator (each
 * product a[i]*b[j] is at most ~51+51=102 bits; a bucket sums at most 5
 * such products, at most ~105 bits -- __int128 has ~127 bits of range,
 * comfortable headroom, confirmed this kernel's toolchains compile
 * __int128 multiply/add/shift-by-constant with no libgcc calls -- see
 * fe25519.h's own doc comment). Fold positions 5-8 into 0-3 (2^255 = 19
 * mod p, and since 51*5=255 exactly, position k for k=5..8 represents
 * 2^(255+51*(k-5)) = 19 * 2^(51*(k-5)) mod p, i.e. folds directly into
 * position k-5 with the plain multiplier 19 -- no extra scale factor,
 * unlike a limb width that doesn't divide 255 evenly). */
void fe25519_mul(fe25519 *r, const fe25519 *a, const fe25519 *b)
{
    u128 wide[9] = {0};
    for (int i = 0; i < 5; i++) {
        for (int j = 0; j < 5; j++) {
            wide[i+j] += (u128)(uint64_t)a->v[i] * (uint64_t)b->v[j];
        }
    }
    for (int k = 8; k >= 5; k--) {
        wide[k-5] += 19 * wide[k];
    }
    /* wide[0..4] can each exceed 51 (even 63) bits after folding --
     * carry-propagate using __int128 arithmetic (roomy enough that no
     * intermediate step here risks overflow) before narrowing to the
     * int64_t storage representation. */
    u128 carry = 0;
    for (int i = 0; i < 5; i++) {
        u128 full = wide[i] + carry;
        r->v[i] = (int64_t)(full & (u128)MASK51);
        carry = full >> WBITS;
    }
    r->v[0] += (int64_t)(carry * 19);
    fe25519_carry(r);
}

void fe25519_sq(fe25519 *r, const fe25519 *a)
{
    fe25519_mul(r, a, a);
}

void fe25519_mul_small(fe25519 *r, const fe25519 *a, uint32_t c)
{
    u128 wide[5];
    for (int i = 0; i < 5; i++) wide[i] = (u128)(uint64_t)a->v[i] * c;
    u128 carry = 0;
    for (int i = 0; i < 5; i++) {
        u128 full = wide[i] + carry;
        r->v[i] = (int64_t)(full & (u128)MASK51);
        carry = full >> WBITS;
    }
    r->v[0] += (int64_t)(carry * 19);
    fe25519_carry(r);
}

/* Modular inverse via Fermat's little theorem: a^(p-2) mod p, using the
 * standard Ed25519 addition chain (222 squarings + 24 extra multiplies)
 * -- the same exponent decomposition every Ed25519 reference
 * implementation uses, not something invented here; verified by the
 * property-based host test (a * inverse(a) == 1) rather than trusted by
 * inspection. */
void fe25519_invert(fe25519 *r, const fe25519 *a)
{
    fe25519 z2, z9, z11, z2_5_0, z2_10_0, z2_20_0, z2_50_0, z2_100_0, t;

    fe25519_sq(&z2, a);                    /* 2 */
    fe25519_sq(&t, &z2);                   /* 4 */
    fe25519_sq(&t, &t);                    /* 8 */
    fe25519_mul(&z9, &t, a);               /* 9 */
    fe25519_mul(&z11, &z9, &z2);           /* 11 */
    fe25519_sq(&t, &z11);                  /* 22 */
    fe25519_mul(&z2_5_0, &t, &z9);         /* 2^5 - 2^0 = 31 */

    fe25519_sq(&t, &z2_5_0);
    for (int i = 1; i < 5; i++) fe25519_sq(&t, &t);
    fe25519_mul(&z2_10_0, &t, &z2_5_0);    /* 2^10 - 2^0 */

    fe25519_sq(&t, &z2_10_0);
    for (int i = 1; i < 10; i++) fe25519_sq(&t, &t);
    fe25519_mul(&z2_20_0, &t, &z2_10_0);   /* 2^20 - 2^0 */

    fe25519_sq(&t, &z2_20_0);
    for (int i = 1; i < 20; i++) fe25519_sq(&t, &t);
    fe25519_mul(&t, &t, &z2_20_0);         /* 2^40 - 2^0 */

    for (int i = 0; i < 10; i++) fe25519_sq(&t, &t);
    fe25519_mul(&z2_50_0, &t, &z2_10_0);   /* 2^50 - 2^0 */

    fe25519_sq(&t, &z2_50_0);
    for (int i = 1; i < 50; i++) fe25519_sq(&t, &t);
    fe25519_mul(&z2_100_0, &t, &z2_50_0);  /* 2^100 - 2^0 */

    fe25519_sq(&t, &z2_100_0);
    for (int i = 1; i < 100; i++) fe25519_sq(&t, &t);
    fe25519_mul(&t, &t, &z2_100_0);        /* 2^200 - 2^0 */

    for (int i = 0; i < 50; i++) fe25519_sq(&t, &t);
    fe25519_mul(&t, &t, &z2_50_0);         /* 2^250 - 2^0 */

    fe25519_sq(&t, &t);
    fe25519_sq(&t, &t);
    fe25519_sq(&t, &t);
    fe25519_sq(&t, &t);
    fe25519_sq(&t, &t);                    /* 2^255 - 2^5 */
    fe25519_mul(r, &t, &z11);              /* 2^255 - 21 = p - 2 */
}

/* Fully reduce a's limbs to a unique representative in [0, p), then pack
 * little-endian into 32 bytes. */
void fe25519_pack(uint8_t out[32], const fe25519 *a)
{
    fe25519 t = *a;
    fe25519_carry(&t);

    /* Force full reduction: subtract p once if t >= p, by subtracting
     * p's exact base-2^51 representation (limb0 = 2^51-19, limbs 1-4 =
     * 2^51-1, computed directly, not memorized) and checking whether
     * that needed a borrow out of the top -- if not, t was already >= p
     * and the subtracted value is the reduced one; if so, t was already
     * < p. Not constant-time; fine here since packing only ever happens
     * on public data (a verified signature's derived point, never a
     * secret) in this verify-only library. */
    static const int64_t P_LIMBS[5] = {
        MASK51 - 18, MASK51, MASK51, MASK51, MASK51
    };
    int64_t borrow = 0;
    int64_t rawlimb[5];
    for (int i = 0; i < 5; i++) {
        int64_t v = t.v[i] - P_LIMBS[i] - borrow;
        if (v < 0) {
            v += (int64_t)1 << WBITS;
            borrow = 1;
        } else {
            borrow = 0;
        }
        rawlimb[i] = v;
    }
    if (!borrow) {
        for (int i = 0; i < 5; i++) t.v[i] = rawlimb[i];
    }

    uint64_t bits = 0;
    int bitcnt = 0;
    int outpos = 0;
    for (int i = 0; i < 5; i++) {
        bits |= ((uint64_t)(uint64_t)t.v[i]) << bitcnt;
        bitcnt += WBITS;
        while (bitcnt >= 8) {
            out[outpos++] = (uint8_t)bits;
            bits >>= 8;
            bitcnt -= 8;
        }
    }
    if (bitcnt > 0 && outpos < 32) {
        out[outpos++] = (uint8_t)bits;
    }
    while (outpos < 32) out[outpos++] = 0;
}

void fe25519_unpack(fe25519 *r, const uint8_t in[32])
{
    uint64_t bits = 0;
    int bitcnt = 0;
    int inpos = 0;
    for (int i = 0; i < 5; i++) {
        while (bitcnt < WBITS && inpos < 32) {
            bits |= ((uint64_t)in[inpos++]) << bitcnt;
            bitcnt += 8;
        }
        r->v[i] = (int64_t)(bits & (uint64_t)MASK51);
        bits >>= WBITS;
        bitcnt -= WBITS;
    }
    /* Bit 255 of the input (the top bit of byte 31) is Ed25519's point-
     * encoding sign bit, not part of y's magnitude -- caller (point
     * decompression) masks it out before calling this. Limb 4 here only
     * ever receives up to 51 significant unpacked bits from a 255-bit
     * input by construction (4*51=204, 255-204=51), so no explicit mask
     * against a stray bit 255 is needed as long as the caller has
     * already cleared it in the input bytes. */
}

int fe25519_eq(const fe25519 *a, const fe25519 *b)
{
    uint8_t pa[32], pb[32];
    fe25519_pack(pa, a);
    fe25519_pack(pb, b);
    return memcmp(pa, pb, 32) == 0;
}

int fe25519_parity(const fe25519 *a)
{
    uint8_t p[32];
    fe25519_pack(p, a);
    return p[0] & 1;
}