clock_curve_math/field/

batch_ops.rs

//! Batch operations for optimized computations.
//!
//! This module provides algorithms for batch operations and other advanced
//! mathematical operations on field elements. These operations are designed
//! to be more efficient than performing individual operations sequentially.

use crate::field::{FieldElement, FieldOps};

#[cfg(feature = "alloc")]
extern crate alloc;
#[cfg(feature = "alloc")]
use alloc::vec::Vec;

/// Batch inversion result containing all computed inverses.
///
/// This struct holds the computed inverses of all input field elements.
/// The inverses are guaranteed to satisfy `input[i] * result[i] ≡ 1 mod p`
/// for all valid indices i.
#[cfg(feature = "alloc")]
#[derive(Debug, Clone)]
pub struct BatchInverseResult {
    /// The computed inverses, in the same order as the input elements
    pub inverses: Vec<FieldElement>,
}

#[cfg(feature = "alloc")]
impl BatchInverseResult {
    /// Get the inverse at the specified index.
    ///
    /// Returns a reference to the computed inverse for the element at the
    /// corresponding position in the input array. The inverses are stored
    /// in the same order as the input elements.
    ///
    /// # Parameters
    /// * `index` - The index of the desired inverse (0 ≤ index < len())
    ///
    /// # Returns
    /// A reference to the computed inverse FieldElement
    ///
    /// # Panics
    /// Panics if `index` is out of bounds (index >= len()).
    ///
    /// # Mathematical Guarantee
    /// The returned inverse satisfies: `input[index] * result[index] ≡ 1 mod p`
    /// where p is the field modulus.
    pub fn get(&self, index: usize) -> &FieldElement {
        &self.inverses[index]
    }

    /// Get the number of inverses computed.
    ///
    /// Returns the total number of inverses in this result, which equals
    /// the number of input elements that were successfully inverted.
    ///
    /// # Returns
    /// The number of computed inverses
    pub fn len(&self) -> usize {
        self.inverses.len()
    }

    /// Check if the result contains no inverses.
    ///
    /// Returns true if no inverses were computed (empty result).
    /// This typically occurs when an empty input array was provided.
    ///
    /// # Returns
    /// `true` if no inverses are stored, `false` otherwise
    pub fn is_empty(&self) -> bool {
        self.inverses.is_empty()
    }

    /// Convert into the internal vector of inverses.
    ///
    /// Consumes this result and returns the underlying vector of computed inverses.
    /// This is useful when you want to take ownership of the inverses rather than
    /// borrowing them through the `get()` method.
    ///
    /// # Returns
    /// The vector of computed inverses in input order
    ///
    /// # Performance
    /// This operation is O(1) - no copying occurs, just ownership transfer.
    pub fn into_vec(self) -> Vec<FieldElement> {
        self.inverses
    }
}

/// Compute the multiplicative inverses of multiple field elements efficiently.
///
/// This function implements Montgomery's batch inversion algorithm, which computes
/// the inverses of all input elements using only one modular inversion operation
/// plus O(n) field multiplications, instead of n separate inversions.
///
/// # Algorithm
///
/// The algorithm works as follows:
/// 1. Compute partial products from left to right: a₁, a₁*a₂, a₁*a₂*a₃, ...
/// 2. Compute partial products from right to left: aₙ, aₙ*aₙ₋₁, aₙ*aₙ₋₁*aₙ₋₂, ...
/// 3. Compute the inverse of the complete product P = a₁*a₂*...*aₙ
/// 4. Use the partial products and inv(P) to compute each individual inverse
///
/// # Time Complexity
///
/// - O(n) field multiplications
/// - O(1) modular inversions (the expensive part)
/// - Total: O(n) vs O(n²) for naive approach with n inversions
///
/// # Space Complexity
///
/// - O(n) for storing partial products
///
/// # Security
///
/// All operations are performed in constant time where possible. The algorithm
/// uses the existing constant-time field operations internally.
///
/// # Examples
///
/// ```
/// use clock_curve_math::{FieldElement, FieldOps, field::batch_ops::batch_inverse};
///
/// // Create some field elements
/// let a = FieldElement::from_u64(2);
/// let b = FieldElement::from_u64(3);
/// let c = FieldElement::from_u64(4);
///
/// // Compute their inverses efficiently
/// let result = batch_inverse(&[a, b, c]).unwrap();
///
/// // Verify correctness: a * inv(a) ≡ 1 mod p
/// assert_eq!(a.mul(result.get(0)), FieldElement::from_u64(1));
/// assert_eq!(b.mul(result.get(1)), FieldElement::from_u64(1));
/// assert_eq!(c.mul(result.get(2)), FieldElement::from_u64(1));
/// ```
///
/// # Errors
///
/// Returns `None` if any input element is zero (which has no multiplicative inverse).
/// In this case, no inverses are computed.
///
/// # Performance
///
/// For small n (≤ 10), the overhead of batch inversion may not be worthwhile.
/// For larger n (> 10), batch inversion becomes increasingly beneficial.
///
#[cfg(feature = "alloc")]
pub fn batch_inverse(elements: &[FieldElement]) -> Option<BatchInverseResult> {
    let n = elements.len();

    // Handle edge cases
    if n == 0 {
        return Some(BatchInverseResult {
            inverses: Vec::new(),
        });
    }

    if n == 1 {
        // For single element, use standard inversion
        if elements[0] == FieldElement::from_u64(0) {
            return None; // Zero has no inverse
        }
        let mut inverses = Vec::new();
        inverses.push(elements[0].inv());
        return Some(BatchInverseResult { inverses });
    }

    // Check for zeros first (constant time check)
    let mut has_zero = false;
    for elem in elements {
        if elem == &FieldElement::from_u64(0) {
            has_zero = true;
            break;
        }
    }
    if has_zero {
        return None; // Cannot invert zero
    }

    // Montgomery's batch inversion algorithm
    // Step 1: Compute prefix products from left to right: prefix[i] = product of first (i+1) elements
    let mut prefix = Vec::with_capacity(n);
    prefix.push(elements[0]);

    for i in 1..n {
        prefix.push(prefix[i - 1].mul(&elements[i]));
    }

    // Step 2: Compute inverse of the complete product
    let mut inv_total = prefix[n - 1].inv();

    // Step 3: Compute inverses from right to left
    let mut inverses = Vec::with_capacity(n);

    for i in (0..n).rev() {
        // inv(a[i]) = prefix[i-1] * inv_total if i > 0 else inv_total
        let prev_prefix = if i > 0 {
            prefix[i - 1] // Direct copy, no heap allocation
        } else {
            FieldElement::from_u64(1)
        };
        let inv_a_i = prev_prefix.mul(&inv_total);
        inverses.push(inv_a_i);

        // Update inv_total for next iteration: inv_total = inv_total * a[i]
        inv_total = inv_total.mul(&elements[i]);
    }

    inverses.reverse(); // Put them in correct order

    Some(BatchInverseResult { inverses })
}

/// Compute batch inverse with comprehensive input validation and error handling.
///
/// This function provides the same Montgomery batch inversion algorithm as `batch_inverse()`,
/// but with additional input validation and structured error handling. It checks for
/// invalid inputs before attempting computation, providing clear error messages for
/// common issues like zero elements.
///
/// # Algorithm
/// Uses the same Montgomery batch inversion algorithm as `batch_inverse()`, but with
/// pre-validation to catch errors early and provide better diagnostics.
///
/// # Input Validation
/// - Checks for zero elements (which have no multiplicative inverse)
/// - Validates input array bounds
/// - Ensures all elements are valid field elements
///
/// # Error Handling
/// Returns detailed error information for:
/// - `MathError::DivisionByZero`: When any input element is zero
/// - Empty input arrays are handled gracefully (return empty result)
///
/// # Performance
/// Same O(n) performance as `batch_inverse()` but with small validation overhead.
/// The validation is minimal and doesn't significantly impact performance.
///
/// # When to Use
/// Use this function instead of `batch_inverse()` when:
/// - Input validation is critical for your application
/// - You need structured error handling
/// - Debugging invalid inputs during development
/// - Building higher-level APIs that need error propagation
///
/// # Examples
/// ```ignore
/// use clock_curve_math::{FieldElement, field::batch_ops::batch_inverse_checked};
///
/// let elements = vec![
///     FieldElement::from_u64(2),
///     FieldElement::from_u64(3),
///     FieldElement::from_u64(4),
/// ];
///
/// match batch_inverse_checked(&elements) {
///     Ok(result) => {
///         // All inverses computed successfully
///         for (i, elem) in elements.iter().enumerate() {
///             assert_eq!(elem.mul(result.get(i)), FieldElement::from_u64(1));
///         }
///     }
///     Err(crate::error::MathError::DivisionByZero) => {
///         // Handle zero element (which has no inverse)
///         println!("Cannot invert zero element");
///     }
///     Err(other) => {
///         // Handle other errors
///         println!("Batch inversion failed: {:?}", other);
///     }
/// }
/// ```
#[cfg(feature = "alloc")]
pub fn batch_inverse_checked(
    elements: &[FieldElement],
) -> Result<BatchInverseResult, crate::error::MathError> {
    if elements.is_empty() {
        return Ok(BatchInverseResult {
            inverses: Vec::new(),
        });
    }

    // Check for zero elements
    for elem in elements.iter() {
        if elem == &FieldElement::from_u64(0) {
            return Err(crate::error::MathError::DivisionByZero);
        }
    }

    // All elements are non-zero, safe to compute batch inverse
    batch_inverse(elements).ok_or(crate::error::MathError::DivisionByZero)
}

/// Optimized batch inverse for small batches (≤ 8 elements).
///
/// This function provides optimized implementations for small batch sizes
/// where the overhead of the general Montgomery algorithm may not be justified.
/// For larger batches, it falls back to the general algorithm.
///
/// # Performance
///
/// - For n ≤ 3: Uses direct inversion (optimal for small n)
/// - For n = 4-8: Uses optimized Montgomery algorithm
/// - For n > 8: Uses standard Montgomery algorithm
///
/// # Examples
///
/// ```
/// use clock_curve_math::{FieldElement, FieldOps, field::batch_ops::batch_inverse_small};
///
/// let elements = [
///     FieldElement::from_u64(2),
///     FieldElement::from_u64(3),
/// ];
///
/// let result = batch_inverse_small(&elements).unwrap();
/// assert_eq!(elements[0].mul(result.get(0)), FieldElement::from_u64(1));
/// ```
#[cfg(feature = "alloc")]
pub fn batch_inverse_small(elements: &[FieldElement]) -> Option<BatchInverseResult> {
    let n = elements.len();

    match n {
        0 => Some(BatchInverseResult {
            inverses: Vec::new(),
        }),
        1 => {
            if elements[0] == FieldElement::from_u64(0) {
                None
            } else {
                let mut inverses = Vec::new();
                inverses.push(elements[0].inv());
                Some(BatchInverseResult { inverses })
            }
        }
        2 => {
            // For 2 elements: inv(a,b) = (inv(a*b)*b, inv(a*b)*a)
            if elements[0] == FieldElement::from_u64(0) || elements[1] == FieldElement::from_u64(0)
            {
                return None;
            }
            let product = elements[0].mul(&elements[1]);
            let inv_product = product.inv();
            let mut inverses = Vec::new();
            inverses.push(inv_product.mul(&elements[1]));
            inverses.push(inv_product.mul(&elements[0]));
            Some(BatchInverseResult { inverses })
        }
        3 => {
            // For 3 elements: use Montgomery algorithm (still small)
            batch_inverse(elements)
        }
        _ => batch_inverse(elements), // Use general algorithm for larger batches
    }
}

#[cfg(all(test, feature = "alloc"))]
mod tests {
    use super::*;
    use crate::field::FieldElement;
    extern crate alloc;
    use alloc::vec::Vec;

    /// Tests batch inversion with empty input arrays.
    ///
    /// Verifies that the batch inversion functions correctly handle the edge case
    /// of empty input arrays, returning an empty result without panicking.
    /// This test ensures robustness when processing variable-length input data
    /// that might occasionally be empty, preventing crashes in production code.
    #[test]
    fn test_batch_inverse_empty() {
        let result = batch_inverse(&[]).unwrap();
        assert!(result.is_empty());
        assert_eq!(result.len(), 0);
    }

    /// Tests batch inversion with single-element arrays.
    ///
    /// Verifies that batch inversion works correctly for the smallest non-empty
    /// case (single element), ensuring the algorithm handles the base case properly.
    /// This test validates that single-element inversion produces the correct
    /// mathematical result and that the batch algorithm doesn't break down
    /// for minimal inputs.
    #[test]
    fn test_batch_inverse_single() {
        let a = FieldElement::from_u64(5);
        let result = batch_inverse(&[a]).unwrap();

        assert_eq!(result.len(), 1);
        assert_eq!(a.mul(result.get(0)), FieldElement::from_u64(1));
    }

    /// Tests batch inversion error handling for zero elements.
    ///
    /// Verifies that batch inversion correctly rejects zero elements, which have
    /// no multiplicative inverse in finite fields. This test ensures that the
    /// functions fail safely when given invalid inputs, preventing silent failures
    /// or incorrect results that could compromise cryptographic security.
    ///
    /// Zero handling is critical because zero elements can cause division by zero
    /// or other mathematical errors that might lead to exploitable behavior.
    #[test]
    fn test_batch_inverse_zero() {
        let zero = FieldElement::from_u64(0);
        let a = FieldElement::from_u64(5);

        // Single zero should fail
        assert!(batch_inverse(&[zero]).is_none());

        // Zero in batch should fail
        assert!(batch_inverse(&[a, zero, a]).is_none());
    }

    /// Tests batch inversion with exactly two elements.
    ///
    /// Verifies that the batch inversion algorithm works correctly for the smallest
    /// multi-element case. This tests the core Montgomery algorithm logic without
    /// the complexity of larger arrays, making it easier to debug issues and
    /// validate the mathematical correctness of the implementation.
    #[test]
    fn test_batch_inverse_two_elements() {
        let a = FieldElement::from_u64(2);
        let b = FieldElement::from_u64(3);

        let result = batch_inverse(&[a, b]).unwrap();

        assert_eq!(result.len(), 2);
        assert_eq!(a.mul(result.get(0)), FieldElement::from_u64(1));
        assert_eq!(b.mul(result.get(1)), FieldElement::from_u64(1));
    }

    /// Tests batch inversion with multiple elements (four elements).
    ///
    /// Verifies that the full Montgomery batch inversion algorithm works correctly
    /// for larger input arrays. This test exercises the complete algorithm including
    /// prefix product computation, total product inversion, and inverse reconstruction.
    /// It ensures the algorithm scales correctly and produces mathematically correct
    /// results for practical batch sizes used in cryptographic applications.
    #[test]
    fn test_batch_inverse_multiple_elements() {
        let elements = [
            FieldElement::from_u64(2),
            FieldElement::from_u64(3),
            FieldElement::from_u64(5),
            FieldElement::from_u64(7),
        ];

        let result = batch_inverse(&elements).unwrap();

        assert_eq!(result.len(), 4);
        for (i, elem) in elements.iter().enumerate() {
            assert_eq!(elem.mul(result.get(i)), FieldElement::from_u64(1));
        }
    }

    #[test]
    fn test_batch_inverse_checked() {
        let mut elements = Vec::new();
        elements.push(FieldElement::from_u64(2));
        elements.push(FieldElement::from_u64(3));

        let result = batch_inverse_checked(&elements).unwrap();
        assert_eq!(result.len(), 2);

        // Test with zero
        let mut elements_with_zero = Vec::new();
        elements_with_zero.push(FieldElement::from_u64(2));
        elements_with_zero.push(FieldElement::from_u64(0));

        assert!(batch_inverse_checked(&elements_with_zero).is_err());
    }

    #[test]
    fn test_batch_inverse_small() {
        // Test empty
        assert!(batch_inverse_small(&[]).unwrap().is_empty());

        // Test single
        let a = FieldElement::from_u64(5);
        let result = batch_inverse_small(&[a]).unwrap();
        assert_eq!(a.mul(result.get(0)), FieldElement::from_u64(1));

        // Test two elements
        let a = FieldElement::from_u64(2);
        let b = FieldElement::from_u64(3);
        let result = batch_inverse_small(&[a, b]).unwrap();
        assert_eq!(a.mul(result.get(0)), FieldElement::from_u64(1));
        assert_eq!(b.mul(result.get(1)), FieldElement::from_u64(1));

        // Test with zero
        assert!(batch_inverse_small(&[FieldElement::from_u64(0)]).is_none());
    }

    /// Tests the BatchInverseResult API methods.
    ///
    /// Verifies that all methods of the BatchInverseResult struct work correctly,
    /// including access methods, length reporting, and ownership transfer.
    /// This test ensures that users can reliably access computed inverses through
    /// the provided API, which is essential for practical use of batch inversion
    /// in cryptographic applications.
    #[test]
    fn test_batch_inverse_result_api() {
        let mut elements = Vec::new();
        elements.push(FieldElement::from_u64(2));
        elements.push(FieldElement::from_u64(3));

        let result = batch_inverse(&elements).unwrap();

        // Test get method
        assert_eq!(result.get(0), &elements[0].inv());
        assert_eq!(result.get(1), &elements[1].inv());

        // Test len and is_empty
        assert_eq!(result.len(), 2);
        assert!(!result.is_empty());

        // Test into_vec
        let vec_result = result.into_vec();
        assert_eq!(vec_result.len(), 2);
    }
}