From 30170c18d26cdcc86d643ac86c453ea971808755 Mon Sep 17 00:00:00 2001
From: Joseph Lwanzo Khausi <josephlkausi@gmail.com>
Date: Thu, 21 Mar 2024 12:03:59 +0200
Subject: [PATCH 1/2] fixed sample and primitive_nth_root

---
 baby-stark/src/field.rs | 43 ++++++++++++++++++++++-------------------
 1 file changed, 23 insertions(+), 20 deletions(-)

diff --git a/baby-stark/src/field.rs b/baby-stark/src/field.rs
index bd0892c..440ad15 100644
--- a/baby-stark/src/field.rs
+++ b/baby-stark/src/field.rs
@@ -157,26 +157,29 @@ impl Field{
         }
     }
 
-    // TBD
-    // pub fn primitive_nth_root(self, n : i128) -> Option<FieldElement>{
-    //     if self.p == 85408008396924667383611388730472331217 {
-    //         let mut root = FieldElement::from(85408008396924667383611388730472331217, self);
-    //         let mut order : i128 = 1 << 119;
-    //         while order != n {
-    //             root = root ^ 2;
-    //             order = order / 2;
-    //         }
-    //         Some(root)
-    //     }else {
-    //         None
-    //     }
-    // }
-
-    // TO BE FIXED
-    pub fn sample(self, byte_array : Vec<u8>) -> FieldElement{
-        let mut acc = 0_i128;
-        for b in byte_array.iter(){
-            acc = (acc << 8) ^ (acc as i128);
+    // done
+    pub fn primitive_nth_root(self, n : i128) -> Option<FieldElement>{
+        if self.p == 85408008396924667383611388730472331217 {
+            if !(n <= 1 << 119 && (n & (n-1)) == 0) {
+                println!("Field does not have nth root of unity where n > 2^119 or not power of two.")
+            }
+            let mut root = FieldElement::from(85408008396924667383611388730472331217, self);
+            let mut order : i128 = 1 << 119;
+            while order != n {
+                root = root ^ 2;
+                order = order / 2;
+            }
+            Some(root)
+        }else {
+            None
+        }
+    }
+
+   // FIXED
+    pub fn sample(self, byte_array: Vec<u8>) -> FieldElement {
+        let mut acc: i128 = 0;
+        for &b in byte_array.iter() {
+            acc = (acc << 8) ^ b as i128;
         }
         FieldElement::from(acc % self.p, self)
     }

From f8a894b431031ab7628fb19523eefb65c116217d Mon Sep 17 00:00:00 2001
From: Joseph Lwanzo Khausi <josephlkausi@gmail.com>
Date: Thu, 21 Mar 2024 12:43:10 +0200
Subject: [PATCH 2/2] Correct the primitive_nth_root

---
 baby-stark/src/field.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/baby-stark/src/field.rs b/baby-stark/src/field.rs
index 440ad15..91b4c4d 100644
--- a/baby-stark/src/field.rs
+++ b/baby-stark/src/field.rs
@@ -166,7 +166,7 @@ impl Field{
             let mut root = FieldElement::from(85408008396924667383611388730472331217, self);
             let mut order : i128 = 1 << 119;
             while order != n {
-                root = root ^ 2;
+                root = root.__mul__(root);
                 order = order / 2;
             }
             Some(root)