crypto_bigint/uint/
sqrt.rs

Help
1//! [`UInt`] square root operations.
2
3use super::UInt;
4use crate::{Limb, Word};
5use subtle::{ConstantTimeEq, CtOption};
6
7impl<const LIMBS: usize> UInt<LIMBS> {
8    /// Computes √(`self`)
9    /// Uses Brent & Zimmermann, Modern Computer Arithmetic, v0.5.9, Algorithm 1.13
10    ///
11    /// Callers can check if `self` is a square by squaring the result
12    pub const fn sqrt(&self) -> Self {
13        let max_bits = (self.bits_vartime() + 1) >> 1;
14        let cap = Self::ONE.shl_vartime(max_bits);
15        let mut guess = cap; // ≥ √(`self`)
16        let mut xn = {
17            let q = self.wrapping_div(&guess);
18            let t = guess.wrapping_add(&q);
19            t.shr_vartime(1)
20        };
21
22        // If guess increased, the initial guess was low.
23        // Repeat until reverse course.
24        while guess.ct_cmp(&xn) == -1 {
25            // Sometimes an increase is too far, especially with large
26            // powers, and then takes a long time to walk back.  The upper
27            // bound is based on bit size, so saturate on that.
28            let res = Limb::ct_cmp(Limb(xn.bits_vartime() as Word), Limb(max_bits as Word)) - 1;
29            let le = Limb::is_nonzero(Limb(res as Word));
30            guess = Self::ct_select(cap, xn, le);
31            xn = {
32                let q = self.wrapping_div(&guess);
33                let t = guess.wrapping_add(&q);
34                t.shr_vartime(1)
35            };
36        }
37
38        // Repeat while guess decreases.
39        while guess.ct_cmp(&xn) == 1 && xn.ct_is_nonzero() == Word::MAX {
40            guess = xn;
41            xn = {
42                let q = self.wrapping_div(&guess);
43                let t = guess.wrapping_add(&q);
44                t.shr_vartime(1)
45            };
46        }
47
48        Self::ct_select(Self::ZERO, guess, self.ct_is_nonzero())
49    }
50
51    /// Wrapped sqrt is just normal √(`self`)
52    /// There’s no way wrapping could ever happen.
53    /// This function exists, so that all operations are accounted for in the wrapping operations.
54    pub const fn wrapping_sqrt(&self) -> Self {
55        self.sqrt()
56    }
57
58    /// Perform checked sqrt, returning a [`CtOption`] which `is_some`
59    /// only if the √(`self`)² == self
60    pub fn checked_sqrt(&self) -> CtOption<Self> {
61        let r = self.sqrt();
62        let s = r.wrapping_mul(&r);
63        CtOption::new(r, self.ct_eq(&s))
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use crate::{Limb, U256};
70
71    #[cfg(feature = "rand")]
72    use {
73        crate::{CheckedMul, Random, U512},
74        rand_chacha::ChaChaRng,
75        rand_core::{RngCore, SeedableRng},
76    };
77
78    #[test]
79    fn edge() {
80        assert_eq!(U256::ZERO.sqrt(), U256::ZERO);
81        assert_eq!(U256::ONE.sqrt(), U256::ONE);
82        let mut half = U256::ZERO;
83        for i in 0..half.limbs.len() / 2 {
84            half.limbs[i] = Limb::MAX;
85        }
86        assert_eq!(U256::MAX.sqrt(), half,);
87    }
88
89    #[test]
90    fn simple() {
91        let tests = [
92            (4u8, 2u8),
93            (9, 3),
94            (16, 4),
95            (25, 5),
96            (36, 6),
97            (49, 7),
98            (64, 8),
99            (81, 9),
100            (100, 10),
101            (121, 11),
102            (144, 12),
103            (169, 13),
104        ];
105        for (a, e) in &tests {
106            let l = U256::from(*a);
107            let r = U256::from(*e);
108            assert_eq!(l.sqrt(), r);
109            assert_eq!(l.checked_sqrt().is_some().unwrap_u8(), 1u8);
110        }
111    }
112
113    #[test]
114    fn nonsquares() {
115        assert_eq!(U256::from(2u8).sqrt(), U256::from(1u8));
116        assert_eq!(U256::from(2u8).checked_sqrt().is_some().unwrap_u8(), 0);
117        assert_eq!(U256::from(3u8).sqrt(), U256::from(1u8));
118        assert_eq!(U256::from(3u8).checked_sqrt().is_some().unwrap_u8(), 0);
119        assert_eq!(U256::from(5u8).sqrt(), U256::from(2u8));
120        assert_eq!(U256::from(6u8).sqrt(), U256::from(2u8));
121        assert_eq!(U256::from(7u8).sqrt(), U256::from(2u8));
122        assert_eq!(U256::from(8u8).sqrt(), U256::from(2u8));
123        assert_eq!(U256::from(10u8).sqrt(), U256::from(3u8));
124    }
125
126    #[cfg(feature = "rand")]
127    #[test]
128    fn fuzz() {
129        let mut rng = ChaChaRng::from_seed([7u8; 32]);
130        for _ in 0..50 {
131            let t = rng.next_u32() as u64;
132            let s = U256::from(t);
133            let s2 = s.checked_mul(&s).unwrap();
134            assert_eq!(s2.sqrt(), s);
135            assert_eq!(s2.checked_sqrt().is_some().unwrap_u8(), 1);
136        }
137
138        for _ in 0..50 {
139            let s = U256::random(&mut rng);
140            let mut s2 = U512::ZERO;
141            s2.limbs[..s.limbs.len()].copy_from_slice(&s.limbs);
142            assert_eq!(s.square().sqrt(), s2);
143        }
144    }
145}
crypto_bigint/uint/sqrt.rs

crypto_bigint/uint/
sqrt.rs
