From dcca363d1b83853d51e256362195569b734b049a Mon Sep 17 00:00:00 2001
From: Sean Bowe <ewillbefull@gmail.com>
Date: Sat, 29 Jul 2017 23:20:59 -0600
Subject: [PATCH] Add muln() to PrimeFieldRepr along with tests for muln/divn.

---
 src/bls12_381/fq.rs | 26 ++++++++++++++++++
 src/bls12_381/fr.rs | 26 ++++++++++++++++++
 src/lib.rs          |  6 ++++-
 src/tests/repr.rs   | 64 ++++++++++++++++++++++++++++++++++++++-------
 4 files changed, 111 insertions(+), 11 deletions(-)

diff --git a/src/bls12_381/fq.rs b/src/bls12_381/fq.rs
index 0ef670657..ea69d0efe 100644
--- a/src/bls12_381/fq.rs
+++ b/src/bls12_381/fq.rs
@@ -325,6 +325,32 @@ impl PrimeFieldRepr for FqRepr {
         }
     }
 
+    #[inline(always)]
+    fn muln(&mut self, mut n: u32) {
+        if n >= 64 * 6 {
+            *self = Self::from(0);
+            return;
+        }
+
+        while n >= 64 {
+            let mut t = 0;
+            for i in self.0.iter_mut() {
+                ::std::mem::swap(&mut t, i);
+            }
+            n -= 64;
+        }
+
+        if n > 0 {
+            let mut t = 0;
+            for i in &mut self.0 {
+                let t2 = *i >> (64 - n);
+                *i <<= n;
+                *i |= t;
+                t = t2;
+            }
+        }
+    }
+
     #[inline(always)]
     fn num_bits(&self) -> u32 {
         let mut ret = (6 as u32) * 64;
diff --git a/src/bls12_381/fr.rs b/src/bls12_381/fr.rs
index a91c49cea..6d6c9eff0 100644
--- a/src/bls12_381/fr.rs
+++ b/src/bls12_381/fr.rs
@@ -161,6 +161,32 @@ impl PrimeFieldRepr for FrRepr {
         }
     }
 
+    #[inline(always)]
+    fn muln(&mut self, mut n: u32) {
+        if n >= 64 * 4 {
+            *self = Self::from(0);
+            return;
+        }
+
+        while n >= 64 {
+            let mut t = 0;
+            for i in self.0.iter_mut() {
+                ::std::mem::swap(&mut t, i);
+            }
+            n -= 64;
+        }
+
+        if n > 0 {
+            let mut t = 0;
+            for i in &mut self.0 {
+                let t2 = *i >> (64 - n);
+                *i <<= n;
+                *i |= t;
+                t = t2;
+            }
+        }
+    }
+
     #[inline(always)]
     fn num_bits(&self) -> u32 {
         let mut ret = (4 as u32) * 64;
diff --git a/src/lib.rs b/src/lib.rs
index 0d58c68a9..a8d72cef0 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -352,7 +352,8 @@ pub trait PrimeFieldRepr: Sized +
     /// Add another representation to this one, returning the carry bit.
     fn add_nocarry(&mut self, other: &Self) -> bool;
 
-    /// Compute the number of bits needed to encode this number.
+    /// Compute the number of bits needed to encode this number. Always a
+    /// multiple of 64.
     fn num_bits(&self) -> u32;
 
     /// Returns true iff this number is zero.
@@ -375,6 +376,9 @@ pub trait PrimeFieldRepr: Sized +
     /// it by 2. Overflow is ignored.
     fn mul2(&mut self);
 
+    /// Performs a leftwise bitshift of this number by some amount.
+    fn muln(&mut self, amt: u32);
+
     /// Writes this `PrimeFieldRepr` as a big endian integer. Always writes
     /// `(num_bits` / 8) bytes.
     fn write_be<W: Write>(&self, mut writer: W) -> io::Result<()> {
diff --git a/src/tests/repr.rs b/src/tests/repr.rs
index b0c119cf3..9403df085 100644
--- a/src/tests/repr.rs
+++ b/src/tests/repr.rs
@@ -2,20 +2,64 @@ use rand::{SeedableRng, XorShiftRng};
 use ::{PrimeFieldRepr};
 
 pub fn random_repr_tests<R: PrimeFieldRepr>() {
-	random_encoding_tests::<R>();
+    random_encoding_tests::<R>();
+    random_muln_tests::<R>();
+    random_divn_tests::<R>();
 }
 
 fn random_encoding_tests<R: PrimeFieldRepr>() {
-	let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
+    let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
 
-	for _ in 0..1000 {
-		let r = R::rand(&mut rng);
-		let mut rdecoded = R::default();
+    for _ in 0..1000 {
+        let r = R::rand(&mut rng);
+        let mut rdecoded = R::default();
 
-		let mut v: Vec<u8> = vec![];
-		r.write_be(&mut v).unwrap();
-		rdecoded.read_be(&v[0..]).unwrap();
+        let mut v: Vec<u8> = vec![];
+        r.write_be(&mut v).unwrap();
+        rdecoded.read_be(&v[0..]).unwrap();
 
-		assert_eq!(r, rdecoded);
-	}
+        assert_eq!(r, rdecoded);
+    }
+}
+
+fn random_muln_tests<R: PrimeFieldRepr>() {
+    let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
+
+    for _ in 0..100 {
+        let r = R::rand(&mut rng);
+
+        for shift in 0..(r.num_bits()+1) {
+            let mut r1 = r;
+            let mut r2 = r;
+
+            for _ in 0..shift {
+                r1.mul2();
+            }
+
+            r2.muln(shift);
+
+            assert_eq!(r1, r2);
+        }
+    }
+}
+
+fn random_divn_tests<R: PrimeFieldRepr>() {
+    let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
+
+    for _ in 0..100 {
+        let r = R::rand(&mut rng);
+
+        for shift in 0..(r.num_bits()+1) {
+            let mut r1 = r;
+            let mut r2 = r;
+
+            for _ in 0..shift {
+                r1.div2();
+            }
+
+            r2.divn(shift);
+
+            assert_eq!(r1, r2);
+        }
+    }
 }