package gnarl

// Montgomery field arithmetic over Z_P.
// 4x uint64 little-endian limbs in Montgomery form.
// P is 216-bit (27 bytes). R = 2^256.

import "math/bits"

type fe [4]uint64

// These are initialized in init().
var pLimbs fe
var feOne fe
var feR2 fe
var feZero fe
var tsOddQ fe
var tsQm1h fe
var thirdPNorm fe
var halfPNorm fe
var pMinus2 fe

const pPrime uint64 = 0xf07d39ef3ea058a3
const tsS = 2
const fieldLen = 27

func initFieldConstants() {
	pLimbs = fe{0x31f9791e0f9ee4f5, 0x121e79ccd682cc99, 0xb6d81bb9b02e5e5d, 0x00000000009563a6}
	feOne = fe{0x761ae156c8e8e77c, 0x023540a7764c229f, 0x315e327188dada36, 0x00000000001012bc}
	feR2 = fe{0xf56927577991be00, 0xf6eb880a87ed3681, 0x18d8b12f949488c0, 0x0000000000684a62}
	tsOddQ = fe{0x4c7e5e4783e7b93d, 0x44879e7335a0b326, 0xadb606ee6c0b9797, 0x00000000002558e9}
	tsQm1h = fe{0x263f2f23c1f3dc9e, 0xa243cf399ad05993, 0xd6db03773605cbcb, 0x000000000012ac74}
	thirdPNorm = fe{0xbb53285f5a8a4c51, 0xb0b4d3444780eedd, 0x3cf2b3e8900f74c9, 0x000000000031cbe2}
	halfPNorm = fe{0x98fcbc8f07cf727a, 0x890f3ce66b41664c, 0x5b6c0ddcd8172f2e, 0x00000000004ab1d3}
	pMinus2 = fe{0x31f9791e0f9ee4f3, 0x121e79ccd682cc99, 0xb6d81bb9b02e5e5d, 0x00000000009563a6}
}

func feSet(dst, src *fe) {
	*dst = *src
}

func feIsZero(a *fe) (result int32) {
	z := a[0] | a[1] | a[2] | a[3]
	z = (z | (0 - z)) >> 63
	return 1 - int32(z)
}

func feEqual(a, b *fe) (result int32) {
	z := (a[0] ^ b[0]) | (a[1] ^ b[1]) | (a[2] ^ b[2]) | (a[3] ^ b[3])
	z = (z | (0 - z)) >> 63
	return 1 - int32(z)
}

func feReduce(a *fe) {
	var d0, d1, d2, d3, borrow uint64
	d0, borrow = bits.Sub64(a[0], pLimbs[0], 0)
	d1, borrow = bits.Sub64(a[1], pLimbs[1], borrow)
	d2, borrow = bits.Sub64(a[2], pLimbs[2], borrow)
	d3, borrow = bits.Sub64(a[3], pLimbs[3], borrow)

	mask := uint64(0) - borrow
	a[0] = (a[0] & mask) | (d0 & ^mask)
	a[1] = (a[1] & mask) | (d1 & ^mask)
	a[2] = (a[2] & mask) | (d2 & ^mask)
	a[3] = (a[3] & mask) | (d3 & ^mask)
}

func feAdd(r, a, b *fe) {
	var carry, borrow, underflow uint64
	var d0, d1, d2, d3 uint64
	r[0], carry = bits.Add64(a[0], b[0], 0)
	r[1], carry = bits.Add64(a[1], b[1], carry)
	r[2], carry = bits.Add64(a[2], b[2], carry)
	r[3], carry = bits.Add64(a[3], b[3], carry)

	d0, borrow = bits.Sub64(r[0], pLimbs[0], 0)
	d1, borrow = bits.Sub64(r[1], pLimbs[1], borrow)
	d2, borrow = bits.Sub64(r[2], pLimbs[2], borrow)
	d3, borrow = bits.Sub64(r[3], pLimbs[3], borrow)

	underflow = borrow &^ carry
	mask := uint64(0) - underflow
	r[0] = (r[0] & mask) | (d0 & ^mask)
	r[1] = (r[1] & mask) | (d1 & ^mask)
	r[2] = (r[2] & mask) | (d2 & ^mask)
	r[3] = (r[3] & mask) | (d3 & ^mask)
}

func feSub(r, a, b *fe) {
	var borrow, carry uint64
	r[0], borrow = bits.Sub64(a[0], b[0], 0)
	r[1], borrow = bits.Sub64(a[1], b[1], borrow)
	r[2], borrow = bits.Sub64(a[2], b[2], borrow)
	r[3], borrow = bits.Sub64(a[3], b[3], borrow)

	mask := uint64(0) - borrow
	r[0], carry = bits.Add64(r[0], pLimbs[0]&mask, 0)
	r[1], carry = bits.Add64(r[1], pLimbs[1]&mask, carry)
	r[2], carry = bits.Add64(r[2], pLimbs[2]&mask, carry)
	r[3], _ = bits.Add64(r[3], pLimbs[3]&mask, carry)
}

func feNeg(r, a *fe) {
	z := a[0] | a[1] | a[2] | a[3]
	nonzero := z | (0 - z)
	mask := 0 - (nonzero >> 63)

	var borrow uint64
	r[0], borrow = bits.Sub64(pLimbs[0], a[0], 0)
	r[1], borrow = bits.Sub64(pLimbs[1], a[1], borrow)
	r[2], borrow = bits.Sub64(pLimbs[2], a[2], borrow)
	r[3], _ = bits.Sub64(pLimbs[3], a[3], borrow)
	r[0] &= mask
	r[1] &= mask
	r[2] &= mask
	r[3] &= mask
}

// montMul computes r = a * b * R^{-1} mod P using fully-unrolled CIOS.
func montMul(r, a, b *fe) {
	var t0, t1, t2, t3, t4 uint64
	var hi, lo, c0, c1, m, carry uint64

	// i = 0
	hi, lo = bits.Mul64(a[0], b[0])
	t0, c0 = bits.Add64(lo, t0, 0)
	carry, _ = bits.Add64(hi, 0, c0)
	hi, lo = bits.Mul64(a[0], b[1])
	lo, c0 = bits.Add64(lo, t1, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t1, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(a[0], b[2])
	lo, c0 = bits.Add64(lo, t2, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t2, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(a[0], b[3])
	lo, c0 = bits.Add64(lo, t3, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t3, c0 = bits.Add64(lo, carry, 0)
	t4, _ = bits.Add64(hi, c1, c0)
	m = t0 * pPrime
	hi, lo = bits.Mul64(m, pLimbs[0])
	_, c0 = bits.Add64(t0, lo, 0)
	carry, _ = bits.Add64(hi, 0, c0)
	hi, lo = bits.Mul64(m, pLimbs[1])
	lo, c0 = bits.Add64(lo, t1, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t0, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(m, pLimbs[2])
	lo, c0 = bits.Add64(lo, t2, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t1, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(m, pLimbs[3])
	lo, c0 = bits.Add64(lo, t3, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t2, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	t3, c0 = bits.Add64(t4, carry, 0)
	t4 = c0

	// i = 1
	hi, lo = bits.Mul64(a[1], b[0])
	t0, c0 = bits.Add64(lo, t0, 0)
	carry, _ = bits.Add64(hi, 0, c0)
	hi, lo = bits.Mul64(a[1], b[1])
	lo, c0 = bits.Add64(lo, t1, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t1, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(a[1], b[2])
	lo, c0 = bits.Add64(lo, t2, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t2, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(a[1], b[3])
	lo, c0 = bits.Add64(lo, t3, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t3, c0 = bits.Add64(lo, carry, 0)
	hi, _ = bits.Add64(hi, c1, c0)
	t4 += hi
	m = t0 * pPrime
	hi, lo = bits.Mul64(m, pLimbs[0])
	_, c0 = bits.Add64(t0, lo, 0)
	carry, _ = bits.Add64(hi, 0, c0)
	hi, lo = bits.Mul64(m, pLimbs[1])
	lo, c0 = bits.Add64(lo, t1, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t0, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(m, pLimbs[2])
	lo, c0 = bits.Add64(lo, t2, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t1, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(m, pLimbs[3])
	lo, c0 = bits.Add64(lo, t3, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t2, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	t3, c0 = bits.Add64(t4, carry, 0)
	t4 = c0

	// i = 2
	hi, lo = bits.Mul64(a[2], b[0])
	t0, c0 = bits.Add64(lo, t0, 0)
	carry, _ = bits.Add64(hi, 0, c0)
	hi, lo = bits.Mul64(a[2], b[1])
	lo, c0 = bits.Add64(lo, t1, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t1, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(a[2], b[2])
	lo, c0 = bits.Add64(lo, t2, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t2, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(a[2], b[3])
	lo, c0 = bits.Add64(lo, t3, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t3, c0 = bits.Add64(lo, carry, 0)
	hi, _ = bits.Add64(hi, c1, c0)
	t4 += hi
	m = t0 * pPrime
	hi, lo = bits.Mul64(m, pLimbs[0])
	_, c0 = bits.Add64(t0, lo, 0)
	carry, _ = bits.Add64(hi, 0, c0)
	hi, lo = bits.Mul64(m, pLimbs[1])
	lo, c0 = bits.Add64(lo, t1, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t0, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(m, pLimbs[2])
	lo, c0 = bits.Add64(lo, t2, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t1, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(m, pLimbs[3])
	lo, c0 = bits.Add64(lo, t3, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t2, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	t3, c0 = bits.Add64(t4, carry, 0)
	t4 = c0

	// i = 3
	hi, lo = bits.Mul64(a[3], b[0])
	t0, c0 = bits.Add64(lo, t0, 0)
	carry, _ = bits.Add64(hi, 0, c0)
	hi, lo = bits.Mul64(a[3], b[1])
	lo, c0 = bits.Add64(lo, t1, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t1, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(a[3], b[2])
	lo, c0 = bits.Add64(lo, t2, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t2, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(a[3], b[3])
	lo, c0 = bits.Add64(lo, t3, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t3, c0 = bits.Add64(lo, carry, 0)
	hi, _ = bits.Add64(hi, c1, c0)
	t4 += hi
	m = t0 * pPrime
	hi, lo = bits.Mul64(m, pLimbs[0])
	_, c0 = bits.Add64(t0, lo, 0)
	carry, _ = bits.Add64(hi, 0, c0)
	hi, lo = bits.Mul64(m, pLimbs[1])
	lo, c0 = bits.Add64(lo, t1, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t0, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(m, pLimbs[2])
	lo, c0 = bits.Add64(lo, t2, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t1, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	hi, lo = bits.Mul64(m, pLimbs[3])
	lo, c0 = bits.Add64(lo, t3, 0)
	hi, c1 = bits.Add64(hi, 0, c0)
	t2, c0 = bits.Add64(lo, carry, 0)
	carry, _ = bits.Add64(hi, c1, c0)
	t3, c0 = bits.Add64(t4, carry, 0)
	t4 = c0

	r[0] = t0
	r[1] = t1
	r[2] = t2
	r[3] = t3
	feReduce(r)
}

func montSquare(r, a *fe) {
	montMul(r, a, a)
}

func feToMont(r, a *fe) {
	montMul(r, a, &feR2)
}

func feFromMont(r, a *fe) {
	one := fe{1, 0, 0, 0}
	montMul(r, a, &one)
}

func feFromSmall(r *fe, v int64) {
	if v >= 0 {
		*r = fe{uint64(v), 0, 0, 0}
	} else {
		*r = pLimbs
		var borrow uint64
		r[0], borrow = bits.Sub64(r[0], uint64(-v), 0)
		r[1], borrow = bits.Sub64(r[1], 0, borrow)
		r[2], borrow = bits.Sub64(r[2], 0, borrow)
		r[3], _ = bits.Sub64(r[3], 0, borrow)
	}
	feToMont(r, r)
}

func feFromBytes27(r *fe, b []byte) (ok bool) {
	if int32(len(b)) < 27 {
		*r = feZero
		return false
	}
	r[3] = uint64(b[0])<<16 | uint64(b[1])<<8 | uint64(b[2])
	r[2] = uint64(b[3])<<56 | uint64(b[4])<<48 | uint64(b[5])<<40 | uint64(b[6])<<32 |
		uint64(b[7])<<24 | uint64(b[8])<<16 | uint64(b[9])<<8 | uint64(b[10])
	r[1] = uint64(b[11])<<56 | uint64(b[12])<<48 | uint64(b[13])<<40 | uint64(b[14])<<32 |
		uint64(b[15])<<24 | uint64(b[16])<<16 | uint64(b[17])<<8 | uint64(b[18])
	r[0] = uint64(b[19])<<56 | uint64(b[20])<<48 | uint64(b[21])<<40 | uint64(b[22])<<32 |
		uint64(b[23])<<24 | uint64(b[24])<<16 | uint64(b[25])<<8 | uint64(b[26])

	var d0, d1, borrow uint64
	d0, borrow = bits.Sub64(r[0], pLimbs[0], 0)
	d1, borrow = bits.Sub64(r[1], pLimbs[1], borrow)
	_, borrow = bits.Sub64(r[2], pLimbs[2], borrow)
	_, borrow = bits.Sub64(r[3], pLimbs[3], borrow)
	_ = d0
	_ = d1
	if borrow == 0 {
		*r = feZero
		return false
	}

	feToMont(r, r)
	return true
}

func feToBytes27(b []byte, a *fe) {
	var norm fe
	feFromMont(&norm, a)
	b[0] = byte(norm[3] >> 16)
	b[1] = byte(norm[3] >> 8)
	b[2] = byte(norm[3])
	b[3] = byte(norm[2] >> 56)
	b[4] = byte(norm[2] >> 48)
	b[5] = byte(norm[2] >> 40)
	b[6] = byte(norm[2] >> 32)
	b[7] = byte(norm[2] >> 24)
	b[8] = byte(norm[2] >> 16)
	b[9] = byte(norm[2] >> 8)
	b[10] = byte(norm[2])
	b[11] = byte(norm[1] >> 56)
	b[12] = byte(norm[1] >> 48)
	b[13] = byte(norm[1] >> 40)
	b[14] = byte(norm[1] >> 32)
	b[15] = byte(norm[1] >> 24)
	b[16] = byte(norm[1] >> 16)
	b[17] = byte(norm[1] >> 8)
	b[18] = byte(norm[1])
	b[19] = byte(norm[0] >> 56)
	b[20] = byte(norm[0] >> 48)
	b[21] = byte(norm[0] >> 40)
	b[22] = byte(norm[0] >> 32)
	b[23] = byte(norm[0] >> 24)
	b[24] = byte(norm[0] >> 16)
	b[25] = byte(norm[0] >> 8)
	b[26] = byte(norm[0])
}

func fePow(r, base, exp *fe) {
	var result fe
	feSet(&result, &feOne)
	var b fe
	feSet(&b, base)

	for i := int32(0); i < 4; i++ {
		word := exp[i]
		for bit := int32(0); bit < 64; bit++ {
			if word&1 == 1 {
				montMul(&result, &result, &b)
			}
			montSquare(&b, &b)
			word >>= 1
		}
	}
	feSet(r, &result)
}

func feInv(r, a *fe) {
	fePow(r, a, &pMinus2)
}

func feSqrt(r, a *fe) (ok bool) {
	if feIsZero(a) == 1 {
		*r = feZero
		return true
	}

	var feFive fe
	feFive[0] = 5
	feToMont(&feFive, &feFive)

	var t fe
	fePow(&t, a, &tsOddQ)

	var rr, tmp fe
	fePow(&tmp, a, &tsQm1h)
	montMul(&rr, &tmp, a)

	var c fe
	fePow(&c, &feFive, &tsOddQ)

	m := int32(tsS)

	for {
		if feEqual(&t, &feOne) == 1 {
			feSet(r, &rr)
			return true
		}

		var bv fe
		feSet(&bv, &t)
		iv := int32(0)
		for j := int32(1); j < m; j++ {
			montSquare(&bv, &bv)
			if feEqual(&bv, &feOne) == 1 {
				iv = j
				break
			}
		}
		if iv == 0 {
			return false
		}

		feSet(&bv, &c)
		for j := int32(0); j < m-iv-1; j++ {
			montSquare(&bv, &bv)
		}

		montMul(&rr, &rr, &bv)
		montSquare(&c, &bv)
		montMul(&t, &t, &c)
		m = iv
	}
}

func feIsGtThirdP(a *fe) (result int32) {
	for i := int32(3); i >= 0; i-- {
		if a[i] > thirdPNorm[i] {
			return 1
		}
		if a[i] < thirdPNorm[i] {
			return 0
		}
	}
	return 0
}

func feIsGtHalfP(a *fe) (result int32) {
	for i := int32(3); i >= 0; i-- {
		if a[i] > halfPNorm[i] {
			return 1
		}
		if a[i] < halfPNorm[i] {
			return 0
		}
	}
	return 0
}