From 6e9f39c9ffaa828bed592837422a52905b77587d Mon Sep 17 00:00:00 2001
From: bhargav <bhargav.annem@gmail.com>
Date: Sat, 18 Nov 2023 19:45:15 -0800
Subject: [PATCH] feat: mod exp in worker threads

---
 src/ntt.rs     | 5 +++--
 src/numbers.rs | 5 +++++
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/ntt.rs b/src/ntt.rs
index 86ae043..dadbf7c 100644
--- a/src/ntt.rs
+++ b/src/ntt.rs
@@ -111,8 +111,9 @@ fn fft(inp: Vec<BigInt>, c: &Constants, w: BigInt) -> Vec<BigInt> {
     let CHUNK_COUNT = 128;
     let chunk_count = BigInt::from(CHUNK_COUNT);
 
-    (1..N / (2 * CHUNK_COUNT))
-        .for_each(|i| pre[i * CHUNK_COUNT] = w.mod_exp(BigInt::from(i) * chunk_count, MOD));
+    pre.par_chunks_mut(CHUNK_COUNT)
+        .enumerate()
+        .for_each(|(i, arr)| arr[0] = w.mod_exp(BigInt::from(i) * chunk_count, MOD));
     pre.par_chunks_mut(CHUNK_COUNT).for_each(|x| {
         (1..x.len()).for_each(|y| {
             let _x = x.to_vec();
diff --git a/src/numbers.rs b/src/numbers.rs
index 770cced..1d3d1ce 100644
--- a/src/numbers.rs
+++ b/src/numbers.rs
@@ -86,6 +86,11 @@ impl BigInt {
             v: U256::random(&mut OsRng),
         }
     }
+
+    pub fn reverse(&self) -> BigInt {
+        let mut v = self.v;
+        BigInt { v }
+    }
 }
 
 impl From<u16> for BigInt {