Merge pull request #2486 from o1-labs/dw/mvpoly-implement-sub

MVPoly: implement Sub
o1-labs · Aug 27, 2024 · 2bff712 · 2bff712
2 parents 1de8b60 + 5733582
commit 2bff712
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 1 deletion.
diff --git a/mvpoly/src/prime.rs b/mvpoly/src/prime.rs
@@ -144,7 +144,7 @@
 use std::{
     collections::HashMap,
     fmt::{Debug, Formatter, Result},
-    ops::{Add, Mul},
+    ops::{Add, Mul, Sub},
 };
 
 use ark_ff::{One, PrimeField, Zero};
@@ -346,6 +346,55 @@ impl<F: PrimeField, const N: usize, const D: usize> Add<&Dense<F, N, D>> for &De
     }
 }
 
+// Subtraction
+impl<F: PrimeField, const N: usize, const D: usize> Sub for Dense<F, N, D> {
+    type Output = Self;
+
+    fn sub(self, other: Self) -> Self {
+        let mut result = Dense::new();
+        for i in 0..self.coeff.len() {
+            result.coeff[i] = self.coeff[i] - other.coeff[i];
+        }
+        result
+    }
+}
+
+impl<F: PrimeField, const N: usize, const D: usize> Sub<&Dense<F, N, D>> for Dense<F, N, D> {
+    type Output = Dense<F, N, D>;
+
+    fn sub(self, other: &Dense<F, N, D>) -> Dense<F, N, D> {
+        let mut result = Dense::new();
+        for i in 0..self.coeff.len() {
+            result.coeff[i] = self.coeff[i] - other.coeff[i];
+        }
+        result
+    }
+}
+
+impl<F: PrimeField, const N: usize, const D: usize> Sub<Dense<F, N, D>> for &Dense<F, N, D> {
+    type Output = Dense<F, N, D>;
+
+    fn sub(self, other: Dense<F, N, D>) -> Dense<F, N, D> {
+        let mut result = Dense::new();
+        for i in 0..self.coeff.len() {
+            result.coeff[i] = self.coeff[i] - other.coeff[i];
+        }
+        result
+    }
+}
+
+impl<F: PrimeField, const N: usize, const D: usize> Sub<&Dense<F, N, D>> for &Dense<F, N, D> {
+    type Output = Dense<F, N, D>;
+
+    fn sub(self, other: &Dense<F, N, D>) -> Dense<F, N, D> {
+        let mut result = Dense::new();
+        for i in 0..self.coeff.len() {
+            result.coeff[i] = self.coeff[i] - other.coeff[i];
+        }
+        result
+    }
+}
+
 // Multiplication
 impl<F: PrimeField, const N: usize, const D: usize> Mul<Dense<F, N, D>> for Dense<F, N, D> {
     type Output = Self;
@@ -422,4 +471,5 @@ impl<F: PrimeField, const N: usize, const D: usize> Debug for Dense<F, N, D> {
     }
 }
 
+impl<F: PrimeField, const N: usize, const D: usize> Dense<F, N, D> {}
 // TODO: implement From/To Expr<F, Column>
diff --git a/mvpoly/tests/prime.rs b/mvpoly/tests/prime.rs
@@ -193,3 +193,12 @@ fn test_double_is_add_twice() {
     let p3 = p1.clone().double();
     assert_eq!(p2, p3);
 }
+
+#[test]
+fn test_sub_zero() {
+    let mut rng = o1_utils::tests::make_test_rng(None);
+    let p1 = Dense::<Fp, 3, 4>::random(&mut rng);
+    let zero = Dense::<Fp, 3, 4>::zero();
+    let p2 = p1.clone() - zero.clone();
+    assert_eq!(p1.clone(), p2);
+}