use super::*;
const PRIME: [u64; 4] = [
0xffff_ffff_ffff_ffff,
0x0000_0000_ffff_ffff,
0x0000_0000_0000_0000,
0xffff_ffff_0000_0001,
];
fn lt<const N: usize>(a: &[u64; N], b: &[u64; N]) -> bool {
let mut borrow: u128 = 0;
for i in 0..N {
let t = (a[i] as u128)
.wrapping_sub(b[i] as u128)
.wrapping_sub(borrow);
borrow = (t >> 64) & 1;
}
borrow == 1
}
pub(super) fn eq<const N: usize>(a: &[u64; N], b: &[u64; N]) -> bool {
let mut same = true;
for i in 0..N {
same &= a[i] == b[i];
}
same
}
fn add_n<const N: usize>(a: [u64; N], b: [u64; N]) -> [u64; N] {
let mut r = [0u64; N];
let mut carry: u128 = 0;
for i in 0..N {
let t = a[i] as u128 + b[i] as u128 + carry;
r[i] = t as u64;
carry = t >> 64;
}
assert_eq!(carry, 0);
r
}
fn widen<const N: usize>(a: [u64; 4]) -> [u64; N] {
let mut r = [0u64; N];
for i in 0..4 {
r[i] = a[i];
}
r
}
fn any_residue() -> [u64; 4] {
let a: [u64; 4] = kani::any();
kani::assume(lt(&a, &PRIME));
a
}
#[kani::proof]
fn add_is_a_plus_b_mod_p() {
let (a, b) = (any_residue(), any_residue());
let out = add(a, b);
let sum: [u64; 5] = add_n(widen(a), widen(b));
assert!(lt(&out, &PRIME));
assert!(sum == widen(out) || sum == add_n(widen(out), widen(PRIME)));
}
#[kani::proof]
fn sub_is_a_minus_b_mod_p() {
let (a, b) = (any_residue(), any_residue());
let out = sub(a, b);
assert!(lt(&out, &PRIME));
let back: [u64; 5] = add_n(widen(out), widen(b));
assert!(back == widen(a) || back == add_n(widen(a), widen(PRIME)));
}
fn add_w<const N: usize>(a: [u64; N], b: [u64; N]) -> [u64; N] {
let mut r = [0u64; N];
let mut carry: u128 = 0;
for i in 0..N {
let t = a[i] as u128 + b[i] as u128 + carry;
r[i] = t as u64;
carry = t >> 64;
}
r
}
fn sub_w<const N: usize>(a: [u64; N], b: [u64; N]) -> [u64; N] {
let mut r = [0u64; N];
let mut borrow: u128 = 0;
for i in 0..N {
let t = (a[i] as u128)
.wrapping_sub(b[i] as u128)
.wrapping_sub(borrow);
r[i] = t as u64;
borrow = (t >> 64) & 1;
}
r
}
fn shl<const N: usize>(u: u64, s: usize) -> [u64; N] {
let mut r = [0u64; N];
let w = s / 64;
if s % 64 == 0 {
r[w] = u;
} else {
r[w] = u << 32;
r[w + 1] = u >> 32;
}
r
}
fn plus_u_p<const N: usize>(x: [u64; N], u: u64, k: usize) -> [u64; N] {
let b = 64 * k;
let pos = add_w(
add_w(add_w(x, shl(u, b + 256)), shl(u, b + 192)),
shl(u, b + 96),
);
sub_w(sub_w(pos, shl(u, b + 224)), shl(u, b))
}
fn at_word<const N: usize>(x: [u64; 4], k: usize) -> [u64; N] {
let mut r = [0u64; N];
for i in 0..4 {
r[k + i] = x[i];
}
r
}
fn montgomery_spec(t: [u64; 9]) -> [u64; 9] {
let mut s = t;
for i in 0..4 {
s = plus_u_p(s, s[i], i);
}
s
}
pub(super) fn mont_step_post(w: [u64; 4], out: [u64; 4]) -> bool {
eq(&plus_u_p::<6>(widen(w), w[0], 0), &at_word(out, 1))
}
pub(super) fn montgomery_low_post(lo: [u64; 4], w: [u64; 4]) -> bool {
eq(&montgomery_spec(widen(lo)), &at_word(w, 4)) & !lt(&PRIME, &w)
}
pub(super) fn reduce_pre(t: [u64; 8]) -> bool {
lt(&[t[4], t[5], t[6], t[7]], &PRIME)
}
pub(super) fn reduce_post(t: [u64; 8], out: [u64; 4]) -> bool {
let lo = [t[0], t[1], t[2], t[3]];
let hi = [t[4], t[5], t[6], t[7]];
let s = add_w(montgomery_spec(widen(lo)), at_word(hi, 4));
let mine: [u64; 9] = at_word(out, 4);
lt(&out, &PRIME) & (eq(&s, &mine) | eq(&s, &add_w(mine, at_word(PRIME, 4))))
}
pub(super) fn finish_pre(hi: [u64; 4], w: [u64; 4]) -> bool {
lt(&hi, &PRIME) & !lt(&PRIME, &w)
}
pub(super) fn finish_post(hi: [u64; 4], w: [u64; 4], out: [u64; 4]) -> bool {
let s: [u64; 5] = add_w(widen(hi), widen(w));
lt(&out, &PRIME) & (eq(&s, &widen(out)) | eq(&s, &add_w(widen(out), widen(PRIME))))
}
#[kani::proof]
fn the_prime_has_the_shape_the_specification_uses() {
assert_eq!(plus_u_p::<5>([0; 5], 1, 0), widen(PRIME));
assert_eq!(PRIME[0], u64::MAX);
}
#[kani::proof_for_contract(mont_step)]
#[kani::solver(bitwuzla)]
fn mont_step_meets_its_contract() {
mont_step(kani::any());
}
#[kani::proof_for_contract(montgomery_low)]
#[kani::stub_verified(mont_step)]
#[kani::solver(bitwuzla)]
fn montgomery_low_meets_its_contract() {
montgomery_low(kani::any());
}
#[kani::proof_for_contract(finish)]
fn finish_meets_its_contract() {
finish(kani::any(), kani::any());
}
#[kani::proof_for_contract(reduce)]
#[kani::stub_verified(montgomery_low)]
#[kani::stub_verified(finish)]
#[kani::solver(bitwuzla)]
fn reduce_meets_its_contract() {
reduce(kani::any());
}
fn word_product(x: u64, y: u64) -> u128 {
(x as u128).wrapping_mul(y as u128)
}
pub(super) fn mac_post(a: u64, b: u64, c: u64, carry: u64, r: (u64, u64)) -> bool {
let want = (a as u128)
.wrapping_add(word_product(b, c))
.wrapping_add(carry as u128);
(r.0 as u128 | ((r.1 as u128) << 64)) == want
}
#[kani::proof_for_contract(mac)]
#[kani::solver(bitwuzla)]
fn mac_meets_its_contract() {
mac(kani::any(), kani::any(), kani::any(), kani::any());
}
const MAX_WORD_PRODUCT: u128 = 0xffff_ffff_ffff_fffe_0000_0000_0000_0001;
static mut LOG_X: [u128; 16] = [0; 16];
static mut LOG_Y: [u128; 16] = [0; 16];
static mut LOG_P: [u128; 16] = [0; 16];
static mut LOG_LEN: usize = 0;
fn product_abstract(x: u128, y: u128) -> u128 {
let p: u128 = kani::any();
kani::assume(p <= MAX_WORD_PRODUCT);
unsafe {
let n = LOG_LEN;
if n < 16 {
LOG_X[n] = x;
LOG_Y[n] = y;
LOG_P[n] = p;
}
LOG_LEN = n + 1;
}
p
}
fn ghost_log() -> (usize, [u128; 16], [u128; 16], [u128; 16]) {
unsafe { (LOG_LEN, LOG_X, LOG_Y, LOG_P) }
}
fn term_at<const N: usize>(x: u128, k: usize) -> [u64; N] {
let mut r = [0u64; N];
r[k] = x as u64;
r[k + 1] = (x >> 64) as u64;
r
}
fn low8(x: [u64; 9]) -> [u64; 8] {
let mut r = [0u64; 8];
for i in 0..8 {
r[i] = x[i];
}
r
}
fn row_from(t: [u128; 4]) -> [u64; 5] {
let mut r = [0u64; 5];
for j in 0..4 {
r = add_w(r, term_at(t[j], j));
}
r
}
fn product_from(t: [[u128; 4]; 4]) -> [u64; 9] {
let mut s = [0u64; 9];
for i in 0..4 {
let row = row_from(t[i]);
let mut term = [0u64; 9];
for k in 0..5 {
term[i + k] = row[k];
}
s = add_w(s, term);
}
s
}
fn square_from(d: [u128; 4], x: [[u128; 4]; 4]) -> [u64; 9] {
let mut s = [0u64; 9];
for i in 0..4 {
s = add_w(s, term_at(d[i], 2 * i));
for j in i + 1..4 {
let t = term_at(x[i][j], i + j);
s = add_w(add_w(s, t), t);
}
}
s
}
fn word_products(a: [u64; 4], b: [u64; 4]) -> [[u128; 4]; 4] {
let mut t = [[0u128; 4]; 4];
for i in 0..4 {
for j in 0..4 {
t[i][j] = word_product(a[i], b[j]);
}
}
t
}
fn product_spec(a: [u64; 4], b: [u64; 4]) -> [u64; 8] {
low8(product_from(word_products(a, b)))
}
fn square_spec(a: [u64; 4]) -> [u64; 8] {
let t = word_products(a, a);
let mut d = [0u128; 4];
for i in 0..4 {
d[i] = t[i][i];
}
low8(square_from(d, t))
}
const MUL_SCHEDULE: [(usize, usize); 16] = [
(0, 0),
(0, 1),
(0, 2),
(0, 3),
(1, 0),
(1, 1),
(1, 2),
(1, 3),
(2, 0),
(2, 1),
(2, 2),
(2, 3),
(3, 0),
(3, 1),
(3, 2),
(3, 3),
];
const SQUARE_SCHEDULE: [(usize, usize); 10] = [
(0, 1),
(0, 2),
(0, 3),
(1, 2),
(1, 3),
(2, 3),
(0, 0),
(1, 1),
(2, 2),
(3, 3),
];
#[kani::proof]
#[kani::stub(u128::wrapping_mul, product_abstract)]
fn the_schoolbook_sums_its_products_at_their_weights() {
let (a, b): ([u64; 4], [u64; 4]) = (kani::any(), kani::any());
let t = product(a, b);
let (n, xs, ys, ps) = ghost_log();
assert!(n == 16);
let mut terms = [[0u128; 4]; 4];
for k in 0..16 {
let (i, j) = MUL_SCHEDULE[k];
assert!(xs[k] == a[i] as u128 && ys[k] == b[j] as u128);
terms[i][j] = ps[k];
}
let want = product_from(terms);
assert!(want[8] == 0);
assert!(eq(&t, &low8(want)));
}
#[kani::proof]
#[kani::stub(u128::wrapping_mul, product_abstract)]
#[kani::solver(bitwuzla)]
fn the_square_schoolbook_sums_its_products_at_their_weights() {
let a: [u64; 4] = kani::any();
let t = square_product(a);
let (n, xs, ys, ps) = ghost_log();
assert!(n == 10);
let mut d = [0u128; 4];
let mut x = [[0u128; 4]; 4];
for k in 0..10 {
let (i, j) = SQUARE_SCHEDULE[k];
assert!(xs[k] == a[i] as u128 && ys[k] == a[j] as u128);
if i == j {
d[i] = ps[k];
} else {
x[i][j] = ps[k];
}
}
let want = square_from(d, x);
assert!(want[8] == 0);
assert!(eq(&t, &low8(want)));
}
#[kani::proof]
#[kani::solver(bitwuzla)]
fn the_product_of_residues_is_in_reduce_s_domain() {
let (a, b) = (any_residue(), any_residue());
assert!(reduce_pre(product_spec(a, b)));
assert!(reduce_pre(square_spec(a)));
}