Compiler Explorer

Source code

use std::array;

/// The Baby Bear prime: 2^31 - 2^27 + 1.
const P: u32 = 2013265921;
// To help read the assembly, note: 2^32 - P = 2281701375

/// Addition modulo P.
///
/// Inputs are assumed to be < P.
/// Assuming this, outputs will also be < P.
#[inline(always)]
fn add(lhs: u32, rhs: u32) -> u32 {
    let mut sum = lhs + rhs; // Never overflows as inputs are < P < 2^31.
    let (corr_sum, over) = sum.overflowing_sub(P);
     // over is false if sum >= P and true if sum < P.
    if !over { 
        sum = corr_sum;
    }
    sum
}

/// Addition modulo P in a degree 4 extension.
///
/// Identical to 4 additions in parallel.
#[unsafe(no_mangle)]
pub fn add_bb_deg_4_ext(lhs: [u32; 4], rhs: [u32; 4]) -> [u32; 4] {
    array::from_fn(|i| add(lhs[i], rhs[i]))
}

/// Compute (lhs + mid) + rhs
#[unsafe(no_mangle)]
pub fn add_bb_assoc_l(lhs: [u32; 4], mid: [u32; 4], rhs: [u32; 4]) -> [u32; 4] {
    let lhs = add_bb_deg_4_ext(lhs, mid);
    add_bb_deg_4_ext(lhs, rhs)
}

/// Compute lhs + (mid + rhs)
#[unsafe(no_mangle)]
pub fn add_bb_assoc_r(lhs: [u32; 4], mid: [u32; 4], rhs: [u32; 4]) -> [u32; 4] {
    let rhs = add_bb_deg_4_ext(mid, rhs);
    add_bb_deg_4_ext(lhs, rhs)
}

/// Check that (lhs + mid) + rhs = lhs + (mid + rhs)
#[unsafe(no_mangle)]
pub fn check_assoc(lhs: [u32; 4], mid: [u32; 4], rhs: [u32; 4]) {
    let assoc_l = add_bb_assoc_l(lhs, mid, rhs);
    let assoc_r = add_bb_assoc_r(lhs, mid, rhs);
    assert_eq!(assoc_l, assoc_r);
}

fn main() -> () {
    let lhs = [252551971, 694974649, 213757600, 1325013984];
    let mid = [506156623, 97664653, 1234719014, 1349792299];
    let rhs = [1626423134, 1338438783, 786682629, 1311208151];

// Check all elements are < P
    assert!(lhs.iter().all(|&x| x < P));
    assert!(mid.iter().all(|&x| x < P));
    assert!(rhs.iter().all(|&x| x < P));

// Let's manually check that
    // (lhs + mid) + rhs = lhs + (mid + rhs)
    let assoc_l = add_bb_assoc_l(lhs, mid, rhs);
    let assoc_r = add_bb_assoc_r(lhs, mid, rhs);

// Check that (x + y) + z = x + (y + z)
    println!("{assoc_l:?}");
    println!("{assoc_r:?}");
    assert_eq!(assoc_l, assoc_r);
    println!("The two assoc's are equal");
    // Everything up until here runs.

// Now let's check the same thing using our
    // function check_assoc
    check_assoc(lhs, mid, rhs);
    // This fails with compilation option:
    // -O -C target-cpu=znver4 -C opt-level=3
    // among others.
    println!("The function check_assoc passed");
}