MulAssign and AddAssign

lambdaclass · Oct 31, 2024 · a6738c7 · a6738c7
1 parent 26d6078
commit a6738c7
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 14 deletions.
diff --git a/math/src/circle/cosets.rs b/math/src/circle/cosets.rs
@@ -3,8 +3,8 @@ use crate::circle::point::CirclePoint;
 use crate::field::fields::mersenne31::field::Mersenne31Field;
 use alloc::vec::Vec;
 
-/// Given g_n, a generator of the subgroup <g_n> of the circle of size n,
-/// and given a shift, that is a another point of the cirvle,
+/// Given g_n, a generator of the subgroup of size n of the circle, i.e. <g_n>,
+/// and given a shift, that is a another point of the circle,
 /// we define the coset shift + <g_n> which is the set of all the points in
 /// <g_n> plus the shift.
 /// For example, if <g_4> = {p1, p2, p3, p4}, then g_8 + <g_4> = {g_8 + p1, g_8 + p2, g_8 + p3, g_8 + p4}.

diff --git a/math/src/circle/point.rs b/math/src/circle/point.rs
@@ -4,7 +4,7 @@ use crate::field::{
     element::FieldElement,
     fields::mersenne31::{extensions::Degree4ExtensionField, field::Mersenne31Field},
 };
-use core::ops::{Add, Mul};
+use core::ops::{Add, AddAssign, Mul, MulAssign};
 
 /// Given a Field F, we implement here the Group which consists of all the points (x, y) such as
 /// x in F, y in F and x^2 + y^2 = 1, i.e. the Circle. The operation of the group will have
@@ -31,10 +31,10 @@ impl<F: IsField + HasCircleParams<F>> CirclePoint<F> {
     }
 
     /// Computes 2(x, y) = (2x^2 - 1, 2xy).
-    pub fn double(self) -> Self {
+    pub fn double(&self) -> Self {
         Self::new(
             self.x.square().double() - FieldElement::one(),
-            self.x.double() * self.y,
+            self.x.double() * self.y.clone(),
         )
         .unwrap()
     }
@@ -129,14 +129,13 @@ impl<F: IsField + HasCircleParams<F>> PartialEq for CirclePoint<F> {
 /// (a, b) + (c, d) = (a * c - b * d, a * d + b * c)
 impl<F: IsField + HasCircleParams<F>> Add for &CirclePoint<F> {
     type Output = CirclePoint<F>;
-
     fn add(self, other: Self) -> Self::Output {
         let x = &self.x * &other.x - &self.y * &other.y;
         let y = &self.x * &other.y + &self.y * &other.x;
         CirclePoint { x, y }
     }
 }
-impl<F: IsField + HasCircleParams<F>> Add<CirclePoint<F>> for CirclePoint<F> {
+impl<F: IsField + HasCircleParams<F>> Add for CirclePoint<F> {
     type Output = CirclePoint<F>;
     fn add(self, rhs: CirclePoint<F>) -> Self::Output {
         &self + &rhs
@@ -154,28 +153,58 @@ impl<F: IsField + HasCircleParams<F>> Add<&CirclePoint<F>> for CirclePoint<F> {
         &self + rhs
     }
 }
-
+impl<F: IsField + HasCircleParams<F>> AddAssign<&CirclePoint<F>> for CirclePoint<F> {
+    fn add_assign(&mut self, rhs: &CirclePoint<F>) {
+        *self = &*self + rhs;
+    }
+}
+impl<F: IsField + HasCircleParams<F>> AddAssign<CirclePoint<F>> for CirclePoint<F> {
+    fn add_assign(&mut self, rhs: CirclePoint<F>) {
+        *self += &rhs;
+    }
+}
 /// Multiplication between a point and a scalar (i.e. group operation repeatedly):
 /// (x, y) * n = (x ,y) + ... + (x, y) n-times.
-impl<F: IsField + HasCircleParams<F>> Mul<u128> for CirclePoint<F> {
+impl<F: IsField + HasCircleParams<F>> Mul<u128> for &CirclePoint<F> {
     type Output = CirclePoint<F>;
-
-    fn mul(self, scalar: u128) -> Self {
+    fn mul(self, scalar: u128) -> Self::Output {
         let mut scalar = scalar;
-        let mut res = Self::zero();
-        let mut cur = self;
+        let mut res = CirclePoint::<F>::zero();
+        let mut cur = self.clone();
         loop {
             if scalar == 0 {
                 return res;
             }
             if scalar & 1 == 1 {
-                res = &res + &cur;
+                res += &cur;
             }
             cur = cur.double();
             scalar >>= 1;
         }
     }
 }
+impl<F: IsField + HasCircleParams<F>> Mul<u128> for CirclePoint<F> {
+    type Output = CirclePoint<F>;
+    fn mul(self, scalar: u128) -> Self::Output {
+        &self * scalar
+    }
+}
+impl<F: IsField + HasCircleParams<F>> MulAssign<u128> for CirclePoint<F> {
+    fn mul_assign(&mut self, scalar: u128) {
+        let mut scalar = scalar;
+        let mut res = CirclePoint::<F>::zero();
+        loop {
+            if scalar == 0 {
+                *self = res.clone();
+            }
+            if scalar & 1 == 1 {
+                res += &*self;
+            }
+            *self = self.double();
+            scalar >>= 1;
+        }
+    }
+}
 
 #[cfg(test)]
 mod tests {