add docstring test for omega

Author: jacksonwalters

src/lib.rs

@@ -21,7 +21,7 @@ fn mod_mul(a: i64, b: i64, p: i64) -> i64 {
 ///
 /// # Returns
 /// The result of the exponentiation modulo `p`.
-fn mod_exp(mut base: i64, mut exp: i64, p: i64) -> i64 {
+pub fn mod_exp(mut base: i64, mut exp: i64, p: i64) -> i64 {
     let mut result = 1;
     base %= p;
     while exp > 0 {
@@ -68,6 +68,20 @@ fn mod_inv(a: i64, modulus: i64) -> i64 {
 /// 
 /// # Returns
 /// The n-th root of unity modulo `modulus`.
+///
+/// # Examples
+///
+/// ```
+/// // For modulus = 17^2 = 289 and n = 8, we compute an 8th root of unity.
+/// let modulus = 17 * 17;
+/// let n = 8;
+///
+/// // Compute the omega for the given modulus and order.
+/// let omega = ntt::omega(modulus, n);
+///
+/// // Verify that omega^n is congruent to 1 modulo modulus.
+/// assert_eq!(ntt::mod_exp(omega, n as i64, modulus), 1);
+/// ```
 pub fn omega(modulus: i64, n: usize) -> i64 {
     let factors = factorize(modulus as i64);
     if factors.len() == 1 {