Merge pull request #196 from str4d/pow-fixed

Speed up Field::invert and SqrtField::sqrt in ff_derive with addition chains

Cargo.lock

@@ -1,5 +1,15 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
+[[package]]
+name = "addchain"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+dependencies = [
+ "num-bigint 0.2.3 (registry+https://github.com/rust-lang/crates.io-index)",
+ "num-integer 0.1.41 (registry+https://github.com/rust-lang/crates.io-index)",
+ "num-traits 0.2.8 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
 [[package]]
 name = "aes"
 version = "0.3.2"
@@ -467,6 +477,7 @@ dependencies = [
 name = "ff_derive"
 version = "0.6.0"
 dependencies = [
+ "addchain 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "num-bigint 0.2.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "num-integer 0.1.41 (registry+https://github.com/rust-lang/crates.io-index)",
  "num-traits 0.2.8 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -1298,6 +1309,7 @@ dependencies = [
 ]
 
 [metadata]
+"checksum addchain 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "1177222c93a7bb492002e9a3cd947c7fd869e085d6e81a9e415ff1be65b3489c"
 "checksum aes 0.3.2 (registry+https://github.com/rust-lang/crates.io-index)" = "54eb1d8fe354e5fc611daf4f2ea97dd45a765f4f1e4512306ec183ae2e8f20c9"
 "checksum aes-soft 0.3.3 (registry+https://github.com/rust-lang/crates.io-index)" = "cfd7e7ae3f9a1fb5c03b389fc6bb9a51400d0c13053f0dca698c832bfd893a0d"
 "checksum aesni 0.6.0 (registry+https://github.com/rust-lang/crates.io-index)" = "2f70a6b5f971e473091ab7cfb5ffac6cde81666c4556751d8d5620ead8abf100"

ff/ff_derive/Cargo.toml

@@ -1,7 +1,10 @@
 [package]
 name = "ff_derive"
 version = "0.6.0"
-authors = ["Sean Bowe <ewillbefull@gmail.com>"]
+authors = [
+    "Sean Bowe <ewillbefull@gmail.com>",
+    "Jack Grigg <thestr4d@gmail.com>",
+]
 description = "Procedural macro library used to build custom prime field implementations"
 documentation = "https://docs.rs/ff/"
 homepage = "https://github.com/ebfull/ff"
@@ -13,6 +16,7 @@ edition = "2018"
 proc-macro = true
 
 [dependencies]
+addchain = "0.1"
 num-bigint = "0.2"
 num-traits = "0.2"
 num-integer = "0.1"

ff/ff_derive/src/lib.rs

@@ -10,6 +10,8 @@ use quote::quote;
 use quote::TokenStreamExt;
 use std::str::FromStr;
 
+mod pow_fixed;
+
 #[proc_macro_derive(PrimeField, attributes(PrimeFieldModulus, PrimeFieldGenerator))]
 pub fn prime_field(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     // Parse the type definition
@@ -407,8 +409,11 @@ fn prime_field_constants_and_sqrt(
     let sqrt_impl = if (modulus % BigUint::from_str("4").unwrap())
         == BigUint::from_str("3").unwrap()
     {
-        let mod_plus_1_over_4 =
+        // Addition chain for (r + 1) // 4
-            biguint_to_u64_vec((modulus + BigUint::from_str("1").unwrap()) >> 2, limbs);
+        let mod_plus_1_over_4 = pow_fixed::generate(
+            &quote! {self},
+            (modulus + BigUint::from_str("1").unwrap()) >> 2,
+        );
 
         quote! {
             impl ::ff::SqrtField for #name {
@@ -418,7 +423,9 @@ fn prime_field_constants_and_sqrt(
                     // Because r = 3 (mod 4)
                     // sqrt can be done with only one exponentiation,
                     // via the computation of  self^((r + 1) // 4) (mod r)
-                    let sqrt = self.pow_vartime(#mod_plus_1_over_4);
+                    let sqrt = {
+                        #mod_plus_1_over_4
+                    };
 
                     ::subtle::CtOption::new(
                         sqrt,
@@ -428,7 +435,8 @@ fn prime_field_constants_and_sqrt(
             }
         }
     } else if (modulus % BigUint::from_str("16").unwrap()) == BigUint::from_str("1").unwrap() {
-        let t_minus_1_over_2 = biguint_to_u64_vec((&t - BigUint::one()) >> 1, limbs);
+        // Addition chain for (t - 1) // 2
+        let t_minus_1_over_2 = pow_fixed::generate(&quote! {self}, (&t - BigUint::one()) >> 1);
 
         quote! {
             impl ::ff::SqrtField for #name {
@@ -438,7 +446,9 @@ fn prime_field_constants_and_sqrt(
                     use ::subtle::{ConditionallySelectable, ConstantTimeEq};
 
                     // w = self^((t - 1) // 2)
-                    let w = self.pow_vartime(#t_minus_1_over_2);
+                    let w = {
+                        #t_minus_1_over_2
+                    };
 
                     let mut v = S;
                     let mut x = *self * &w;
@@ -742,11 +752,10 @@ fn prime_field_impl(
         a: proc_macro2::TokenStream,
         name: &syn::Ident,
         modulus: &BigUint,
-        limbs: usize,
     ) -> proc_macro2::TokenStream {
-        let mod_minus_2 = biguint_to_u64_vec(modulus - BigUint::from(2u64), limbs);
+        // Addition chain for p - 2
+        let mod_minus_2 = pow_fixed::generate(&a, modulus - BigUint::from(2u64));
 
-        // TODO: Improve on this by computing an addition chain for mod_minus_two
         quote! {
             use ::subtle::ConstantTimeEq;
 
@@ -756,7 +765,9 @@ fn prime_field_impl(
             // `ff_derive` requires that `p` is prime; in this case, `phi(p) = p - 1`, and
             // thus:
             //     a^-1 ≡ a^(p - 2) mod p
-            let inv = #a.pow_vartime(#mod_minus_2);
+            let inv = {
+                #mod_minus_2
+            };
 
             ::subtle::CtOption::new(inv, !#a.ct_eq(&#name::zero()))
         }
@@ -764,7 +775,7 @@ fn prime_field_impl(
 
     let squaring_impl = sqr_impl(quote! {self}, limbs);
     let multiply_impl = mul_impl(quote! {self}, quote! {other}, limbs);
-    let invert_impl = inv_impl(quote! {self}, name, modulus, limbs);
+    let invert_impl = inv_impl(quote! {self}, name, modulus);
     let montgomery_impl = mont_impl(limbs);
 
     // (self.0).0[0].ct_eq(&(other.0).0[0]) & (self.0).0[1].ct_eq(&(other.0).0[1]) & ...

ff/ff_derive/src/pow_fixed.rs

@@ -0,0 +1,56 @@
+//! Fixed-exponent variable-base exponentiation using addition chains.
+
+use addchain::{build_addition_chain, Step};
+use num_bigint::BigUint;
+use quote::quote;
+use syn::Ident;
+
+/// Returns t{n} as an ident.
+fn get_temp(n: usize) -> Ident {
+    Ident::new(&format!("t{}", n), proc_macro2::Span::call_site())
+}
+
+pub(crate) fn generate(
+    base: &proc_macro2::TokenStream,
+    exponent: BigUint,
+) -> proc_macro2::TokenStream {
+    let steps = build_addition_chain(exponent);
+
+    let mut gen = proc_macro2::TokenStream::new();
+
+    // First entry in chain is one, i.e. the base.
+    let start = get_temp(0);
+    gen.extend(quote! {
+        let #start = #base;
+    });
+
+    let mut tmps = vec![start];
+    for (i, step) in steps.into_iter().enumerate() {
+        let out = get_temp(i + 1);
+
+        gen.extend(match step {
+            Step::Double { index } => {
+                let val = &tmps[index];
+                quote! {
+                    let #out = #val.square();
+                }
+            }
+            Step::Add { left, right } => {
+                let left = &tmps[left];
+                let right = &tmps[right];
+                quote! {
+                    let #out = #left * #right;
+                }
+            }
+        });
+
+        tmps.push(out.clone());
+    }
+
+    let end = tmps.last().expect("have last");
+    gen.extend(quote! {
+        #end
+    });
+
+    gen
+}