//! Benchmarks for basic arithmetic operations.
//!
//! This module contains performance benchmarks for fundamental BigInt
//! arithmetic operations including addition, multiplication, and squaring.
//! Benchmarks cover both normal cases and edge cases like carry propagation.

use clock_bigint::{BigIntCore, U256, add_magnitude, mul, sqr, sub};
use criterion::{Criterion, criterion_group, criterion_main};

fn bench_u256_add(c: &mut Criterion) {
    let mut group = c.benchmark_group("u256_arithmetic");

    // Benchmark 256-bit addition
    group.bench_function("u256_add", |bencher| {
        let a = U256::from_u64(0x123456789ABCDEF0);
        let b = U256::from_u64(0xFEDCBA9876543210);

        bencher.iter(|| {
            // We need to use the dynamic BigInt for add_magnitude since U256 doesn't have the trait
            let a_big = clock_bigint::BigInt::from_limbs(a.limbs(), 4).unwrap();
            let b_big = clock_bigint::BigInt::from_limbs(b.limbs(), 4).unwrap();
            let mut result_big = clock_bigint::BigInt::new(4);

            std::hint::black_box(add_magnitude(&a_big, &b_big, &mut result_big)).unwrap();
        });
    });

    // Benchmark 256-bit multiplication
    group.bench_function("u256_mul", |bencher| {
        let a = U256::from_u64(0x123456789ABCDEF0);
        let b = U256::from_u64(0xFEDCBA9876543210);

        bencher.iter(|| {
            let a_big = clock_bigint::BigInt::from_limbs(a.limbs(), 4).unwrap();
            let b_big = clock_bigint::BigInt::from_limbs(b.limbs(), 4).unwrap();

            std::hint::black_box(mul(&a_big, &b_big)).unwrap();
        });
    });

    // Benchmark 256-bit squaring (special case of multiplication)
    group.bench_function("u256_sqr", |bencher| {
        let a = U256::from_u64(0x123456789ABCDEF0);

        bencher.iter(|| {
            let a_big = clock_bigint::BigInt::from_limbs(a.limbs(), 4).unwrap();

            std::hint::black_box(sqr(&a_big)).unwrap();
        });
    });

    // Benchmark 256-bit subtraction
    group.bench_function("u256_sub", |bencher| {
        let a = U256::from_u64(0xFEDCBA9876543210);
        let b = U256::from_u64(0x123456789ABCDEF0);

        bencher.iter(|| {
            let a_big = clock_bigint::BigInt::from_limbs(a.limbs(), 4).unwrap();
            let b_big = clock_bigint::BigInt::from_limbs(b.limbs(), 4).unwrap();

            std::hint::black_box(clock_bigint::sub(&a_big, &b_big)).unwrap();
        });
    });

    group.finish();
}

fn bench_edge_cases(c: &mut Criterion) {
    let mut group = c.benchmark_group("edge_cases");

    // Benchmark addition with carry propagation
    group.bench_function("u256_add_carry_chain", |bencher| {
        let a = U256::from_limbs(&[u64::MAX, u64::MAX, u64::MAX, 0]).unwrap();
        let b = U256::from_limbs(&[1, 0, 0, 0]).unwrap();
        let mut result = clock_bigint::BigInt::new(5); // Allow for carry

        bencher.iter(|| {
            let a_big = clock_bigint::BigInt::from_limbs(a.limbs(), 5).unwrap();
            let b_big = clock_bigint::BigInt::from_limbs(b.limbs(), 5).unwrap();

            std::hint::black_box(add_magnitude(&a_big, &b_big, &mut result)).unwrap();
        });
    });

    // Benchmark multiplication with large operands
    group.bench_function("u256_mul_large", |bencher| {
        let a = U256::from_limbs(&[u64::MAX, u64::MAX, u64::MAX, u64::MAX]).unwrap();
        let b = U256::from_limbs(&[u64::MAX, u64::MAX, u64::MAX, u64::MAX]).unwrap();

        bencher.iter(|| {
            let a_big = clock_bigint::BigInt::from_limbs(a.limbs(), 4).unwrap();
            let b_big = clock_bigint::BigInt::from_limbs(b.limbs(), 4).unwrap();

            std::hint::black_box(mul(&a_big, &b_big)).unwrap();
        });
    });

    // Benchmark subtraction with borrow propagation
    group.bench_function("u256_sub_borrow_chain", |bencher| {
        let a = U256::from_limbs(&[0, 0, 0, 1]).unwrap(); // Small number
        let b = U256::from_limbs(&[1, 0, 0, 0]).unwrap(); // Will cause borrow across limbs

        bencher.iter(|| {
            let a_big = clock_bigint::BigInt::from_limbs(a.limbs(), 4).unwrap();
            let b_big = clock_bigint::BigInt::from_limbs(b.limbs(), 4).unwrap();

            std::hint::black_box(sub(&a_big, &b_big)).unwrap();
        });
    });

    group.finish();
}

criterion_group!(benches, bench_u256_add, bench_edge_cases);
criterion_main!(benches);
