clock_bigint/

timeboxed.rs

//! Time-sliced (pausable) operations for long-running computations.
//!
//! Enables asynchronous VM execution by allowing operations to be paused
//! and resumed, preventing blocking during long computations.

use crate::ct;
use crate::error::{BigIntError, Result};
use crate::limbs::{Limb, limb_mac, limb_mul_wide};
use crate::montgomery::{MontgomeryContext, from_mont, mont_mul, to_mont};
use crate::types::BigIntCore;

#[cfg(feature = "alloc")]
use crate::types::BigInt;
#[cfg(feature = "alloc")]
use alloc::{vec, vec::Vec};

/// Partial result from a time-sliced operation.
///
/// Contains the current state and a flag indicating whether the operation
/// is complete or needs to be resumed.
#[derive(Clone, Debug)]
pub enum PartialResult<T> {
    /// Operation completed successfully.
    Complete(T),
    /// Operation needs to be resumed with the given continuation state.
    Partial(ContinuationState),
}

/// Continuation state for resuming a time-sliced operation.
#[derive(Clone, Debug)]
pub struct ContinuationState {
    /// Operation type identifier.
    pub op_type: OpType,
    /// Current intermediate values.
    pub state: OperationState,
    /// Number of cycles consumed so far.
    pub cycles_consumed: u64,
    /// Maximum cycles allowed.
    pub max_cycles: u64,
}

/// Type of operation being performed.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum OpType {
    /// Multiplication operation.
    Mul,
    /// Modular exponentiation operation.
    ModExp,
}

/// Internal state for different operation types.
#[derive(Clone, Debug)]
pub enum OperationState {
    /// State for multiplication.
    MulState {
        /// First operand.
        a: BigInt,
        /// Second operand.
        b: BigInt,
        /// Current result buffer.
        result: Vec<Limb>,
        /// Current iteration index.
        i: usize,
        /// Current j index (for nested loops).
        j: usize,
    },
    /// State for modular exponentiation.
    ModExpState {
        /// Modulus for Montgomery arithmetic.
        modulus: BigInt,
        /// Exponent.
        exponent: BigInt,
        /// R0 value (Montgomery ladder).
        r0: BigInt,
        /// R1 value (Montgomery ladder).
        r1: BigInt,
        /// Current bit position.
        bit_pos: usize,
        /// Total bits to process.
        total_bits: usize,
    },
}

/// Start a time-sliced multiplication operation.
///
/// Returns a PartialResult that can be resumed if the operation
/// exceeds the maximum cycle count.
#[cfg(feature = "alloc")]
pub fn mul_timeboxed(a: &BigInt, b: &BigInt, max_cycles: u64) -> Result<PartialResult<BigInt>> {
    // Handle zero cases
    if a.is_zero() || b.is_zero() {
        let max_limbs = a.max_limbs().max(b.max_limbs());
        return Ok(PartialResult::Complete(BigInt::new(max_limbs)));
    }

    let n = a.limb_count().max(b.limb_count());
    let result_limbs = 2 * n;

    // Check for overflow
    if result_limbs > a.max_limbs().max(b.max_limbs()) {
        return Err(BigIntError::Overflow);
    }

    let max_limbs = a.max_limbs().max(b.max_limbs()).max(result_limbs);
    let mut result = BigInt::new(max_limbs);
    result.ensure_capacity(result_limbs)?;

    let a_limbs = a.limbs();
    let b_limbs = b.limbs();

    // Normalize to same length for multiplication
    let mut a_normalized = vec![0u64; n];
    let mut b_normalized = vec![0u64; n];
    a_normalized[..a_limbs.len()].copy_from_slice(a_limbs);
    b_normalized[..b_limbs.len()].copy_from_slice(b_limbs);

    let result_limbs_mut = result.limbs_mut();

    // Initialize result to zero
    for i in 0..(2 * n) {
        result_limbs_mut[i] = 0;
    }

    // Start timeboxed Comba multiplication
    let mut result_vec = result_limbs_mut.to_vec();
    mul_timeboxed_comba(
        &a_normalized,
        &b_normalized,
        &mut result_vec,
        0,
        0,
        0,
        max_cycles,
    )
}

fn mul_timeboxed_comba(
    a: &[Limb],
    b: &[Limb],
    result: &mut Vec<Limb>,
    start_i: usize,
    start_j: usize,
    cycles_consumed: u64,
    max_cycles: u64,
) -> Result<PartialResult<BigInt>> {
    let n = a.len();
    debug_assert_eq!(b.len(), n);
    debug_assert!(result.len() >= 2 * n);

    let mut cycles = cycles_consumed;
    let mut i = start_i;
    let mut j = start_j;

    // Comba multiplication algorithm with cycle tracking
    while i < n {
        let mut carry = 0u64;
        let _j_start = if i == start_i { start_j } else { 0 };

        while j < n {
            // Each limb multiplication costs some cycles
            cycles += 10; // Estimate cycles for limb operations
            if cycles >= max_cycles {
                // Create continuation state
                let max_limbs = result.len();
                let mut result_bigint = BigInt::new(max_limbs);
                result_bigint.limbs_mut().copy_from_slice(result);

                let state = OperationState::MulState {
                    a: BigInt::from_limbs(a, max_limbs)?,
                    b: BigInt::from_limbs(b, max_limbs)?,
                    result: result.to_vec(),
                    i,
                    j,
                };

                return Ok(PartialResult::Partial(ContinuationState {
                    op_type: OpType::Mul,
                    state,
                    cycles_consumed: cycles,
                    max_cycles,
                }));
            }

            let wide = limb_mul_wide(a[i], b[j]);
            let (lo, hi) = (wide.lo, wide.hi);
            let (sum_lo, carry1) = limb_mac(result[i + j], lo, carry);
            result[i + j] = sum_lo;
            carry = carry1 + hi;
            j += 1;
        }

        // Add final carry
        cycles += 2; // Estimate cycles for carry addition
        if i + n < result.len() {
            let (sum, _) = limb_mac(result[i + n], 0, carry);
            result[i + n] = sum;
        }

        i += 1;
        j = 0;
    }

    // Operation complete - create final BigInt
    let max_limbs = result.len();
    let mut result_bigint = BigInt::new(max_limbs);
    result_bigint.ensure_capacity(result.len())?;
    result_bigint.limbs_mut().copy_from_slice(result);
    result_bigint.canonicalize()?;

    Ok(PartialResult::Complete(result_bigint))
}

/// Start a time-sliced modular exponentiation operation.
///
/// Returns a PartialResult that can be resumed if the operation
/// exceeds the maximum cycle count.
#[cfg(feature = "alloc")]
pub fn modexp_timeboxed(
    ctx: &MontgomeryContext,
    base: &BigInt,
    exponent: &BigInt,
    max_cycles: u64,
) -> Result<PartialResult<BigInt>> {
    // Handle special cases
    if exponent.is_zero() {
        // base^0 = 1 mod m
        let one = BigInt::from_u64(1, ctx.modulus().max_limbs());
        return Ok(PartialResult::Complete(one));
    }

    if base.is_zero() {
        // 0^exponent = 0 mod m (if exponent > 0)
        return Ok(PartialResult::Complete(BigInt::from_u64(
            0,
            ctx.modulus().max_limbs(),
        )));
    }

    // Convert base to Montgomery form
    let x_mont = to_mont(ctx, base)?;

    // Convert 1 to Montgomery form (R̃ = to_mont(1))
    let one = BigInt::from_u64(1, ctx.modulus().max_limbs());
    let r_mont = to_mont(ctx, &one)?;

    // Initialize Montgomery ladder: R0 = R̃, R1 = x̃
    let r0 = r_mont;
    let r1 = x_mont;

    // Find the most significant bit
    let exp_limbs = exponent.limbs();
    let mut msb_found = false;
    let mut bit_pos = 0;

    for i in (0..exp_limbs.len()).rev() {
        if exp_limbs[i] != 0 {
            bit_pos = (i * 64) + (63 - exp_limbs[i].leading_zeros() as usize);
            msb_found = true;
            break;
        }
    }

    if !msb_found {
        // Exponent is zero (should have been caught above, but handle it)
        let one = BigInt::from_u64(1, ctx.modulus().max_limbs());
        return Ok(PartialResult::Complete(one));
    }

    // Start timeboxed Montgomery ladder
    modexp_timeboxed_ladder(ctx, exponent, r0, r1, bit_pos, 0, max_cycles)
}

fn modexp_timeboxed_ladder(
    ctx: &MontgomeryContext,
    exponent: &BigInt,
    mut r0: BigInt,
    mut r1: BigInt,
    total_bits: usize,
    start_bit: usize,
    max_cycles: u64,
) -> Result<PartialResult<BigInt>> {
    let mut cycles_consumed = 0u64;
    let exp_limbs = exponent.limbs();

    // Process bits from MSB to LSB
    for bit_pos in (0..=total_bits).rev().skip(start_bit) {
        // Check cycle limit before each step
        cycles_consumed += 5; // Estimate cycles for bit processing setup
        if cycles_consumed >= max_cycles {
            let state = OperationState::ModExpState {
                modulus: ctx.modulus().clone(),
                exponent: exponent.clone(),
                r0: r0.clone(),
                r1: r1.clone(),
                bit_pos,
                total_bits,
            };

            return Ok(PartialResult::Partial(ContinuationState {
                op_type: OpType::ModExp,
                state,
                cycles_consumed,
                max_cycles,
            }));
        }

        let limb_idx = bit_pos / 64;
        let bit_idx = bit_pos % 64;
        let bit = if limb_idx < exp_limbs.len() {
            (exp_limbs[limb_idx] >> bit_idx) & 1
        } else {
            0
        };

        // Constant-time conditional swap
        cycles_consumed += 20; // Estimate cycles for swap operations
        ct::ct_swap(bit, &mut r0.limbs_mut()[0], &mut r1.limbs_mut()[0]);
        swap_bigint(bit, &mut r0, &mut r1);

        // Montgomery ladder step
        cycles_consumed += 50; // Estimate cycles for Montgomery operations
        r1 = mont_mul(ctx, &r0, &r1)?;
        r0 = mont_mul(ctx, &r0, &r0)?;

        // Swap again
        cycles_consumed += 20; // Estimate cycles for swap operations
        swap_bigint(bit, &mut r0, &mut r1);
    }

    // Convert result from Montgomery form
    let result = from_mont(ctx, &r0)?;

    Ok(PartialResult::Complete(result))
}

/// Constant-time swap of two BigInt values.
fn swap_bigint(choice: u64, a: &mut BigInt, b: &mut BigInt) {
    let a_limbs = a.limbs_mut();
    let b_limbs = b.limbs_mut();
    let len = a_limbs.len().min(b_limbs.len());

    for i in 0..len {
        ct::ct_swap(choice, &mut a_limbs[i], &mut b_limbs[i]);
    }

    // Also swap signs
    let a_sign = a.sign();
    let b_sign = b.sign();
    let new_a_sign = ct::ct_select(choice, b_sign as u64, a_sign as u64) != 0;
    let new_b_sign = ct::ct_select(choice, a_sign as u64, b_sign as u64) != 0;
    a.set_sign(new_a_sign);
    b.set_sign(new_b_sign);
}

/// Resume a time-sliced operation from a continuation state.
///
/// Continues execution until completion or until max_cycles is reached.
#[cfg(feature = "alloc")]
pub fn resume(state: ContinuationState) -> Result<PartialResult<BigInt>> {
    match state.op_type {
        OpType::Mul => {
            // Resume multiplication
            if let OperationState::MulState {
                a,
                b,
                mut result,
                i,
                j,
            } = state.state
            {
                let a_limbs = a.limbs();
                let b_limbs = b.limbs();

                mul_timeboxed_comba(
                    a_limbs,
                    b_limbs,
                    &mut result,
                    i,
                    j,
                    state.cycles_consumed,
                    state.max_cycles,
                )
            } else {
                Err(BigIntError::InvalidState)
            }
        }
        OpType::ModExp => {
            // Resume modular exponentiation
            if let OperationState::ModExpState {
                modulus,
                exponent,
                r0,
                r1,
                bit_pos,
                total_bits,
            } = state.state
            {
                // Reconstruct Montgomery context
                let ctx = MontgomeryContext::new(modulus)?;

                modexp_timeboxed_ladder(
                    &ctx,
                    &exponent,
                    r0,
                    r1,
                    total_bits,
                    bit_pos,
                    state.max_cycles - state.cycles_consumed,
                )
            } else {
                Err(BigIntError::InvalidState)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_mul_timeboxed() {
        let a = BigInt::from_u64(10, 10);
        let b = BigInt::from_u64(20, 10);
        let result = mul_timeboxed(&a, &b, 1000).unwrap();
        match result {
            PartialResult::Complete(r) => {
                assert_eq!(r.limbs()[0], 200);
            }
            PartialResult::Partial(_) => {
                panic!("Expected complete result");
            }
        }
    }

    #[test]
    fn test_mul_timeboxed_with_high_cycles() {
        let a = BigInt::from_u64(100, 10);
        let b = BigInt::from_u64(200, 10);

        // Use high cycle limit to ensure completion
        let result = mul_timeboxed(&a, &b, 10000).unwrap();
        match result {
            PartialResult::Complete(r) => {
                assert_eq!(r.limbs()[0], 20000);
            }
            PartialResult::Partial(_) => {
                panic!("Expected complete result with high cycle limit");
            }
        }
    }
}