module poly1305

import math.bits
import math.unsigned

const uint192_zero = Uint192{}

// Uint192 is a structure representing 192 bits of unsigned integer.
// It serves as a custom allocator for the poly1305 implementation to use.
// However, it can be used to complement and expand the standard `math.unsigned`
// module with a more complete and extensive range of handling of unsigned integers.
struct Uint192 {
mut:
	lo u64
	mi u64
	hi u64
}

// Here, we define several required functionalities of the custom allocator.

// add_checked returns u+v with carry
fn (u Uint192) add_checked(v Uint192, c u64) (Uint192, u64) {
	lo, c0 := bits.add_64(u.lo, v.lo, c)
	mi, c1 := bits.add_64(u.mi, v.mi, c0)
	hi, c2 := bits.add_64(u.hi, v.hi, c1)
	x := Uint192{lo, mi, hi}
	return x, c2
}

// add_128_checked returns u+v with carry
fn (u Uint192) add_128_checked(v unsigned.Uint128, c u64) (Uint192, u64) {
	lo, c0 := bits.add_64(u.lo, v.lo, c)
	mi, c1 := bits.add_64(u.mi, v.hi, c0)
	hi, c2 := bits.add_64(u.hi, 0, c1)
	x := Uint192{lo, mi, hi}
	return x, c2
}

fn (u Uint192) add_64_checked(v u64, c u64) (Uint192, u64) {
	lo, c0 := bits.add_64(u.lo, v, c)
	mi, c1 := bits.add_64(u.mi, 0, c0)
	hi, c2 := bits.add_64(u.hi, 0, c1)
	x := Uint192{lo, mi, hi}
	return x, c2
}

fn (u Uint192) sub_checked(v Uint192) (Uint192, u64) {
	lo, b0 := bits.sub_64(u.lo, v.lo, 0)
	mi, b1 := bits.sub_64(u.mi, v.mi, b0)
	hi, b2 := bits.sub_64(u.hi, v.hi, b1)
	return Uint192{lo, mi, hi}, b2
}

// mul_64_checked returns `u*v` even if the result size is over > 192 bits.
// It returns a `(Uin192, u64)` pair where the former stores the low 192 bits and
// the rest of high bits stored in the `u64` part. You can check the value of the `u64` part
// for `c != 0`, it's mean, the product of `u*v` is overflowing 192 bits.
fn (u Uint192) mul_64_checked(v u64) (Uint192, u64) {
	// see mul_128_checked for the logic
	m0 := u128_from_64_mul(u.lo, v)
	m1 := u128_from_64_mul(u.mi, v)
	m2 := u128_from_64_mul(u.hi, v)
	// propagates carry
	t0, c0 := bits.add_64(m0.lo, 0, 0)
	t1, c1 := bits.add_64(m0.hi, m1.lo, c0)
	t2, c2 := bits.add_64(m1.hi, m2.lo, c1)
	t3, c3 := bits.add_64(m2.hi, 0, c2)
	// something bad happen if the last c3 is non null
	if c3 != 0 {
		panic('Uint192: unexpected overflow')
	}
	x := Uint192{
		lo: t0
		mi: t1
		hi: t2
	}
	return x, t3
}

// mul_128_checked is a generic product of 192 bits u by 128 bits v.
// Its returns u*v even if the result is over > 192 bits, and stores the remaining
// high bits of the result into the Uint128 structure.
fn (u Uint192) mul_128_checked(v unsigned.Uint128) (Uint192, unsigned.Uint128) {
	//                  u.hi         u.mi        u.lo
	//                               v.hi        v.lo
	// ---------------------------------------------------------- x
	//                  uhi*vlo      umi*vlo     ulo*vlo
	//      uhi*vhi     umi*vhi      ulo*vhi
	// ==========================================================
	//      m3          m2           m1          m0                // Uint128
	//
	// ---------------------------------------------------------- +
	//      m3.hi       m2.hi       m1.hi        m0.hi
	//                  m3.lo       m2.lo        m1.lo     m0.lo
	// ==========================================================
	//      t4          t3          t2           t1        t0
	//
	ulovlo := u128_from_64_mul(u.lo, v.lo)
	umivlo := u128_from_64_mul(u.mi, v.lo)
	uhivlo := u128_from_64_mul(u.hi, v.lo)

	ulovhi := u128_from_64_mul(u.lo, v.hi)
	umivhi := u128_from_64_mul(u.mi, v.hi)
	uhivhi := u128_from_64_mul(u.hi, v.hi)

	m0 := ulovlo
	m1, c1 := unsigned.add_128(umivlo, ulovhi, 0)
	m2, c2 := unsigned.add_128(uhivlo, umivhi, c1)
	m3, c3 := unsigned.add_128(uhivhi, unsigned.uint128_zero, c2)
	if c3 != 0 {
		panic('Uint192: unexpected overflow')
	}
	// Note about the multiplication results in Poly1305 context.
	// In the properly clamped 128 bits of v, (called "r" in the poly1305 context) and
	// safely reduced form of high part of the 192 bits accumulator u (u.hi), where only
	// maximum of four low bits of u.hi is set, and we can assume (and confirm with tests)
	// if the high bit part of the product of uhi*vhi and uhi*vlo is not set,
	// ie, x = (uhi*vhi).hi == 0 and y = (uhi*vlo).hi == 0.
	// and the 128 bits addition of `m2 = uhi*vlo + umi*vhi`, would also not overflow 128 bits,
	// thats also mean, the last carry is null for the reason and m3 = (uhi*vhi).hi is also null
	//
	t0 := m0.lo
	t1, c4 := bits.add_64(m0.hi, m1.lo, 0)
	t2, c5 := bits.add_64(m1.hi, m2.lo, c4)
	t3, c6 := bits.add_64(m2.hi, m3.lo, c5)
	t4, c7 := bits.add_64(m3.hi, 0, c6)
	if c7 != 0 {
		panic('Uint192: unexpected overflow')
	}
	// based on previous notes, for poly1305 context, it tells us if the product
	// doesn't have a fitfh limb (t4), ie t4 == null, and we can safely ignore it.
	//
	x := Uint192{
		lo: t0
		mi: t1
		hi: t2
	}
	hb := unsigned.uint128_new(t3, t4)
	return x, hb
}

// u128_from_64_mul creates new Uint128 from 64x64 bit product of x*y
fn u128_from_64_mul(x u64, y u64) unsigned.Uint128 {
	hi, lo := bits.mul_64(x, y)
	return unsigned.uint128_new(lo, hi)
}

// select_64 returns x if v == 1 and y if v == 0, in constant time.
fn select_64(v u64, x u64, y u64) u64 {
	return ~(v - 1) & x | (v - 1) & y
}

fn shift_right_by2(mut a unsigned.Uint128) unsigned.Uint128 {
	a.lo = a.lo >> 2 | (a.hi & 3) << 62
	a.hi = a.hi >> 2
	return a
}