From e9d3923829affc219bda28b31f2845cce4ba9e38 Mon Sep 17 00:00:00 2001
From: Sean Bowe <ewillbefull@gmail.com>
Date: Mon, 11 Dec 2017 23:06:05 -0700
Subject: [PATCH] Refactor jubjub implementation to be abstract over field,
 parameters.

---
 src/jubjub/edwards.rs    | 237 +++-------------------------
 src/jubjub/mod.rs        |  97 +++++-------
 src/jubjub/montgomery.rs | 328 +++------------------------------------
 src/jubjub/tests.rs      | 299 +++++++++++++++++++++++++++++++++++
 4 files changed, 385 insertions(+), 576 deletions(-)
 create mode 100644 src/jubjub/tests.rs

diff --git a/src/jubjub/edwards.rs b/src/jubjub/edwards.rs
index 68d323b..bf9b3f4 100644
--- a/src/jubjub/edwards.rs
+++ b/src/jubjub/edwards.rs
@@ -1,5 +1,4 @@
 use pairing::{
-    Engine,
     Field,
     SqrtField,
     PrimeField,
@@ -8,11 +7,10 @@ use pairing::{
 };
 
 use super::{
+    JubjubEngine,
     JubjubParams,
     Unknown,
     PrimeOrder,
-    Fs,
-    FsRepr,
     montgomery
 };
 
@@ -24,7 +22,7 @@ use std::marker::PhantomData;
 
 // Represents the affine point (X/Z, Y/Z) via the extended
 // twisted Edwards coordinates.
-pub struct Point<E: Engine, Subgroup> {
+pub struct Point<E: JubjubEngine, Subgroup> {
     x: E::Fr,
     y: E::Fr,
     t: E::Fr,
@@ -32,7 +30,7 @@ pub struct Point<E: Engine, Subgroup> {
     _marker: PhantomData<Subgroup>
 }
 
-fn convert_subgroup<E: Engine, S1, S2>(from: &Point<E, S1>) -> Point<E, S2>
+fn convert_subgroup<E: JubjubEngine, S1, S2>(from: &Point<E, S1>) -> Point<E, S2>
 {
     Point {
         x: from.x,
@@ -43,7 +41,7 @@ fn convert_subgroup<E: Engine, S1, S2>(from: &Point<E, S1>) -> Point<E, S2>
     }
 }
 
-impl<E: Engine> From<Point<E, PrimeOrder>> for Point<E, Unknown>
+impl<E: JubjubEngine> From<Point<E, PrimeOrder>> for Point<E, Unknown>
 {
     fn from(p: Point<E, PrimeOrder>) -> Point<E, Unknown>
     {
@@ -51,14 +49,14 @@ impl<E: Engine> From<Point<E, PrimeOrder>> for Point<E, Unknown>
     }
 }
 
-impl<E: Engine, Subgroup> Clone for Point<E, Subgroup>
+impl<E: JubjubEngine, Subgroup> Clone for Point<E, Subgroup>
 {
     fn clone(&self) -> Self {
         convert_subgroup(self)
     }
 }
 
-impl<E: Engine, Subgroup> PartialEq for Point<E, Subgroup> {
+impl<E: JubjubEngine, Subgroup> PartialEq for Point<E, Subgroup> {
     fn eq(&self, other: &Point<E, Subgroup>) -> bool {
         // p1 = (x1/z1, y1/z1)
         // p2 = (x2/z2, y2/z2)
@@ -82,9 +80,9 @@ impl<E: Engine, Subgroup> PartialEq for Point<E, Subgroup> {
     }
 }
 
-impl<E: Engine> Point<E, Unknown> {
+impl<E: JubjubEngine> Point<E, Unknown> {
     /// This guarantees the point is in the prime order subgroup
-    pub fn mul_by_cofactor(&self, params: &JubjubParams<E>) -> Point<E, PrimeOrder>
+    pub fn mul_by_cofactor(&self, params: &E::Params) -> Point<E, PrimeOrder>
     {
         let tmp = self.double(params)
                       .double(params)
@@ -93,7 +91,7 @@ impl<E: Engine> Point<E, Unknown> {
         convert_subgroup(&tmp)
     }
 
-    pub fn rand<R: Rng>(rng: &mut R, params: &JubjubParams<E>) -> Self
+    pub fn rand<R: Rng>(rng: &mut R, params: &E::Params) -> Self
     {
         loop {
             // given an x on the curve, y^2 = (1 + x^2) / (1 - dx^2)
@@ -104,7 +102,7 @@ impl<E: Engine> Point<E, Unknown> {
             let mut num = E::Fr::one();
             num.add_assign(&x2);
 
-            x2.mul_assign(&params.edwards_d);
+            x2.mul_assign(params.edwards_d());
 
             let mut den = E::Fr::one();
             den.sub_assign(&x2);
@@ -139,11 +137,11 @@ impl<E: Engine> Point<E, Unknown> {
     }
 }
 
-impl<E: Engine, Subgroup> Point<E, Subgroup> {
+impl<E: JubjubEngine, Subgroup> Point<E, Subgroup> {
     /// Convert from a Montgomery point
     pub fn from_montgomery(
         m: &montgomery::Point<E, Subgroup>,
-        params: &JubjubParams<E>
+        params: &E::Params
     ) -> Self
     {
         match m.into_xy() {
@@ -214,7 +212,7 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
 
                     // u = xs
                     let mut u = x;
-                    u.mul_assign(&params.scale);
+                    u.mul_assign(params.scale());
 
                     // v = x - 1
                     let mut v = x;
@@ -251,8 +249,8 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
 
     /// Attempts to cast this as a prime order element, failing if it's
     /// not in the prime order subgroup.
-    pub fn as_prime_order(&self, params: &JubjubParams<E>) -> Option<Point<E, PrimeOrder>> {
-        if self.mul(Fs::char(), params) == Point::zero() {
+    pub fn as_prime_order(&self, params: &E::Params) -> Option<Point<E, PrimeOrder>> {
+        if self.mul(E::Fs::char(), params) == Point::zero() {
             Some(convert_subgroup(self))
         } else {
             None
@@ -291,11 +289,11 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
         p
     }
 
-    pub fn double(&self, params: &JubjubParams<E>) -> Self {
+    pub fn double(&self, params: &E::Params) -> Self {
         self.add(self, params)
     }
 
-    pub fn add(&self, other: &Self, params: &JubjubParams<E>) -> Self
+    pub fn add(&self, other: &Self, params: &E::Params) -> Self
     {
         // A = x1 * x2
         let mut a = self.x;
@@ -306,7 +304,7 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
         b.mul_assign(&other.y);
 
         // C = d * t1 * t2
-        let mut c = params.edwards_d;
+        let mut c = params.edwards_d().clone();
         c.mul_assign(&self.t);
         c.mul_assign(&other.t);
 
@@ -363,7 +361,11 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
         }
     }
 
-    pub fn mul<S: Into<FsRepr>>(&self, scalar: S, params: &JubjubParams<E>) -> Self
+    pub fn mul<S: Into<<E::Fs as PrimeField>::Repr>>(
+        &self,
+        scalar: S,
+        params: &E::Params
+    ) -> Self
     {
         let mut res = Self::zero();
 
@@ -378,196 +380,3 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
         res
     }
 }
-
-#[cfg(test)]
-mod test {
-    use rand::{XorShiftRng, SeedableRng, Rand};
-    use super::{JubjubParams, Point, PrimeOrder, Fs};
-    use pairing::bls12_381::{Bls12};
-    use pairing::{Engine, Field};
-
-    fn is_on_curve<E: Engine>(
-        x: E::Fr,
-        y: E::Fr,
-        params: &JubjubParams<E>
-    ) -> bool
-    {
-        let mut x2 = x;
-        x2.square();
-
-        let mut y2 = y;
-        y2.square();
-
-        // -x^2 + y^2
-        let mut lhs = y2;
-        lhs.sub_assign(&x2);
-
-        // 1 + d x^2 y^2
-        let mut rhs = y2;
-        rhs.mul_assign(&x2);
-        rhs.mul_assign(&params.edwards_d);
-        rhs.add_assign(&E::Fr::one());
-
-        lhs == rhs
-    }
-
-    #[test]
-    fn test_rand() {
-        let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = JubjubParams::new();
-
-        for _ in 0..100 {
-            let (x, y) = Point::rand(&mut rng, &params).into_xy();
-
-            assert!(is_on_curve(x, y, &params));
-        }
-    }
-
-    #[test]
-    fn test_identities() {
-        let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = JubjubParams::new();
-
-        let z = Point::<Bls12, PrimeOrder>::zero();
-        assert!(z.double(&params) == z);
-        assert!(z.negate() == z);
-
-        for _ in 0..100 {
-            let r = Point::rand(&mut rng, &params);
-
-            assert!(r.add(&Point::zero(), &params) == r);
-            assert!(r.add(&r.negate(), &params) == Point::zero());
-        }
-    }
-
-    #[test]
-    fn test_associativity() {
-        let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = JubjubParams::new();
-
-        for _ in 0..1000 {
-            let a = Point::rand(&mut rng, &params);
-            let b = Point::rand(&mut rng, &params);
-            let c = Point::rand(&mut rng, &params);
-
-            assert!(a.add(&b, &params).add(&c, &params) == c.add(&a, &params).add(&b, &params));
-        }
-    }
-
-    #[test]
-    fn test_order() {
-        let rng = &mut XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = &JubjubParams::new();
-
-        // The neutral element is in the prime order subgroup.
-        assert!(Point::<Bls12, PrimeOrder>::zero().as_prime_order(params).is_some());
-
-        for _ in 0..50 {
-            // Pick a random point and multiply it by the cofactor
-            let base = Point::rand(rng, params).mul_by_cofactor(params);
-
-            // Any point multiplied by the cofactor will be in the prime
-            // order subgroup
-            assert!(base.as_prime_order(params).is_some());
-        }
-
-        // It's very likely that at least one out of 50 random points on the curve
-        // is not in the prime order subgroup.
-        let mut at_least_one_not_in_prime_order_subgroup = false;
-        for _ in 0..50 {
-            // Pick a random point.
-            let base = Point::rand(rng, params);
-
-            at_least_one_not_in_prime_order_subgroup |= base.as_prime_order(params).is_none();
-        }
-        assert!(at_least_one_not_in_prime_order_subgroup);
-    }
-
-    #[test]
-    fn test_mul_associativity() {
-        let rng = &mut XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = &JubjubParams::new();
-
-        for _ in 0..100 {
-            // Pick a random point and multiply it by the cofactor
-            let base = Point::rand(rng, params).mul_by_cofactor(params);
-
-            let mut a = Fs::rand(rng);
-            let b = Fs::rand(rng);
-            let c = Fs::rand(rng);
-
-            let res1 = base.mul(a, params).mul(b, params).mul(c, params);
-            let res2 = base.mul(b, params).mul(c, params).mul(a, params);
-            let res3 = base.mul(c, params).mul(a, params).mul(b, params);
-            a.mul_assign(&b);
-            a.mul_assign(&c);
-            let res4 = base.mul(a, params);
-
-            assert!(res1 == res2);
-            assert!(res2 == res3);
-            assert!(res3 == res4);
-
-            let (x, y) = res1.into_xy();
-            assert!(is_on_curve(x, y, params));
-
-            let (x, y) = res2.into_xy();
-            assert!(is_on_curve(x, y, params));
-
-            let (x, y) = res3.into_xy();
-            assert!(is_on_curve(x, y, params));
-        }
-    }
-
-    #[test]
-    fn test_montgomery_conversion() {
-        use super::montgomery;
-
-        let rng = &mut XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = &JubjubParams::new();
-
-        for _ in 0..200 {
-            // compute base in montgomery
-            let base = montgomery::Point::rand(rng, params);
-
-            // sample random exponent
-            let exp = Fs::rand(rng);
-
-            // exponentiate in montgomery, convert to edwards
-            let ed_expected = Point::from_montgomery(&base.mul(exp, params), params);
-
-            // convert to edwards and exponentiate
-            let ed_exponentiated = Point::from_montgomery(&base, params).mul(exp, params);
-
-            let (x, y) = ed_expected.into_xy();
-            assert!(is_on_curve(x, y, params));
-
-            assert!(ed_exponentiated == ed_expected);
-        }
-    }
-
-    #[test]
-    fn test_back_and_forth() {
-        use super::montgomery;
-
-        let rng = &mut XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = &JubjubParams::new();
-
-        for _ in 0..200 {
-            // compute base in montgomery
-            let base = montgomery::Point::rand(rng, params);
-
-            // convert to edwards
-            let base_ed = Point::from_montgomery(&base, params);
-
-            {
-                let (x, y) = base_ed.into_xy();
-                assert!(is_on_curve(x, y, params));
-            }
-
-            // convert back to montgomery
-            let base_mont = montgomery::Point::from_edwards(&base_ed, params);
-
-            assert!(base == base_mont);
-        }
-    }
-}
diff --git a/src/jubjub/mod.rs b/src/jubjub/mod.rs
index 61b9992..e510eb8 100644
--- a/src/jubjub/mod.rs
+++ b/src/jubjub/mod.rs
@@ -16,7 +16,8 @@
 
 use pairing::{
     Engine,
-    PrimeField
+    PrimeField,
+    SqrtField
 };
 
 use pairing::bls12_381::{
@@ -24,28 +25,48 @@ use pairing::bls12_381::{
     Fr
 };
 
-mod fs;
-
-pub use self::fs::{Fs, FsRepr};
-
 pub mod edwards;
 pub mod montgomery;
 
-/// These are the pre-computed parameters of the Jubjub
-/// curve.
-pub struct JubjubParams<E: Engine> {
-    edwards_d: E::Fr,
-    montgomery_a: E::Fr,
+#[cfg(test)]
+pub mod tests;
 
-    scale: E::Fr
+pub trait JubjubEngine: Engine {
+    type Fs: PrimeField + SqrtField;
+    type Params: JubjubParams<Self>;
+}
+
+pub trait JubjubParams<E: JubjubEngine>: Sized {
+    fn edwards_d(&self) -> &E::Fr;
+    fn montgomery_a(&self) -> &E::Fr;
+    fn scale(&self) -> &E::Fr;
 }
 
 pub enum Unknown { }
 pub enum PrimeOrder { }
 
-impl JubjubParams<Bls12> {
+pub mod fs;
+
+impl JubjubEngine for Bls12 {
+    type Fs = self::fs::Fs;
+    type Params = JubjubBls12;
+}
+
+pub struct JubjubBls12 {
+    edwards_d: Fr,
+    montgomery_a: Fr,
+    scale: Fr
+}
+
+impl JubjubParams<Bls12> for JubjubBls12 {
+    fn edwards_d(&self) -> &Fr { &self.edwards_d }
+    fn montgomery_a(&self) -> &Fr { &self.montgomery_a }
+    fn scale(&self) -> &Fr { &self.scale }
+}
+
+impl JubjubBls12 {
     pub fn new() -> Self {
-        JubjubParams {
+        JubjubBls12 {
             // d = -(10240/10241)
             edwards_d: Fr::from_str("19257038036680949359750312669786877991949435402254120286184196891950884077233").unwrap(),
             // A = 40962
@@ -56,51 +77,9 @@ impl JubjubParams<Bls12> {
     }
 }
 
-#[cfg(test)]
-mod test {
-    use pairing::{Field, SqrtField, LegendreSymbol, PrimeField};
-    use pairing::bls12_381::{Fr};
-    use super::JubjubParams;
+#[test]
+fn test_jubjub_bls12() {
+    let params = JubjubBls12::new();
 
-    #[test]
-    fn test_params() {
-        let params = JubjubParams::new();
-
-        // a = -1
-        let mut a = Fr::one();
-        a.negate();
-
-        {
-            // The twisted Edwards addition law is complete when d is nonsquare
-            // and a is square.
-
-            assert!(params.edwards_d.legendre() == LegendreSymbol::QuadraticNonResidue);
-            assert!(a.legendre() == LegendreSymbol::QuadraticResidue);
-        }
-
-        {
-            // Check that A^2 - 4 is nonsquare:
-            let mut tmp = params.montgomery_a;
-            tmp.square();
-            tmp.sub_assign(&Fr::from_str("4").unwrap());
-            assert!(tmp.legendre() == LegendreSymbol::QuadraticNonResidue);
-        }
-
-        {
-            // Check that A - 2 is nonsquare:
-            let mut tmp = params.montgomery_a;
-            tmp.sub_assign(&Fr::from_str("2").unwrap());
-            assert!(tmp.legendre() == LegendreSymbol::QuadraticNonResidue);
-        }
-
-        {
-            // Check the validity of the scaling factor
-            let mut tmp = a;
-            tmp.sub_assign(&params.edwards_d);
-            tmp = tmp.inverse().unwrap();
-            tmp.mul_assign(&Fr::from_str("4").unwrap());
-            tmp = tmp.sqrt().unwrap();
-            assert_eq!(tmp, params.scale);
-        }
-    }
+    tests::test_suite::<Bls12>(&params);
 }
diff --git a/src/jubjub/montgomery.rs b/src/jubjub/montgomery.rs
index 9d7c944..6af0e6d 100644
--- a/src/jubjub/montgomery.rs
+++ b/src/jubjub/montgomery.rs
@@ -1,5 +1,4 @@
 use pairing::{
-    Engine,
     Field,
     SqrtField,
     PrimeField,
@@ -8,11 +7,10 @@ use pairing::{
 };
 
 use super::{
+    JubjubEngine,
     JubjubParams,
     Unknown,
     PrimeOrder,
-    Fs,
-    FsRepr,
     edwards
 };
 
@@ -24,14 +22,14 @@ use std::marker::PhantomData;
 
 // Represents the affine point (X/Z, Y/Z) via the extended
 // twisted Edwards coordinates.
-pub struct Point<E: Engine, Subgroup> {
+pub struct Point<E: JubjubEngine, Subgroup> {
     x: E::Fr,
     y: E::Fr,
     infinity: bool,
     _marker: PhantomData<Subgroup>
 }
 
-fn convert_subgroup<E: Engine, S1, S2>(from: &Point<E, S1>) -> Point<E, S2>
+fn convert_subgroup<E: JubjubEngine, S1, S2>(from: &Point<E, S1>) -> Point<E, S2>
 {
     Point {
         x: from.x,
@@ -41,7 +39,7 @@ fn convert_subgroup<E: Engine, S1, S2>(from: &Point<E, S1>) -> Point<E, S2>
     }
 }
 
-impl<E: Engine> From<Point<E, PrimeOrder>> for Point<E, Unknown>
+impl<E: JubjubEngine> From<Point<E, PrimeOrder>> for Point<E, Unknown>
 {
     fn from(p: Point<E, PrimeOrder>) -> Point<E, Unknown>
     {
@@ -49,14 +47,14 @@ impl<E: Engine> From<Point<E, PrimeOrder>> for Point<E, Unknown>
     }
 }
 
-impl<E: Engine, Subgroup> Clone for Point<E, Subgroup>
+impl<E: JubjubEngine, Subgroup> Clone for Point<E, Subgroup>
 {
     fn clone(&self) -> Self {
         convert_subgroup(self)
     }
 }
 
-impl<E: Engine, Subgroup> PartialEq for Point<E, Subgroup> {
+impl<E: JubjubEngine, Subgroup> PartialEq for Point<E, Subgroup> {
     fn eq(&self, other: &Point<E, Subgroup>) -> bool {
         match (self.infinity, other.infinity) {
             (true, true) => true,
@@ -68,9 +66,9 @@ impl<E: Engine, Subgroup> PartialEq for Point<E, Subgroup> {
     }
 }
 
-impl<E: Engine> Point<E, Unknown> {
+impl<E: JubjubEngine> Point<E, Unknown> {
     /// This guarantees the point is in the prime order subgroup
-    pub fn mul_by_cofactor(&self, params: &JubjubParams<E>) -> Point<E, PrimeOrder>
+    pub fn mul_by_cofactor(&self, params: &E::Params) -> Point<E, PrimeOrder>
     {
         let tmp = self.double(params)
                       .double(params)
@@ -79,7 +77,7 @@ impl<E: Engine> Point<E, Unknown> {
         convert_subgroup(&tmp)
     }
 
-    pub fn rand<R: Rng>(rng: &mut R, params: &JubjubParams<E>) -> Self
+    pub fn rand<R: Rng>(rng: &mut R, params: &E::Params) -> Self
     {
         loop {
             // given an x on the curve, y^2 = x^3 + A*x^2 + x
@@ -89,7 +87,7 @@ impl<E: Engine> Point<E, Unknown> {
             x2.square();
 
             let mut rhs = x2;
-            rhs.mul_assign(&params.montgomery_a);
+            rhs.mul_assign(params.montgomery_a());
             rhs.add_assign(&x);
             x2.mul_assign(&x);
             rhs.add_assign(&x2);
@@ -113,11 +111,11 @@ impl<E: Engine> Point<E, Unknown> {
     }
 }
 
-impl<E: Engine, Subgroup> Point<E, Subgroup> {
+impl<E: JubjubEngine, Subgroup> Point<E, Subgroup> {
     /// Convert from an Edwards point
     pub fn from_edwards(
         e: &edwards::Point<E, Subgroup>,
-        params: &JubjubParams<E>
+        params: &E::Params
     ) -> Self
     {
         let (x, y) = e.into_xy();
@@ -168,7 +166,7 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
                 v.mul_assign(&x.inverse().unwrap());
 
                 // Scale it into the correct curve constants
-                v.mul_assign(&params.scale);
+                v.mul_assign(params.scale());
 
                 Point {
                     x: u,
@@ -182,8 +180,8 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
 
     /// Attempts to cast this as a prime order element, failing if it's
     /// not in the prime order subgroup.
-    pub fn as_prime_order(&self, params: &JubjubParams<E>) -> Option<Point<E, PrimeOrder>> {
-        if self.mul(Fs::char(), params) == Point::zero() {
+    pub fn as_prime_order(&self, params: &E::Params) -> Option<Point<E, PrimeOrder>> {
+        if self.mul(E::Fs::char(), params) == Point::zero() {
             Some(convert_subgroup(self))
         } else {
             None
@@ -216,7 +214,7 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
         p
     }
 
-    pub fn double(&self, params: &JubjubParams<E>) -> Self {
+    pub fn double(&self, params: &E::Params) -> Self {
         if self.infinity {
             return Point::zero();
         }
@@ -227,7 +225,7 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
 
         let mut delta = E::Fr::one();
         {
-            let mut tmp = params.montgomery_a;
+            let mut tmp = params.montgomery_a().clone();
             tmp.mul_assign(&self.x);
             tmp.double();
             delta.add_assign(&tmp);
@@ -247,7 +245,7 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
 
         let mut x3 = delta;
         x3.square();
-        x3.sub_assign(&params.montgomery_a);
+        x3.sub_assign(params.montgomery_a());
         x3.sub_assign(&self.x);
         x3.sub_assign(&self.x);
 
@@ -265,7 +263,7 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
         }
     }
 
-    pub fn add(&self, other: &Self, params: &JubjubParams<E>) -> Self
+    pub fn add(&self, other: &Self, params: &E::Params) -> Self
     {
         match (self.infinity, other.infinity) {
             (true, true) => Point::zero(),
@@ -289,7 +287,7 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
 
                     let mut x3 = delta;
                     x3.square();
-                    x3.sub_assign(&params.montgomery_a);
+                    x3.sub_assign(params.montgomery_a());
                     x3.sub_assign(&self.x);
                     x3.sub_assign(&other.x);
 
@@ -310,7 +308,11 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
         }
     }
 
-    pub fn mul<S: Into<FsRepr>>(&self, scalar: S, params: &JubjubParams<E>) -> Self
+    pub fn mul<S: Into<<E::Fs as PrimeField>::Repr>>(
+        &self,
+        scalar: S,
+        params: &E::Params
+    ) -> Self
     {
         let mut res = Self::zero();
 
@@ -325,283 +327,3 @@ impl<E: Engine, Subgroup> Point<E, Subgroup> {
         res
     }
 }
-
-#[cfg(test)]
-mod test {
-    use rand::{XorShiftRng, SeedableRng, Rand};
-    use super::{JubjubParams, Point, PrimeOrder, Unknown, Fs};
-    use pairing::bls12_381::{Bls12, Fr};
-    use pairing::{Engine, Field, PrimeField};
-    use std::marker::PhantomData;
-
-    fn is_on_curve<E: Engine>(
-        x: E::Fr,
-        y: E::Fr,
-        params: &JubjubParams<E>
-    ) -> bool
-    {
-        let mut lhs = y;
-        lhs.square();
-
-        let mut x2 = x;
-        x2.square();
-
-        let mut x3 = x2;
-        x3.mul_assign(&x);
-
-        let mut rhs = x2;
-        rhs.mul_assign(&params.montgomery_a);
-        rhs.add_assign(&x);
-        rhs.add_assign(&x3);
-
-        lhs == rhs
-    }
-
-    #[test]
-    fn test_rand() {
-        let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = JubjubParams::new();
-
-        for _ in 0..100 {
-            let (x, y) = Point::rand(&mut rng, &params).into_xy().unwrap();
-
-            assert!(is_on_curve(x, y, &params));
-        }
-    }
-
-    #[test]
-    fn test_identities() {
-        let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = JubjubParams::new();
-
-        let z = Point::<Bls12, PrimeOrder>::zero();
-        assert!(z.double(&params) == z);
-        assert!(z.negate() == z);
-
-        for _ in 0..100 {
-            let r = Point::rand(&mut rng, &params);
-
-            assert!(r.add(&Point::zero(), &params) == r);
-            assert!(r.add(&r.negate(), &params) == Point::zero());
-        }
-    }
-
-    #[test]
-    fn test_associativity() {
-        let mut rng = XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = JubjubParams::new();
-
-        for _ in 0..1000 {
-            let a = Point::rand(&mut rng, &params);
-            let b = Point::rand(&mut rng, &params);
-            let c = Point::rand(&mut rng, &params);
-
-            assert!(a.add(&b, &params).add(&c, &params) == c.add(&a, &params).add(&b, &params));
-        }
-    }
-
-    #[test]
-    fn test_order() {
-        let rng = &mut XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = &JubjubParams::new();
-
-        // The neutral element is in the prime order subgroup.
-        assert!(Point::<Bls12, PrimeOrder>::zero().as_prime_order(params).is_some());
-
-        for _ in 0..50 {
-            // Pick a random point and multiply it by the cofactor
-            let base = Point::rand(rng, params).mul_by_cofactor(params);
-
-            // Any point multiplied by the cofactor will be in the prime
-            // order subgroup
-            assert!(base.as_prime_order(params).is_some());
-        }
-
-        // It's very likely that at least one out of 50 random points on the curve
-        // is not in the prime order subgroup.
-        let mut at_least_one_not_in_prime_order_subgroup = false;
-        for _ in 0..50 {
-            // Pick a random point.
-            let base = Point::rand(rng, params);
-
-            at_least_one_not_in_prime_order_subgroup |= base.as_prime_order(params).is_none();
-        }
-        assert!(at_least_one_not_in_prime_order_subgroup);
-    }
-
-    #[test]
-    fn test_mul_associativity() {
-        let rng = &mut XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = &JubjubParams::new();
-
-        for _ in 0..100 {
-            // Pick a random point and multiply it by the cofactor
-            let base = Point::rand(rng, params).mul_by_cofactor(params);
-
-            let mut a = Fs::rand(rng);
-            let b = Fs::rand(rng);
-            let c = Fs::rand(rng);
-
-            let res1 = base.mul(a, params).mul(b, params).mul(c, params);
-            let res2 = base.mul(b, params).mul(c, params).mul(a, params);
-            let res3 = base.mul(c, params).mul(a, params).mul(b, params);
-            a.mul_assign(&b);
-            a.mul_assign(&c);
-            let res4 = base.mul(a, params);
-
-            assert!(res1 == res2);
-            assert!(res2 == res3);
-            assert!(res3 == res4);
-
-            let (x, y) = res1.into_xy().unwrap();
-            assert!(is_on_curve(x, y, params));
-
-            let (x, y) = res2.into_xy().unwrap();
-            assert!(is_on_curve(x, y, params));
-
-            let (x, y) = res3.into_xy().unwrap();
-            assert!(is_on_curve(x, y, params));
-        }
-    }
-
-    #[test]
-    fn test_edwards_conversion() {
-        use super::edwards;
-
-        let rng = &mut XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = &JubjubParams::new();
-
-        for _ in 0..100 {
-            // compute base in edwards
-            let base = edwards::Point::rand(rng, params);
-
-            // sample random exponent
-            let exp = Fs::rand(rng);
-
-            // exponentiate in edwards
-            let mont_expected = Point::from_edwards(&base.mul(exp, params), params);
-
-            // convert to montgomery and exponentiate
-            let mont_exp = Point::from_edwards(&base, params).mul(exp, params);
-
-            assert!(mont_exp == mont_expected);
-
-            let (x, y) = mont_expected.into_xy().unwrap();
-            assert!(is_on_curve(x, y, params));
-        }
-    }
-
-    #[test]
-    fn test_back_and_forth() {
-        use super::edwards;
-
-        let rng = &mut XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
-        let params = &JubjubParams::new();
-
-        for _ in 0..100 {
-            // compute base in edwards
-            let base = edwards::Point::rand(rng, params);
-
-            // convert to montgomery
-            let base_mont = Point::from_edwards(&base, params);
-
-            {
-                let (x, y) = base_mont.into_xy().unwrap();
-                assert!(is_on_curve(x, y, params));
-            }
-
-            // convert back to edwards
-            let base_ed = edwards::Point::from_montgomery(&base_mont, params);
-
-            assert!(base == base_ed);
-        }
-    }
-
-    #[test]
-    fn test_low_order_points() {
-        use super::edwards;
-
-        let params = &JubjubParams::new();
-
-        let mut low_order_points: Vec<Point<Bls12, Unknown>> = vec![];
-
-        {
-            let mut push_point = |x, y| {
-                let x = Fr::from_str(x).unwrap();
-                let y = Fr::from_str(y).unwrap();
-
-                assert!(is_on_curve(x, y, params));
-
-                low_order_points.push(Point {
-                    x: x,
-                    y: y,
-                    infinity: false,
-                    _marker: PhantomData
-                });
-            };
-
-            // p is a point of order 8
-
-            // push p
-            push_point(
-                "26700795483254565448379661158233243896148151268643422869645920428793919977699",
-                "38240351061652197568958466618399906060451208175623222883988435386266133962140"
-            );
-
-            // push 2p
-            push_point(
-                "1",
-                "40876724960280933289965479552128619538703197557433544801868355907127087029496"
-            );
-
-            // push 3p
-            push_point(
-                "48853380121562139410032601262067414539517111118072400994428343856767649516850",
-                "32041076745907035847439769934443325418710075447471957144325987857573529479623"
-            );
-
-            // push 4p
-            push_point(
-                "0",
-                "0"
-            );
-
-            // push 5p
-            push_point(
-                "48853380121562139410032601262067414539517111118072400994428343856767649516850",
-                "20394798429219154632007970573742640418980477053055680678277670842365051704890"
-            );
-
-            // push 6p
-            push_point(
-                "1",
-                "11559150214845257189482260956057346298987354943094093020735302792811494155017"
-            );
-
-            // push 7p
-            push_point(
-                "26700795483254565448379661158233243896148151268643422869645920428793919977699",
-                "14195524113473992910489273889786059777239344324904414938615223313672447222373"
-            );
-        }
-
-        // push 8p (point at infinity)
-        low_order_points.push(Point::zero());
-
-        for point in &low_order_points {
-            let ed = edwards::Point::from_montgomery(point, params);
-            let mut ed_tmp = ed.clone();
-            let mut mont_tmp = point.clone();
-            for _ in 0..8 {
-                let mont_again = Point::from_edwards(&ed_tmp, params);
-                assert!(mont_again == mont_tmp);
-
-                let ed_again = edwards::Point::from_montgomery(&mont_tmp, params);
-                assert!(ed_again == ed_tmp);
-
-                ed_tmp = ed_tmp.add(&ed, params);
-                mont_tmp = mont_tmp.add(point, params);
-            }
-        }
-    }
-}
diff --git a/src/jubjub/tests.rs b/src/jubjub/tests.rs
new file mode 100644
index 0000000..d594383
--- /dev/null
+++ b/src/jubjub/tests.rs
@@ -0,0 +1,299 @@
+// TODO
+use super::*;
+
+use pairing::{
+    Field,
+    LegendreSymbol
+};
+
+use rand::{XorShiftRng, SeedableRng, Rand};
+
+pub fn test_suite<E: JubjubEngine>(params: &E::Params) {
+    test_back_and_forth::<E>(params);
+    test_jubjub_params::<E>(params);
+    test_rand::<E>(params);
+    test_identities::<E>(params);
+    test_addition_associativity::<E>(params);
+    test_order::<E>(params);
+    test_mul_associativity::<E>(params);
+    test_loworder::<E>(params);
+}
+
+fn is_on_mont_curve<E: JubjubEngine, P: JubjubParams<E>>(
+    x: E::Fr,
+    y: E::Fr,
+    params: &P
+) -> bool
+{
+    let mut lhs = y;
+    lhs.square();
+
+    let mut x2 = x;
+    x2.square();
+
+    let mut x3 = x2;
+    x3.mul_assign(&x);
+
+    let mut rhs = x2;
+    rhs.mul_assign(params.montgomery_a());
+    rhs.add_assign(&x);
+    rhs.add_assign(&x3);
+
+    lhs == rhs
+}
+
+fn is_on_twisted_edwards_curve<E: JubjubEngine, P: JubjubParams<E>>(
+    x: E::Fr,
+    y: E::Fr,
+    params: &P
+) -> bool
+{
+    let mut x2 = x;
+    x2.square();
+
+    let mut y2 = y;
+    y2.square();
+
+    // -x^2 + y^2
+    let mut lhs = y2;
+    lhs.sub_assign(&x2);
+
+    // 1 + d x^2 y^2
+    let mut rhs = y2;
+    rhs.mul_assign(&x2);
+    rhs.mul_assign(params.edwards_d());
+    rhs.add_assign(&E::Fr::one());
+
+    lhs == rhs
+}
+
+fn test_loworder<E: JubjubEngine>(params: &E::Params) {
+    let rng = &mut XorShiftRng::from_seed([0x3dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
+    let inf = montgomery::Point::zero();
+
+    // try to find a point of order 8
+    let p = loop {
+        let r = montgomery::Point::<E, _>::rand(rng, params).mul(E::Fs::char(), params);
+
+        let r2 = r.double(params);
+        let r4 = r2.double(params);
+        let r8 = r4.double(params);
+
+        if r2 != inf && r4 != inf && r8 == inf {
+            break r;
+        }
+    };
+
+    let mut loworder_points = vec![];
+    {
+        let mut tmp = p.clone();
+
+        for _ in 0..8 {
+            assert!(!loworder_points.contains(&tmp));
+            loworder_points.push(tmp.clone());
+            tmp = tmp.add(&p, params);
+        }
+    }
+    assert!(loworder_points[7] == inf);
+}
+
+fn test_mul_associativity<E: JubjubEngine>(params: &E::Params) {
+    use self::edwards::Point;
+    let rng = &mut XorShiftRng::from_seed([0x3dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
+
+    for _ in 0..100 {
+        // Pick a random point and multiply it by the cofactor
+        let base = Point::<E, _>::rand(rng, params).mul_by_cofactor(params);
+
+        let mut a = E::Fs::rand(rng);
+        let b = E::Fs::rand(rng);
+        let c = E::Fs::rand(rng);
+
+        let res1 = base.mul(a, params).mul(b, params).mul(c, params);
+        let res2 = base.mul(b, params).mul(c, params).mul(a, params);
+        let res3 = base.mul(c, params).mul(a, params).mul(b, params);
+        a.mul_assign(&b);
+        a.mul_assign(&c);
+        let res4 = base.mul(a, params);
+
+        assert!(res1 == res2);
+        assert!(res2 == res3);
+        assert!(res3 == res4);
+
+        let (x, y) = res1.into_xy();
+        assert!(is_on_twisted_edwards_curve(x, y, params));
+
+        let (x, y) = res2.into_xy();
+        assert!(is_on_twisted_edwards_curve(x, y, params));
+
+        let (x, y) = res3.into_xy();
+        assert!(is_on_twisted_edwards_curve(x, y, params));
+    }
+}
+
+fn test_order<E: JubjubEngine>(params: &E::Params) {
+    use self::edwards::Point;
+    let rng = &mut XorShiftRng::from_seed([0x3dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
+
+    // The neutral element is in the prime order subgroup.
+    assert!(Point::<E, PrimeOrder>::zero().as_prime_order(params).is_some());
+
+    for _ in 0..50 {
+        // Pick a random point and multiply it by the cofactor
+        let base = Point::<E, _>::rand(rng, params).mul_by_cofactor(params);
+
+        // Any point multiplied by the cofactor will be in the prime
+        // order subgroup
+        assert!(base.as_prime_order(params).is_some());
+    }
+
+    // It's very likely that at least one out of 50 random points on the curve
+    // is not in the prime order subgroup.
+    let mut at_least_one_not_in_prime_order_subgroup = false;
+    for _ in 0..50 {
+        // Pick a random point.
+        let base = Point::<E, _>::rand(rng, params);
+
+        at_least_one_not_in_prime_order_subgroup |= base.as_prime_order(params).is_none();
+    }
+    assert!(at_least_one_not_in_prime_order_subgroup);
+}
+
+fn test_addition_associativity<E: JubjubEngine>(params: &E::Params) {
+    let rng = &mut XorShiftRng::from_seed([0x3dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
+
+    for _ in 0..1000 {
+        use self::montgomery::Point;
+
+        let a = Point::<E, _>::rand(rng, params);
+        let b = Point::<E, _>::rand(rng, params);
+        let c = Point::<E, _>::rand(rng, params);
+
+        assert!(a.add(&b, &params).add(&c, &params) == c.add(&a, &params).add(&b, &params));
+    }
+
+    for _ in 0..1000 {
+        use self::edwards::Point;
+
+        let a = Point::<E, _>::rand(rng, params);
+        let b = Point::<E, _>::rand(rng, params);
+        let c = Point::<E, _>::rand(rng, params);
+
+        assert!(a.add(&b, &params).add(&c, &params) == c.add(&a, &params).add(&b, &params));
+    }
+}
+
+fn test_identities<E: JubjubEngine>(params: &E::Params) {
+    let rng = &mut XorShiftRng::from_seed([0x3dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
+
+    {
+        use self::edwards::Point;
+
+        let z = Point::<E, PrimeOrder>::zero();
+        assert!(z.double(&params) == z);
+        assert!(z.negate() == z);
+
+        for _ in 0..100 {
+            let r = Point::<E, _>::rand(rng, params);
+
+            assert!(r.add(&Point::zero(), &params) == r);
+            assert!(r.add(&r.negate(), &params) == Point::zero());
+        }
+    }
+
+    {
+        use self::montgomery::Point;
+
+        let z = Point::<E, PrimeOrder>::zero();
+        assert!(z.double(&params) == z);
+        assert!(z.negate() == z);
+
+        for _ in 0..100 {
+            let r = Point::<E, _>::rand(rng, params);
+
+            assert!(r.add(&Point::zero(), &params) == r);
+            assert!(r.add(&r.negate(), &params) == Point::zero());
+        }
+    }
+}
+
+fn test_rand<E: JubjubEngine>(params: &E::Params) {
+    let rng = &mut XorShiftRng::from_seed([0x3dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
+
+    for _ in 0..1000 {
+        let p = montgomery::Point::<E, _>::rand(rng, params);
+        let e = edwards::Point::<E, _>::rand(rng, params);
+
+        {
+            let (x, y) = p.into_xy().unwrap();
+            assert!(is_on_mont_curve(x, y, params));
+        }
+
+        {
+            let (x, y) = e.into_xy();
+            assert!(is_on_twisted_edwards_curve(x, y, params));
+        }
+    }
+}
+
+fn test_back_and_forth<E: JubjubEngine>(params: &E::Params) {
+    let rng = &mut XorShiftRng::from_seed([0x5dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
+
+    for _ in 0..1000 {
+        let s = E::Fs::rand(rng);
+        let edwards_p1 = edwards::Point::<E, _>::rand(rng, params);
+        let mont_p1 = montgomery::Point::from_edwards(&edwards_p1, params);
+        let mont_p2 = montgomery::Point::<E, _>::rand(rng, params);
+        let edwards_p2 = edwards::Point::from_montgomery(&mont_p2, params);
+
+        let mont = mont_p1.add(&mont_p2, params).mul(s, params);
+        let edwards = edwards_p1.add(&edwards_p2, params).mul(s, params);
+
+        assert!(
+            montgomery::Point::from_edwards(&edwards, params) == mont
+        );
+
+        assert!(
+            edwards::Point::from_montgomery(&mont, params) == edwards
+        );
+    }
+}
+
+fn test_jubjub_params<E: JubjubEngine>(params: &E::Params) {
+    // a = -1
+    let mut a = E::Fr::one();
+    a.negate();
+
+    {
+        // The twisted Edwards addition law is complete when d is nonsquare
+        // and a is square.
+
+        assert!(params.edwards_d().legendre() == LegendreSymbol::QuadraticNonResidue);
+        assert!(a.legendre() == LegendreSymbol::QuadraticResidue);
+    }
+
+    {
+        // Check that A^2 - 4 is nonsquare:
+        let mut tmp = params.montgomery_a().clone();
+        tmp.square();
+        tmp.sub_assign(&E::Fr::from_str("4").unwrap());
+        assert!(tmp.legendre() == LegendreSymbol::QuadraticNonResidue);
+    }
+
+    {
+        // Check that A - 2 is nonsquare:
+        let mut tmp = params.montgomery_a().clone();
+        tmp.sub_assign(&E::Fr::from_str("2").unwrap());
+        assert!(tmp.legendre() == LegendreSymbol::QuadraticNonResidue);
+    }
+
+    {
+        // Check the validity of the scaling factor
+        let mut tmp = a;
+        tmp.sub_assign(&params.edwards_d());
+        tmp = tmp.inverse().unwrap();
+        tmp.mul_assign(&E::Fr::from_str("4").unwrap());
+        tmp = tmp.sqrt().unwrap();
+        assert_eq!(&tmp, params.scale());
+    }
+}