chore: simplify add and sub (#84)

jialinli98 · TomAFrench · web-flow · commit 1ac323346420 · 2025-07-16T17:46:59.000+01:00
Co-authored-by: Tom French &lt;15848336+TomAFrench@users.noreply.github.com&gt;
diff --git a/src/bigcurve_test.nr b/src/bigcurve_test.nr
@@ -322,7 +322,7 @@ fn test_incomplete_add_with_hint() {
 }
 
 #[test]
-fn test_add() {
+fn test_constrained_add() {
     unsafe {
         let P = BN254::one();
 
@@ -380,6 +380,73 @@ fn test_add() {
     }
 }
 
+#[test]
+unconstrained fn test_unconstrained_add() {
+    let P = BN254::one();
+
+    // Q = 2P
+    let Q = BN254 {
+        x: BN254_Fq::from_limbs([
+            0x7816a916871ca8d3c208c16d87cfd3,
+            0x44e72e131a029b85045b68181585d9,
+            0x0306,
+        ]),
+        y: BN254_Fq::from_limbs([
+            0xa6a449e3538fc7ff3ebf7a5a18a2c4,
+            0x738c0e0a7c92e7845f96b2ae9c0a68,
+            0x15ed,
+        ]),
+        is_infinity: false,
+    };
+
+    let result = BN254J::from(P.add(Q));
+
+    let P_j = BN254J::from(P);
+    let Q_j = BN254J::from(Q);
+    let expected = P_j.add(Q_j).0;
+
+    assert(result.eq(expected));
+}
+#[test]
+unconstrained fn test_unconstrained_add_one_side_infinity() {
+    //Q + infinity and infinity + Q
+    let Q = BN254::one();
+    let P = BN254::point_at_infinity();
+    let expected = BN254J::from(Q);
+    let result = BN254J::from(Q.add(P));
+    assert(result.eq(expected));
+    let result = BN254J::from(P.add(Q));
+    assert(result.eq(expected));
+}
+
+#[test]
+unconstrained fn test_unconstrained_add_doubling() {
+    // P + P
+    let P = BN254::one();
+    let Q = BN254::one();
+    let P_j = BN254J::from(P);
+    let result = BN254J::from(P.add(Q));
+    let expected = P_j.dbl().0;
+    assert(result.eq(expected));
+}
+
+#[test]
+unconstrained fn test_unconstrained_add_infinity() {
+    // P = -Q
+    let P = BN254::one();
+    let Q = P.neg();
+    let result = BN254J::from(P.add(Q));
+    let expected = BN254J::point_at_infinity();
+    assert(result.eq(expected));
+
+    // both at infinity
+    let P = BN254::point_at_infinity();
+    let Q = BN254::point_at_infinity();
+    let result = BN254J::from(Q.add(P));
+    let expected = BN254J::point_at_infinity();
+    assert(result.eq(expected));
+}
+
 #[test]
 fn test_sub() {
     unsafe {
@@ -439,6 +506,78 @@ fn test_sub() {
     }
 }
 
+#[test]
+unconstrained fn test_unconstrained_sub() {
+    let P = BN254::one();
+
+    // Q = 2P
+    let Q = BN254 {
+        x: BN254_Fq::from_limbs([
+            0x7816a916871ca8d3c208c16d87cfd3,
+            0x44e72e131a029b85045b68181585d9,
+            0x0306,
+        ]),
+        y: BN254_Fq::from_limbs([
+            0xa6a449e3538fc7ff3ebf7a5a18a2c4,
+            0x738c0e0a7c92e7845f96b2ae9c0a68,
+            0x15ed,
+        ]),
+        is_infinity: false,
+    };
+
+    let result = BN254J::from(P.sub(Q));
+
+    let P_j = BN254J::from(P);
+    let Q_j = BN254J::from(Q);
+    let expected = P_j.sub(Q_j).0;
+
+    assert(result.eq(expected));
+}
+
+#[test]
+unconstrained fn test_unconstrained_sub_doubling() {
+    // doubling
+    let P = BN254::one();
+    let Q = BN254::one();
+
+    let result = BN254J::from(P.sub(Q.neg()));
+    let P_j = BN254J::from(P);
+    let expected = P_j.dbl().0;
+    assert(result.eq(expected));
+}
+
+#[test]
+unconstrained fn test_unconstrained_sub_infinity() {
+    // both at infinity
+    let P = BN254::one();
+    let Q = BN254::one();
+    let result = BN254J::from(P.sub(Q));
+    let expected = BN254J::point_at_infinity();
+    assert(result.eq(expected));
+    // both infinity
+    let P = BN254::point_at_infinity();
+    let Q = BN254::point_at_infinity();
+    let result = BN254J::from(Q.sub(P));
+    let expected = BN254J::point_at_infinity();
+    assert(result.eq(expected));
+}
+
+#[test]
+unconstrained fn test_unconstrained_sub_one_side_infinity() {
+    let P = BN254::point_at_infinity();
+    let Q = BN254::one();
+
+    // lhs infinity
+    let result = BN254J::from(P.sub(Q));
+    let expected = BN254J::from(Q.neg());
+    assert(result.eq(expected));
+
+    // rhs infinity
+    let result = BN254J::from(Q.sub(P));
+    let expected = expected.neg();
+    assert(result.eq(expected));
+}
+
 #[test]
 fn test_make_table() {
     unsafe {
diff --git a/src/curve_jac.nr b/src/curve_jac.nr
@@ -304,6 +304,7 @@ where
         let y_equal_predicate = S2.eq(S1);
         let lhs_infinity = self.is_infinity;
         let rhs_infinity = p2.is_infinity;
+
         let double_predicate =
             x_equal_predicate & y_equal_predicate & !lhs_infinity & !rhs_infinity;
         let add_predicate = !x_equal_predicate & !lhs_infinity & !rhs_infinity;
@@ -317,9 +318,13 @@ where
         } else if (infinity_predicate) {
             result = (CurveJ::point_at_infinity(), JTranscript::new());
         } else if (lhs_infinity & !rhs_infinity) {
-            result = (p2, JTranscript::new());
+            let new_transcript =
+                JTranscript { x3: p2.x, y3: p2.y, z3: p2.z, lambda_numerator: B::zero() };
+            result = (p2, new_transcript);
         } else if (rhs_infinity & !lhs_infinity) {
-            result = (self, JTranscript::new());
+            let new_transcript =
+                JTranscript { x3: self.x, y3: self.y, z3: self.z, lambda_numerator: B::zero() };
+            result = (self, new_transcript);
         }
         result
         // let (_, PP): (B, B ) = B::compute_quadratic_expression([[U2, U1]], [[false, true]], [[U2, U1]], [[false, true]], [], []);
diff --git a/src/lib.nr b/src/lib.nr
@@ -184,12 +184,16 @@ pub comptime fn derive_curve_impl(
             // Expensive witness generation! Avoid if possible
             impl std::ops::Add for $typ {
                 fn add(self, other: Self) -> Self {
-                    let lhsJ = $CurveJ::<$field_type, $typ>::from(self);
-                    let rhsJ = $CurveJ::<$field_type, $typ>::from(other);
-
-                    let transcript = unsafe { $AffineTranscript::from_j(lhsJ.add(rhsJ).1) };
-
-                    $crate::add_with_hint::<$field_type, $typ>(self,other, transcript)
+                    let lhsJ = $crate::curve_jac::CurveJ::<$field_type, $typ>::from(self);
+                    let rhsJ = $crate::curve_jac::CurveJ::<$field_type, $typ>::from(other);
+                    let (result_jac, j_transcript) = unsafe { lhsJ.add(rhsJ) };
+                    let transcript = unsafe { $crate::curve_jac::AffineTranscript::from_j(j_transcript) };
+                    if std::runtime::is_unconstrained() {
+                        $typ::from_coordinates(transcript.x3, transcript.y3, result_jac.is_infinity)
+
+                    } else {
+                        $crate::add_with_hint::<$field_type, $typ>(self, other, transcript)           
+                    }
                 }
             }
 
@@ -204,12 +208,17 @@ pub comptime fn derive_curve_impl(
 
             impl std::ops::Sub for $typ {
                 fn sub(self, other: Self) -> Self {
-                    let lhsJ = $CurveJ::<$field_type, $typ>::from(self);
-                    let rhsJ = $CurveJ::<$field_type, $typ>::from(other);
-
-                    let transcript = unsafe { $AffineTranscript::from_j(lhsJ.sub(rhsJ).1) };
-
-                    $crate::sub_with_hint::<$field_type, $typ>(self, other, transcript)
+                    let lhsJ = $crate::curve_jac::CurveJ::<$field_type, $typ>::from(self);
+                    let rhsJ = $crate::curve_jac::CurveJ::<$field_type, $typ>::from(other);
+                    let (result_jac, j_transcript) = unsafe { lhsJ.sub(rhsJ) };
+                    
+                    // Convert back to affine coordinates using the transcript
+                    let transcript = unsafe { $crate::curve_jac::AffineTranscript::from_j(j_transcript) };
+                    if std::runtime::is_unconstrained() {
+                        $typ::from_coordinates(transcript.x3, transcript.y3, result_jac.is_infinity)
+                    } else {
+                        $crate::sub_with_hint::<$field_type, $typ>(self, other, transcript)
+                    }
                 }
             }
 
@@ -549,10 +558,10 @@ pub(crate) fn add_with_hint<B: BigNum, P: BigCurve<B>>(
 
     // If we are skipping the evaluation of a group operation (x2 == x1, y2 == -y1 OR any input points are at infinity),
     // set input operands to 0!
-    let (x1, y1, x2, y2) = if evaluate_group_operation_predicate {
-        (x1, y1, x2, y2)
+    let (x1, y1, x2, y2, x3, y3) = if evaluate_group_operation_predicate {
+        (x1, y1, x2, y2, x3, y3)
     } else {
-        (B::zero(), B::zero(), B::zero(), B::zero())
+        (B::zero(), B::zero(), B::zero(), B::zero(), B::zero(), B::zero())
     };
 
     // lambda * 2y - 3x * x = 0
@@ -678,10 +687,10 @@ pub(crate) fn sub_with_hint<B: BigNum, P: BigCurve<B>>(
 
     // If we are skipping the evaluation of a group operation (x2 == x1, y2 == -y1 OR any input points are at infinity),
     // set input operands to 0!
-    let (x1, y1, x2, y2) = if evaluate_group_operation_predicate {
-        (x1, y1, x2, y2)
+    let (x1, y1, x2, y2, x3, y3) = if evaluate_group_operation_predicate {
+        (x1, y1, x2, y2, x3, y3)
     } else {
-        (B::zero(), B::zero(), B::zero(), B::zero())
+        (B::zero(), B::zero(), B::zero(), B::zero(), B::zero(), B::zero())
     };
 
     // lambda * 2y - 3x*x = 0
