crypto_bigint/uint/ref_type/
sqrt.rs

1//! Square root calculation.
2
3use super::UintRef;
4use crate::{Choice, Limb, NonZero, Uint, WideWord, Word, word};
5
6impl UintRef {
7    #[inline]
8    pub(crate) const fn sqrt_assign(&mut self, buf: (&mut UintRef, &mut UintRef)) -> Choice {
9        let len = self.nlimbs();
10        assert!(buf.0.nlimbs() >= len && buf.1.nlimbs() >= len);
11
12        match len {
13            0 => Choice::TRUE,
14            1 => {
15                let rem;
16                (self.limbs[0].0, rem) = sqrt_rem_word(self.limbs[0].0);
17                Limb(rem).is_zero()
18            }
19            _ => {
20                let mut bits = self.bits();
21                let is_zero = Choice::from_u32_nz(bits).not();
22                // Set first limb to 1 if the input is zero
23                self.limbs[0] = Limb::select(self.limbs[0], Limb::ONE, is_zero);
24                bits = is_zero.select_u32(bits, 1);
25
26                // Shift such that at least one of the top two bits are non-zero
27                let s_shift = (self.bits_precision() - bits) >> 1;
28                self.shl_assign(s_shift * 2);
29
30                // Compute root and uncorrected remainder, placing them in buf.0 and buf.1
31                self.sqrt_rem_normalized::<false>(buf.0, buf.1);
32
33                // Copy the shifted root to self and correct it
34                self.copy_from(buf.0.leading(self.nlimbs()));
35                self.shr_assign(s_shift);
36
37                // Set root to zero if self was zero
38                self.limbs[0] = Limb::select(self.limbs[0], Limb::ZERO, is_zero);
39
40                // Check if there is a remainder
41                buf.1.is_zero()
42            }
43        }
44    }
45
46    #[inline]
47    pub(crate) const fn sqrt_assign_vartime(&mut self, buf: (&mut UintRef, &mut UintRef)) -> bool {
48        let bits = self.bits_vartime();
49        if bits == 0 {
50            return true;
51        }
52
53        let words = bits.div_ceil(Limb::BITS);
54        assert!(buf.0.nlimbs() >= words as usize && buf.1.nlimbs() >= words as usize);
55
56        // Only consider the populated limbs
57        let out = self.leading_mut(words as usize);
58
59        if words <= 2 {
60            // No shifts needed
61            sqrt_rem_small::<true>(out.as_limbs(), buf.0.as_mut_limbs(), buf.1.as_mut_limbs());
62            out.copy_from(buf.0.leading(out.nlimbs()));
63        } else {
64            // Shift such that at least one of the top two bits are non-zero
65            let lz = words * Limb::BITS - bits;
66            let s_shift = lz >> 1;
67            out.shl_assign_limb_vartime(s_shift * 2);
68
69            // Compute root and uncorrected remainder, placing them in buf.0 and buf.1
70            out.sqrt_rem_normalized::<true>(buf.0, buf.1);
71
72            // Copy the shifted root to self and correct it
73            out.copy_from(buf.0.leading(out.nlimbs()));
74            out.shr_assign_limb_vartime(s_shift);
75        }
76
77        // Check if there is a remainder
78        buf.1.is_zero_vartime()
79    }
80
81    /// Corresponds to Brent & Zimmermann, Modern Computer Arithmetic, v0.5.9, Algorithm 1.12
82    #[allow(clippy::cast_possible_truncation)]
83    const fn sqrt_rem_normalized<const VARTIME: bool>(&self, s: &mut UintRef, r: &mut UintRef) {
84        let len = self.nlimbs();
85
86        // Handle base case of a square root < 4 limbs
87        // Unlike the source material, we do not handle 4 limbs in the base case sqrt function
88        if len < 4 {
89            sqrt_rem_small::<VARTIME>(self.as_limbs(), s.as_mut_limbs(), r.as_mut_limbs());
90            return;
91        }
92
93        let l = len >> 2;
94        let rt_len = (len - l * 2).div_ceil(2);
95        let (a_lo, a_hi) = self.split_at(l * 2);
96        let (a0, a1) = a_lo.split_at(l);
97
98        // (s', r') = SqrtRem(a3•B + a2), leaving `l` spare low limbs in each
99        a_hi.sqrt_rem_normalized::<VARTIME>(s.trailing_mut(l), r.trailing_mut(l));
100
101        // Split up buffers
102        let (s, s_tail) = s.split_at_mut(l + rt_len);
103        let (s_lo, s_hi) = s.split_at_mut(l);
104        let (r, r_tail) = r.split_at_mut(l + rt_len + 1);
105
106        // Set r = r'B + a1
107        r.leading_mut(l).copy_from(a1);
108        // Set u = s'
109        let u = s_tail.leading_mut(rt_len);
110        u.copy_from(s_hi);
111
112        // We wish to divide r/2s', setting r to the quotient and u to the remainder, but
113        // 2s' doesn't fit within the available buffer. Compute r/s' and adjust the result
114        r.div_rem(u);
115        // Adjust the result for divisor 2s'
116        let r_mod2 = r.shr1_assign();
117        let u_hi = u.conditional_add_assign(s_hi, Limb::ZERO, r_mod2.is_nonzero());
118
119        let (r_lo, r_hi) = r.split_at_mut(l);
120
121        // s = s'B + q
122        s_lo.copy_from(r_lo);
123        let q_hi = r_hi.limbs[0];
124        s_hi.add_assign_limb(q_hi);
125        debug_assert!(q_hi.0 & !1 == 0);
126
127        // r = uB + a0
128        r_lo.copy_from(a0);
129        r_hi.leading_mut(rt_len).copy_from(u);
130        r_hi.limbs[rt_len] = u_hi;
131
132        // Compute q^2
133        let q2 = s_tail.leading_mut(l * 2);
134        q2.fill(Limb::ZERO);
135        s_lo.wrapping_square(q2);
136
137        // r -= q^2, producing a borrow if r < 0
138        let (r0, r1) = r.split_at_mut(l * 2);
139        let borrow = r0.borrowing_sub_assign(q2, Limb::ZERO);
140        let swap = r1.borrowing_sub_assign_limb(q_hi, borrow).lsb_to_choice();
141        let s_carry = Limb::select(Limb::ZERO, Limb::ONE, swap);
142
143        // s -= 1 if r < 0
144        s.borrowing_sub_assign_limb(s_carry, Limb::ZERO);
145
146        // r += 2s + 1 if r < 0
147        let s_mul = s_carry.shl(1);
148        r.carrying_add_assign_mul_limb(s, s_mul, s_carry);
149
150        // Clear upper limbs of the result buffers
151        s_tail.fill(Limb::ZERO);
152        r_tail.fill(Limb::ZERO);
153    }
154}
155
156/// Compute the square root and remainder of a [`Word`].
157///
158/// Adapted from Hacker's Delight 2E by Henry S. Warren Jr.,
159/// Fig. 11-4, "Integer square root, hardware algorithm"
160/// based on Toepler's algorithm.
161#[inline]
162const fn sqrt_rem_word(value: Word) -> (Word, Word) {
163    let mut m = 1 << (Word::BITS - 2);
164    let mut x = value;
165    let mut y = 0;
166    while m != 0 {
167        let b = y | m;
168        y >>= 1;
169        let mask = word::choice_to_mask(word::choice_from_le(b, x));
170        x = x.wrapping_sub(b & mask);
171        y |= m & mask;
172        m >>= 2;
173    }
174    (y, x)
175}
176
177/// Compute the root and remainder for `n+1` limbs, given the root and remainder
178/// for `n` limbs (n=1 or 2). The input value must be normalized such that at least one of the
179/// top two bits is set, producing an `s1` with at least `Word::BITS/2` bits.
180///
181/// This is essentially a specialized version of the root expansion portion of `sqrt_rem_normalized`.
182#[allow(clippy::cast_possible_truncation)]
183#[inline]
184const fn sqrt_rem_expand_by_word(s1: Word, r1: WideWord, next: Word) -> (WideWord, WideWord) {
185    const HALF_WIDTH: u32 = Word::BITS >> 1;
186    debug_assert!((s1 >> (HALF_WIDTH - 1)) > 0);
187
188    // Split the lower word into lower and upper halves
189    let (a0, a1) = (
190        (next & ((1 << HALF_WIDTH) - 1)) as WideWord,
191        (next >> HALF_WIDTH) as WideWord,
192    );
193    let d = (r1 << HALF_WIDTH) | a1;
194
195    // Divide by (r1B + a1) by s1 (not 2s1 which could overflow a limb)
196    let (q, _) = Uint::<2>::from_wide_word(d).div_rem_limb(NonZero::<Limb>::new_unwrap(Limb(s1)));
197    // Correct the quotient
198    let q = (q.limbs[0].0 >> 1) as WideWord;
199    // Recompute the remainder u
200    let u = d - (q << 1) * s1 as WideWord;
201    // Set s = s1B + q
202    let s = ((s1 as WideWord) << HALF_WIDTH) + q;
203
204    // Set r' = uB + a0
205    let r_pre = (u << HALF_WIDTH) | a0;
206    let q2 = q.pow(2);
207    let swap = word::choice_to_wide_mask(word::choice_from_wide_lt(r_pre, q2));
208
209    // Subtract 1 from s if r' < q^2
210    let s = s - (swap & 1);
211    // Set r = r' - q2, adding back 2s + 1 if r would be negative
212    let r = r_pre.wrapping_add(((s << 1) | 1) & swap) - q2;
213
214    (s, r)
215}
216
217/// Compute the square root and remainder for a 1 to 3 limb input, which must be normalized.
218/// This is the base case square root calculation.
219#[allow(clippy::cast_possible_truncation)]
220#[inline]
221const fn sqrt_rem_small<const VARTIME: bool>(value: &[Limb], s: &mut [Limb], r: &mut [Limb]) {
222    let len = value.len();
223    assert!(len < 4, "value exceeds maximum size for sqrt_rem_small");
224
225    if len == 0 {
226        return;
227    }
228
229    if len == 1 {
230        let base = value[0].0;
231        (s[0].0, r[0].0) = if VARTIME {
232            let root = base.isqrt();
233            (root, base.wrapping_sub(root.wrapping_pow(2)))
234        } else {
235            sqrt_rem_word(base)
236        };
237        return;
238    }
239
240    // Compute root, remainder for two words
241    let (s1, r1) = {
242        if VARTIME {
243            let base = word::join(value[len - 2].0, value[len - 1].0);
244            let root = base.isqrt();
245            (root as Word, base.wrapping_sub(root.wrapping_pow(2)))
246        } else {
247            let (s0, r0) = sqrt_rem_word(value[len - 1].0);
248            let (s1, r1) = sqrt_rem_expand_by_word(s0, r0 as WideWord, value[len - 2].0);
249            (s1 as Word, r1)
250        }
251    };
252    if len == 2 {
253        s[0].0 = s1;
254        (r[0].0, r[1].0) = word::split_wide(r1);
255        return;
256    }
257
258    // Expand to root, remainder for three words
259    let (s1, r1) = sqrt_rem_expand_by_word(s1, r1, value[len - 3].0);
260    (s[0].0, s[1].0) = word::split_wide(s1);
261    (r[0].0, r[1].0) = word::split_wide(r1);
262}
263
264#[cfg(test)]
265mod tests {
266    use super::{sqrt_rem_expand_by_word, sqrt_rem_word};
267    use crate::{WideWord, Word, word};
268
269    #[test]
270    fn sqrt_rem_word_expected() {
271        fn check(val: Word) {
272            let (root, rem) = sqrt_rem_word(val);
273            let check = root.pow(2) + rem;
274            assert_eq!(check, val, "val: {val}, root: {root}, rem: {rem}");
275        }
276
277        check(0);
278        check(1);
279        check(2);
280        check(Word::MAX);
281    }
282
283    #[test]
284    fn sqrt_rem_expand_by_word_expected() {
285        fn check(val: WideWord) {
286            let (lo, hi) = word::split_wide(val);
287            let s1 = hi.isqrt();
288            let r1 = WideWord::from(hi - s1.pow(2));
289            let (s, r) = sqrt_rem_expand_by_word(s1, r1, lo);
290            assert_eq!(s, val.isqrt());
291            assert_eq!(r, val - s.pow(2));
292        }
293
294        check(1 << (WideWord::BITS - 1));
295        check(2 << (WideWord::BITS - 2));
296        check(WideWord::MAX);
297    }
298}
crypto_bigint/uint/ref_type/sqrt.rs

crypto_bigint/uint/ref_type/
sqrt.rs
