decaf377/ark_curve/
invsqrt.rs

1use core::convert::TryInto;
2use hashbrown::HashMap;
3
4use crate::Fq;
5use ark_ff::{BigInteger256, BigInteger64, Field, Zero};
6use ark_std::boxed::Box;
7use ark_std::vec::Vec;
8use once_cell::sync::Lazy;
9
10use crate::ark_curve::constants::{
11    G, M_MINUS_ONE_DIV_TWO, N, ONE, SQRT_W, ZETA_TO_ONE_MINUS_M_DIV_TWO,
12};
13
14struct SquareRootTables {
15    pub s_lookup: HashMap<Fq, u64>,
16    pub nonsquare_lookup: [Fq; 2],
17    pub g0: Box<[Fq; 256]>,
18    pub g8: Box<[Fq; 256]>,
19    pub g16: Box<[Fq; 256]>,
20    pub g24: Box<[Fq; 256]>,
21    pub g32: Box<[Fq; 256]>,
22    pub g40: Box<[Fq; 256]>,
23}
24
25impl SquareRootTables {
26    fn new() -> Self {
27        let mut s_lookup = HashMap::new();
28        for nu in 0..256 {
29            // These entries should be g**(-1 * nu * 2**(n-w)) so:
30            // 1. We compute g**(nu * 2**(n-w))
31            // 2. Then take the inverse of the quantity from step 1
32            let exp: BigInteger256 = (nu * 2u64.pow(N - SQRT_W)).into();
33            let g_inv = G.pow(exp);
34            s_lookup.insert(
35                g_inv.inverse().expect("inverse exists for these elements"),
36                nu,
37            );
38        }
39
40        let powers_of_two = [0, 8, 16, 24, 32, 40];
41        let mut gtab = Vec::new();
42        for power_of_two in powers_of_two {
43            let mut gtab_i = Vec::<Fq>::new();
44            for nu in 0..256 {
45                let exp: BigInteger256 = (nu * 2u64.pow(power_of_two)).into();
46                gtab_i.push(G.pow(exp))
47            }
48            gtab.push(gtab_i);
49        }
50
51        let nonsquare_lookup = [*ONE, *ZETA_TO_ONE_MINUS_M_DIV_TWO];
52
53        Self {
54            s_lookup,
55            nonsquare_lookup,
56            g40: gtab.pop().unwrap().into_boxed_slice().try_into().unwrap(),
57            g32: gtab.pop().unwrap().into_boxed_slice().try_into().unwrap(),
58            g24: gtab.pop().unwrap().into_boxed_slice().try_into().unwrap(),
59            g16: gtab.pop().unwrap().into_boxed_slice().try_into().unwrap(),
60            g8: gtab.pop().unwrap().into_boxed_slice().try_into().unwrap(),
61            g0: gtab.pop().unwrap().into_boxed_slice().try_into().unwrap(),
62        }
63    }
64}
65
66static SQRT_LOOKUP_TABLES: Lazy<SquareRootTables> = Lazy::new(|| SquareRootTables::new());
67
68impl Fq {
69    /// Computes the square root of a ratio of field elements, returning:
70    ///
71    /// - `(true, sqrt(num/den))` if `num` and `den` are both nonzero and `num/den` is square;
72    /// - `(true, 0)` if `num` is zero;
73    /// - `(false, 0)` if `den` is zero;
74    /// - `(false, sqrt(zeta*num/den))` if `num` and `den` are both nonzero and `num/den` is nonsquare;
75    pub fn sqrt_ratio_zeta(num: &Self, den: &Self) -> (bool, Self) {
76        // This square root method is based on:
77        // * [Sarkar2020](https://eprint.iacr.org/2020/1407)
78        // * [Zcash Pasta](https://github.com/zcash/pasta_curves)
79        //
80        // See the Penumbra protocol specification for details.
81        if num.is_zero() {
82            return (true, *num);
83        }
84        if den.is_zero() {
85            return (false, *den);
86        }
87
88        let s_exp: BigInteger256 = (2u64.pow(N) - 1).into();
89        let s = den.pow(s_exp);
90        let t = s.square() * den;
91        let w = (*num * t).pow(*M_MINUS_ONE_DIV_TWO) * s;
92
93        let v = w * den;
94        let uv = w * num;
95
96        // x = u * v^2 = x5
97        let x5 = uv * v;
98        let pow2_8: BigInteger64 = 2u64.pow(8).into();
99
100        // x4 = x5^{2^8}
101        let x4 = x5.pow(pow2_8);
102        // x3 = x4^{2^8}
103        let x3 = x4.pow(pow2_8);
104        // x2 = x3^{2^8}
105        let x2 = x3.pow(pow2_8);
106        // x1 = x2^{2^8}
107        let x1 = x2.pow(pow2_8);
108        let pow2_7: BigInteger64 = 2u64.pow(7).into();
109        // x0 = x1^{2^7}
110        let x0 = x1.pow(pow2_7);
111
112        // i = 0
113        let q0_prime = SQRT_LOOKUP_TABLES.s_lookup[&x0];
114        let mut t = q0_prime;
115
116        // i = 1
117        let alpha_1 = x1 * SQRT_LOOKUP_TABLES.g32[(t & 0xFF) as usize];
118        let q1_prime = SQRT_LOOKUP_TABLES.s_lookup[&alpha_1];
119        t += q1_prime << 7;
120
121        // i = 2
122        let alpha_2 = x2
123            * SQRT_LOOKUP_TABLES.g24[(t & 0xFF) as usize]
124            * SQRT_LOOKUP_TABLES.g32[((t >> 8) & 0xFF) as usize];
125        let q2 = SQRT_LOOKUP_TABLES.s_lookup[&alpha_2];
126        t += q2 << 15;
127
128        // i = 3
129        let alpha_3 = x3
130            * SQRT_LOOKUP_TABLES.g16[(t & 0xFF) as usize]
131            * SQRT_LOOKUP_TABLES.g24[((t >> 8) & 0xFF) as usize]
132            * SQRT_LOOKUP_TABLES.g32[((t >> 16) & 0xFF) as usize];
133        let q3 = SQRT_LOOKUP_TABLES.s_lookup[&alpha_3];
134        t += q3 << 23;
135
136        // i = 4
137        let alpha_4 = x4
138            * SQRT_LOOKUP_TABLES.g8[(t & 0xFF) as usize]
139            * SQRT_LOOKUP_TABLES.g16[((t >> 8) & 0xFF) as usize]
140            * SQRT_LOOKUP_TABLES.g24[((t >> 16) & 0xFF) as usize]
141            * SQRT_LOOKUP_TABLES.g32[((t >> 24) & 0xFF) as usize];
142        let q4 = SQRT_LOOKUP_TABLES.s_lookup[&alpha_4];
143        t += q4 << 31;
144
145        // i = 5
146        let alpha_5 = x5
147            * SQRT_LOOKUP_TABLES.g0[(t & 0xFF) as usize]
148            * SQRT_LOOKUP_TABLES.g8[((t >> 8) & 0xFF) as usize]
149            * SQRT_LOOKUP_TABLES.g16[((t >> 16) & 0xFF) as usize]
150            * SQRT_LOOKUP_TABLES.g24[((t >> 24) & 0xFF) as usize]
151            * SQRT_LOOKUP_TABLES.g32[((t >> 32) & 0xFF) as usize];
152        let q5 = SQRT_LOOKUP_TABLES.s_lookup[&alpha_5];
153        t += q5 << 39;
154
155        t = (t + 1) >> 1;
156        let res: Fq = uv
157            * SQRT_LOOKUP_TABLES.nonsquare_lookup[(q0_prime & 0b1) as usize]
158            * SQRT_LOOKUP_TABLES.g0[(t & 0xFF) as usize]
159            * SQRT_LOOKUP_TABLES.g8[((t >> 8) & 0xFF) as usize]
160            * SQRT_LOOKUP_TABLES.g16[((t >> 16) & 0xFF) as usize]
161            * SQRT_LOOKUP_TABLES.g24[((t >> 24) & 0xFF) as usize]
162            * SQRT_LOOKUP_TABLES.g32[((t >> 32) & 0xFF) as usize]
163            * SQRT_LOOKUP_TABLES.g40[((t >> 40) & 0xFF) as usize];
164
165        ((q0_prime & 0b1) == 0, res)
166    }
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172    use crate::ark_curve::constants::ZETA;
173
174    use proptest::prelude::*;
175
176    fn fq_strategy() -> BoxedStrategy<Fq> {
177        any::<[u8; 32]>()
178            .prop_map(|bytes| Fq::from_le_bytes_mod_order(&bytes[..]))
179            .boxed()
180    }
181
182    proptest! {
183        #![proptest_config(ProptestConfig::with_cases(10000))]
184        #[test]
185        fn sqrt_ratio_zeta(u in fq_strategy(), v in fq_strategy()) {
186            if u.is_zero() {
187                assert_eq!(Fq::sqrt_ratio_zeta(&u, &v), (true, u));
188            } else if v.is_zero() {
189                assert_eq!(Fq::sqrt_ratio_zeta(&u, &v), (false, v));
190            } else {
191                let (was_square, sqrt_zeta_uv) = Fq::sqrt_ratio_zeta(&u, &v);
192                let zeta_uv = sqrt_zeta_uv * sqrt_zeta_uv;
193                if was_square {
194                    // check zeta_uv = u/v
195                    assert_eq!(u, v * zeta_uv);
196                } else {
197                    // check zeta_uv = zeta * u / v
198                    assert_eq!(ZETA * u, v * zeta_uv);
199                }
200            }
201        }
202    }
203
204    #[test]
205    fn sqrt_ratio_edge_cases() {
206        // u = 0
207        assert_eq!(Fq::sqrt_ratio_zeta(&Fq::zero(), &ONE), (true, Fq::zero()));
208
209        // v = 0
210        assert_eq!(Fq::sqrt_ratio_zeta(&ONE, &Fq::zero()), (false, Fq::zero()));
211    }
212}