Add assignment operators

Addresses part of #116
LLFourn · Sep 19, 2022 · 0ddb1b8 · 0ddb1b8
1 parent d37adfc
commit 0ddb1b8
Show file tree

Hide file tree

Showing 2 changed files with 105 additions and 2 deletions.
diff --git a/secp256kfun/src/point.rs b/secp256kfun/src/point.rs
@@ -4,7 +4,10 @@ use crate::{
     marker::*,
     op, Scalar,
 };
-use core::marker::PhantomData;
+use core::{
+    marker::PhantomData,
+    ops::{AddAssign, SubAssign},
+};
 use rand_core::RngCore;
 
 /// A point on the secp256k1 elliptic curve.
@@ -498,6 +501,30 @@ crate::impl_fromstr_deserialize! {
     }
 }
 
+impl<TR, SL, SR, ZR> AddAssign<Point<TR, SR, ZR>> for Point<NonNormal, SL, Zero> {
+    fn add_assign(&mut self, rhs: Point<TR, SR, ZR>) {
+        *self = crate::op::point_add(self, &rhs).set_secrecy::<SL>()
+    }
+}
+
+impl<TR, SL, SR, ZR> AddAssign<&Point<TR, SR, ZR>> for Point<NonNormal, SL, Zero> {
+    fn add_assign(&mut self, rhs: &Point<TR, SR, ZR>) {
+        *self = crate::op::point_add(self, rhs).set_secrecy::<SL>()
+    }
+}
+
+impl<TR, SL, SR, ZR> SubAssign<&Point<TR, SR, ZR>> for Point<NonNormal, SL, Zero> {
+    fn sub_assign(&mut self, rhs: &Point<TR, SR, ZR>) {
+        *self = crate::op::point_sub(self, rhs).set_secrecy::<SL>()
+    }
+}
+
+impl<TR, SL, SR, ZR> SubAssign<Point<TR, SR, ZR>> for Point<NonNormal, SL, Zero> {
+    fn sub_assign(&mut self, rhs: Point<TR, SR, ZR>) {
+        *self = crate::op::point_sub(self, &rhs).set_secrecy::<SL>()
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -732,4 +759,17 @@ mod test {
         let mult_point = g!({ Scalar::random(&mut rand::thread_rng()) } * G);
         assert!(format!("{:?}", mult_point).starts_with("Point<NonNormal,Public,NonZero>"));
     }
+
+    #[test]
+    fn assign_tests() {
+        let a_orig = Point::random(&mut rand::thread_rng())
+            .mark_zero()
+            .non_normal();
+        let mut a = a_orig.clone();
+        let b = Point::random(&mut rand::thread_rng());
+        a += b;
+        assert_eq!(a, op::point_add(&a_orig, &b));
+        a -= b;
+        assert_eq!(a, a_orig);
+    }
 }
diff --git a/secp256kfun/src/scalar.rs b/secp256kfun/src/scalar.rs
@@ -1,6 +1,9 @@
 //! Scalar arithmetic (integers mod the secp256k1 group order)
 use crate::{backend, hash::HashInto, marker::*, op};
-use core::marker::PhantomData;
+use core::{
+    marker::PhantomData,
+    ops::{AddAssign, MulAssign, SubAssign},
+};
 use digest::{generic_array::typenum::U32, Digest};
 use rand_core::RngCore;
 
@@ -348,6 +351,54 @@ where
     }
 }
 
+impl<SL, SR, ZR> AddAssign<Scalar<SR, ZR>> for Scalar<SL, Zero> {
+    fn add_assign(&mut self, rhs: Scalar<SR, ZR>) {
+        *self = crate::op::scalar_add(&self, &rhs).set_secrecy::<SL>();
+    }
+}
+
+impl<SL, SR, ZR> AddAssign<&Scalar<SR, ZR>> for Scalar<SL, Zero> {
+    fn add_assign(&mut self, rhs: &Scalar<SR, ZR>) {
+        *self = crate::op::scalar_add(&self, rhs).set_secrecy::<SL>();
+    }
+}
+
+impl<SL, SR, ZR> SubAssign<&Scalar<SR, ZR>> for Scalar<SL, Zero> {
+    fn sub_assign(&mut self, rhs: &Scalar<SR, ZR>) {
+        *self = crate::op::scalar_sub(&self, rhs).set_secrecy::<SL>();
+    }
+}
+
+impl<SL, SR, ZR> SubAssign<Scalar<SR, ZR>> for Scalar<SL, Zero> {
+    fn sub_assign(&mut self, rhs: Scalar<SR, ZR>) {
+        *self = crate::op::scalar_sub(&self, &rhs).set_secrecy::<SL>();
+    }
+}
+
+impl<SL, SR> MulAssign<Scalar<SR, NonZero>> for Scalar<SL, NonZero> {
+    fn mul_assign(&mut self, rhs: Scalar<SR, NonZero>) {
+        *self = crate::op::scalar_mul(&self, &rhs).set_secrecy::<SL>();
+    }
+}
+
+impl<SL, SR> MulAssign<&Scalar<SR, NonZero>> for Scalar<SL, NonZero> {
+    fn mul_assign(&mut self, rhs: &Scalar<SR, NonZero>) {
+        *self = crate::op::scalar_mul(&self, rhs).set_secrecy::<SL>();
+    }
+}
+
+impl<SL, SR, ZR: ZeroChoice> MulAssign<Scalar<SR, ZR>> for Scalar<SL, Zero> {
+    fn mul_assign(&mut self, rhs: Scalar<SR, ZR>) {
+        *self = crate::op::scalar_mul(&self, &rhs).set_secrecy::<SL>();
+    }
+}
+
+impl<SL, SR, ZR: ZeroChoice> MulAssign<&Scalar<SR, ZR>> for Scalar<SL, Zero> {
+    fn mul_assign(&mut self, rhs: &Scalar<SR, ZR>) {
+        *self = crate::op::scalar_mul(&self, rhs).set_secrecy::<SL>();
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -484,4 +535,16 @@ mod test {
             )
         );
     }
+
+    #[test]
+    fn assign_tests() {
+        let mut a = Scalar::from(42);
+        let b = Scalar::from(1337).public();
+        a += b;
+        assert_eq!(a, Scalar::from(1379));
+        a -= b;
+        assert_eq!(a, Scalar::from(42));
+        a *= b;
+        assert_eq!(a, Scalar::from(42 * 1337));
+    }
 }