From 8be8dc7a5464842dae3cbb90cbb295a0f3f40d66 Mon Sep 17 00:00:00 2001
From: Alivecow <mail@ta-lotz.de>
Date: Sat, 23 Nov 2024 10:26:32 +0100
Subject: [PATCH] feat: Add edge case handling for poly diff

Add handling for cases in which poly is of degree 0 or 1
---
 src/utils/poly.rs | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/src/utils/poly.rs b/src/utils/poly.rs
index 8dc0105..8c5c51a 100644
--- a/src/utils/poly.rs
+++ b/src/utils/poly.rs
@@ -295,9 +295,17 @@ impl Polynomial {
 
     pub fn diff(mut self) -> Self {
         // Pop first element
-        self.polynomial.remove(0);
+        // Check if the polynomial is 1 or less. In this case, output would be [] without check
+        // Output should be [0; 16] however
+        if self.polynomial.len() > 1 {
+            self.polynomial.remove(0);
+        } else {
+            return Polynomial::new(vec![FieldElement::new(vec![0; 16])]);
+        }
 
         for (position, element) in self.polynomial.iter_mut().enumerate() {
+            // Set all uneven degrees to 0, as they were the even degrees before
+            // As we are in GF128, this means they become 0 after mul with even number
             if position % 2 == 1 {
                 *element = FieldElement::new(vec![0; 16]);
             }
@@ -1206,4 +1214,16 @@ mod tests {
 
         assert_eq!(json!(result.to_c_array()), expected);
     }
+
+    #[test]
+    fn test_poly_diff_len1() {
+        let json1 = json!(["IJustWannaTellYouAAAAA==",]);
+        let expected = json!(["AAAAAAAAAAAAAAAAAAAAAA==",]);
+        let element1: Polynomial = Polynomial::from_c_array(&json1);
+        eprintln!("Starting poly sqrt");
+
+        let result = element1.diff();
+
+        assert_eq!(json!(result.to_c_array()), expected);
+    }
 }