add documentation

add arguments and return value docstrings

Author: jacksonwalters

src/lib.rs

@@ -2,16 +2,26 @@ use reikna::totient::totient;
 use reikna::factor::quick_factorize;
 use std::collections::HashMap;
 
-// Modular arithmetic functions using i64
+/// Modular arithmetic functions using i64
 fn mod_add(a: i64, b: i64, p: i64) -> i64 {
     (a + b) % p
 }
 
+/// Modular multiplication
 fn mod_mul(a: i64, b: i64, p: i64) -> i64 {
     (a * b) % p
 }
 
-pub fn mod_exp(mut base: i64, mut exp: i64, p: i64) -> i64 {
+/// Modular exponentiation
+/// # Arguments
+///
+/// * `base` - Base of the exponentiation.
+/// * `exp` - Exponent.
+/// * `p` - Prime modulus for the operations.
+///
+/// # Returns
+/// The result of the exponentiation modulo `p`.
+fn mod_exp(mut base: i64, mut exp: i64, p: i64) -> i64 {
     let mut result = 1;
     base %= p;
     while exp > 0 {
@@ -24,6 +34,7 @@ pub fn mod_exp(mut base: i64, mut exp: i64, p: i64) -> i64 {
     result
 }
 
+/// Extended Euclidean algorithm
 fn extended_gcd(a: i64, b: i64) -> (i64, i64, i64) {
     if b == 0 {
         (a, 1, 0)  // gcd, x, y
@@ -33,15 +44,23 @@ fn extended_gcd(a: i64, b: i64) -> (i64, i64, i64) {
     }
 }
 
-pub fn mod_inv(a: i64, modulus: i64) -> i64 {
+/// Compute the modular inverse of a modulo modulus
+fn mod_inv(a: i64, modulus: i64) -> i64 {
     let (gcd, x, _) = extended_gcd(a, modulus);
     if gcd != 1 {
         panic!("{} and {} are not coprime, no inverse exists", a, modulus);
     }
     (x % modulus + modulus) % modulus  // Ensure a positive result
 }
 
-// Compute n-th root of unity (omega) for p not necessarily prime
+/// Compute n-th root of unity (omega) for p not necessarily prime
+/// # Arguments
+///
+/// * `modulus` - Modulus. n must divide each prime power factor.
+/// * `n` - Order of the root of unity.
+/// 
+/// # Returns
+/// The n-th root of unity modulo `modulus`.
 pub fn omega(modulus: i64, n: usize) -> i64 {
     let factors = factorize(modulus as i64);
     if factors.len() == 1 {
@@ -56,7 +75,15 @@ pub fn omega(modulus: i64, n: usize) -> i64 {
     }
 }
 
-// Forward transform using NTT, output bit-reversed 
+/// Forward transform using NTT, output bit-reversed
+/// # Arguments
+///
+/// * `a` - Input vector.
+/// * `omega` - Primitive root of unity modulo `p`.
+/// * `n` - Length of the input vector and the result.
+/// * `p` - Prime modulus for the operations.
+///
+/// # Returns
 pub fn ntt(a: &[i64], omega: i64, n: usize, p: i64) -> Vec<i64> {
     let mut result = a.to_vec();
     let mut step = n/2;
@@ -77,7 +104,16 @@ pub fn ntt(a: &[i64], omega: i64, n: usize, p: i64) -> Vec<i64> {
 	result
 }
 
-// Inverse transform using INTT, input bit-reversed 
+/// Inverse transform using INTT, input bit-reversed
+/// # Arguments
+/// 
+/// * `a` - Input vector (bit-reversed).
+/// * `omega` - Primitive root of unity modulo `p`.
+/// * `n` - Length of the input vector and the result.
+/// * `p` - Prime modulus for the operations.
+///
+/// # Returns
+/// A vector representing the inverse NTT of the input vector.
 pub fn intt(a: &[i64], omega: i64, n: usize, p: i64) -> Vec<i64> {
     let omega_inv = mod_inv(omega, p);
     let n_inv = mod_inv(n as i64, p);
@@ -103,7 +139,16 @@ pub fn intt(a: &[i64], omega: i64, n: usize, p: i64) -> Vec<i64> {
         .collect()
 }
 
-// Naive polynomial multiplication
+/// Naive polynomial multiplication
+/// # Arguments
+///
+/// * `a` - First polynomial (as a vector of coefficients).
+/// * `b` - Second polynomial (as a vector of coefficients).
+/// * `n` - Length of the polynomials and the result.
+/// * `p` - Prime modulus for the operations.
+///
+/// # Returns
+/// A vector representing the polynomial product modulo `p`.
 pub fn polymul(a: &Vec<i64>, b: &Vec<i64>, n: i64, p: i64) -> Vec<i64> {
     let mut result = vec![0; n as usize];
     for i in 0..a.len() {
@@ -145,7 +190,14 @@ pub fn polymul_ntt(a: &[i64], b: &[i64], n: usize, p: i64, omega: i64) -> Vec<i6
     c
 }
 
-/// Compute the prime factorization of `n` (with multiplicities).
+/// Compute the prime factorization of `n` (with multiplicities)
+/// Uses reikna::quick_factorize internally
+/// # Arguments
+/// 
+/// * `n` - Number to factorize.
+/// 
+/// # Returns
+/// A HashMap with the prime factors of `n` as keys and their multiplicities as values.
 fn factorize(n: i64) -> HashMap<i64, u32> {
     let mut factors = HashMap::new();
     for factor in quick_factorize(n as u64) {
@@ -167,6 +219,12 @@ pub fn primitive_root(p: i64, e: u32) -> i64 {
 }
 
 /// Finds a primitive root modulo a prime p
+/// # Arguments
+///
+/// * `p` - Prime modulus.
+///
+/// # Returns
+/// A primitive root modulo `p`.
 fn primitive_root_mod_p(p: i64) -> i64 {
     let phi = p - 1;
     let factors = factorize(phi); // Reusing factorize to get both prime factors and multiplicities
@@ -179,7 +237,16 @@ fn primitive_root_mod_p(p: i64) -> i64 {
     0 // Should never happen
 }
 
-// the Chinese remainder theorem for two moduli
+/// the Chinese remainder theorem for two moduli
+/// # Arguments
+///
+/// * `a1` - First residue.
+/// * `n1` - First modulus.
+/// * `a2` - Second residue.
+/// * `n2` - Second modulus.
+///
+/// # Returns
+/// The solution to the system of congruences x = a1 (mod n1) and x = a2 (mod n2).
 pub fn crt(a1: i64, n1: i64, a2: i64, n2: i64) -> i64 {
     let n = n1 * n2;
     let m1 = mod_inv(n1, n2); // Inverse of n1 mod n2
@@ -188,10 +255,17 @@ pub fn crt(a1: i64, n1: i64, a2: i64, n2: i64) -> i64 {
     if x < 0 { x + n } else { x }
 }
 
-// computes an n^th root of unity modulo a composite modulus
-// note we require that an n^th root of unity exists for each multiplicative group modulo p^e
-// use the CRT isomorphism to pull back each n^th root of unity to the composite modulus
-// for the NTT, we require than a 2n^th root of unity exists
+/// computes an n^th root of unity modulo a composite modulus
+/// note we require that an n^th root of unity exists for each multiplicative group modulo p^e
+/// use the CRT isomorphism to pull back each n^th root of unity to the composite modulus
+/// for the NTT, we require than a 2n^th root of unity exists
+/// # Arguments
+///
+/// * `modulus` - Modulus. n must divide each prime power factor.
+/// * `n` - Order of the root of unity.
+///
+/// # Returns
+/// The n-th root of unity modulo `modulus`.
 pub fn root_of_unity(modulus: i64, n: i64) -> i64 {
     let factors = factorize(modulus);
     let mut result = 1;
@@ -202,7 +276,15 @@ pub fn root_of_unity(modulus: i64, n: i64) -> i64 {
 	result
 }
 
-//ensure the root of unity satisfies sum_{j=0}^{n-1} omega^{jk} = 0 for 1 \le k < n
+/// ensure the root of unity satisfies sum_{j=0}^{n-1} omega^{jk} = 0 for 1 \le k < n
+/// # Arguments
+///
+/// * `omega` - n-th root of unity.
+/// * `n` - Order of the root of unity.
+/// * `modulus` - Modulus.
+///
+/// # Returns
+/// True if the root of unity satisfies the condition.
 pub fn verify_root_of_unity(omega: i64, n: i64, modulus: i64) -> bool {
     assert!(mod_exp(omega, n, modulus as i64) == 1, "omega is not an n-th root of unity");
     assert!(mod_exp(omega, n/2, modulus as i64) == modulus-1, "omgea^(n/2) != -1 (mod modulus)");