fixed_dsp/basic/
sqrt.rs

1use crate::common::error::{Error, Result};
2use crate::common::tables::{SQRT_INITIAL_LUT_U16, SQRT_INITIAL_LUT_U32};
3
4#[inline]
5fn sqrt_newton_q15(number: i32, mut var1: i32) -> i32 {
6    for _ in 0..3 {
7        let mut temp = (var1 * var1) >> 12;
8        temp = (number * temp) >> 15;
9        temp = 0x3000 - temp;
10        var1 = (var1 * temp) >> 13;
11    }
12
13    (number * var1) >> 12
14}
15
16#[inline]
17fn sqrt_newton_q31(number: i64, mut var1: i64) -> i64 {
18    for _ in 0..3 {
19        let mut temp = (var1 * var1) >> 28;
20        temp = (number * temp) >> 31;
21        temp = 0x30000000 - temp;
22        var1 = (var1 * temp) >> 29;
23    }
24
25    (number * var1) >> 28
26}
27
28/// Q15 square root using the CMSIS lookup-table + Newton iteration.
29pub fn sqrt_i16(input: i16) -> Result<i16> {
30    if input < 0 {
31        return Err(Error::NanInf);
32    }
33
34    if input == 0 {
35        return Ok(0);
36    }
37
38    let input_u = input as u16 as u32;
39    let sign_bits = input_u.leading_zeros() as i32 - 17;
40    let shift = if sign_bits & 1 == 0 {
41        sign_bits
42    } else {
43        sign_bits - 1
44    } as u32;
45
46    let number = (input as i32).wrapping_shl(shift);
47    let index = (((number >> 11) as usize) - (0x2000usize >> 11)) as usize;
48    let var1 = SQRT_INITIAL_LUT_U16[index] as i32;
49    let mut result = sqrt_newton_q15(number, var1);
50
51    let final_shift = if sign_bits & 1 == 0 {
52        (sign_bits / 2) as u32
53    } else {
54        ((sign_bits - 1) / 2) as u32
55    };
56    result >>= final_shift;
57
58    Ok(result as i16)
59}
60
61/// Q31 square root using the CMSIS lookup-table + Newton iteration.
62pub fn sqrt_i32(input: i32) -> Result<i32> {
63    if input < 0 {
64        return Err(Error::NanInf);
65    }
66
67    if input == 0 {
68        return Ok(0);
69    }
70
71    let input_u = input as u32;
72    let sign_bits = input_u.leading_zeros() as i32 - 1;
73    let shift = if sign_bits & 1 == 0 {
74        sign_bits
75    } else {
76        sign_bits - 1
77    } as u32;
78
79    let number = (input as i64).wrapping_shl(shift);
80    let index = (((number >> 26) as usize) - (0x20000000usize >> 26)) as usize;
81    let var1 = SQRT_INITIAL_LUT_U32[index] as i64;
82    let mut result = sqrt_newton_q31(number, var1);
83
84    let final_shift = if sign_bits & 1 == 0 {
85        (sign_bits / 2) as u32
86    } else {
87        ((sign_bits - 1) / 2) as u32
88    };
89    result >>= final_shift;
90
91    Ok(result as i32)
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn sqrt_q15_handles_sign_and_zero() {
100        assert_eq!(sqrt_i16(-1), Err(Error::NanInf));
101        assert_eq!(sqrt_i16(0), Ok(0));
102    }
103
104    #[test]
105    fn sqrt_q31_handles_sign_and_zero() {
106        assert_eq!(sqrt_i32(-1), Err(Error::NanInf));
107        assert_eq!(sqrt_i32(0), Ok(0));
108    }
109}
fixed_dsp/basic/sqrt.rs

fixed_dsp/basic/
sqrt.rs
