num_modular/
preinv.rs

1use crate::{DivExact, DivExactAssign, ModularUnaryOps};
2
3/// Pre-computing the modular inverse for fast divisibility check.
4///
5/// This struct stores the modular inverse of a divisor, and a limit for divisibility check.
6/// See <https://math.stackexchange.com/a/1251328> for the explanation of the trick
7#[must_use]
8#[derive(Debug, Clone, Copy)]
9pub struct PreModInv<T> {
10    d_inv: T, // modular inverse of divisor
11    q_lim: T, // limit of residue
12}
13
14macro_rules! impl_preinv_for_prim_int {
15    ($t:ident, $ns:ident) => {
16        mod $ns {
17            use super::*;
18            use crate::word::$t::*;
19
20            impl PreModInv<$t> {
21                /// Construct the preinv instance with raw values.
22                ///
23                /// This function can be used to initialize preinv in a constant context, the divisor d
24                /// is required only for verification of d_inv and q_lim.
25                #[inline]
26                pub const fn new(d_inv: $t, q_lim: $t) -> Self {
27                    Self { d_inv, q_lim }
28                }
29
30                // check if the divisor is consistent in debug mode
31                #[inline]
32                fn debug_check(&self, d: $t) {
33                    debug_assert!(d % 2 != 0, "only odd divisors are supported");
34                    debug_assert!(d.wrapping_mul(self.d_inv) == 1);
35                    debug_assert!(self.q_lim * d > (<$t>::MAX - d));
36                }
37            }
38
39            impl From<$t> for PreModInv<$t> {
40                #[inline]
41                fn from(v: $t) -> Self {
42                    use crate::word::$t::*;
43
44                    debug_assert!(v % 2 != 0, "only odd divisors are supported");
45                    let d_inv = extend(v).invm(&merge(0, 1)).unwrap() as $t;
46                    let q_lim = <$t>::MAX / v;
47                    Self { d_inv, q_lim }
48                }
49            }
50
51            impl DivExact<$t, PreModInv<$t>> for $t {
52                type Output = $t;
53                #[inline]
54                fn div_exact(self, d: $t, pre: &PreModInv<$t>) -> Option<Self> {
55                    pre.debug_check(d);
56                    let q = self.wrapping_mul(pre.d_inv);
57                    if q <= pre.q_lim {
58                        Some(q)
59                    } else {
60                        None
61                    }
62                }
63            }
64
65            impl DivExact<$t, PreModInv<$t>> for DoubleWord {
66                type Output = DoubleWord;
67
68                #[inline]
69                fn div_exact(self, d: $t, pre: &PreModInv<$t>) -> Option<Self::Output> {
70                    pre.debug_check(d);
71
72                    // this implementation comes from GNU factor,
73                    // see https://math.stackexchange.com/q/4436380/815652 for explanation
74
75                    let (n0, n1) = split(self);
76                    let q0 = n0.wrapping_mul(pre.d_inv);
77                    let nr0 = wmul(q0, d);
78                    let nr0 = split(nr0).1;
79                    if nr0 > n1 {
80                        return None;
81                    }
82                    let nr1 = n1 - nr0;
83                    let q1 = nr1.wrapping_mul(pre.d_inv);
84                    if q1 > pre.q_lim {
85                        return None;
86                    }
87                    Some(merge(q0, q1))
88                }
89            }
90
91            impl DivExactAssign<$t, PreModInv<$t>> for $t {
92                #[inline]
93                fn div_exact_assign(&mut self, d: $t, pre: &PreModInv<$t>) -> bool {
94                    match DivExact::div_exact(*self, d, pre) {
95                        Some(q) => {
96                            *self = q;
97                            true
98                        }
99                        None => false,
100                    }
101                }
102            }
103
104            impl DivExactAssign<$t, PreModInv<$t>> for DoubleWord {
105                #[inline]
106                fn div_exact_assign(&mut self, d: $t, pre: &PreModInv<$t>) -> bool {
107                    match DivExact::div_exact(*self, d, pre) {
108                        Some(q) => {
109                            *self = q;
110                            true
111                        }
112                        None => false,
113                    }
114                }
115            }
116        }
117    };
118}
119impl_preinv_for_prim_int!(u8, u8_impl);
120impl_preinv_for_prim_int!(u16, u16_impl);
121impl_preinv_for_prim_int!(u32, u32_impl);
122impl_preinv_for_prim_int!(u64, u64_impl);
123impl_preinv_for_prim_int!(usize, usize_impl);
124
125// XXX: unchecked div_exact can be introduced by not checking the q_lim,
126//      investigate this after `exact_div` is introduced or removed from core lib
127//      https://github.com/rust-lang/rust/issues/85122
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132    use rand::random;
133
134    #[test]
135    #[allow(unstable_name_collisions)]
136    fn div_exact_test() {
137        const N: u8 = 100;
138        for _ in 0..N {
139            // u8 test
140            let d = random::<u8>() | 1;
141            let pre: PreModInv<_> = d.into();
142
143            let n: u8 = random();
144            let expect = if n % d == 0 { Some(n / d) } else { None };
145            assert_eq!(n.div_exact(d, &pre), expect, "{} / {}", n, d);
146            let n: u16 = random();
147            let expect = if n % (d as u16) == 0 {
148                Some(n / (d as u16))
149            } else {
150                None
151            };
152            assert_eq!(n.div_exact(d, &pre), expect, "{} / {}", n, d);
153
154            // u16 test
155            let d = random::<u16>() | 1;
156            let pre: PreModInv<_> = d.into();
157
158            let n: u16 = random();
159            let expect = if n % d == 0 { Some(n / d) } else { None };
160            assert_eq!(n.div_exact(d, &pre), expect, "{} / {}", n, d);
161            let n: u32 = random();
162            let expect = if n % (d as u32) == 0 {
163                Some(n / (d as u32))
164            } else {
165                None
166            };
167            assert_eq!(n.div_exact(d, &pre), expect, "{} / {}", n, d);
168
169            // u32 test
170            let d = random::<u32>() | 1;
171            let pre: PreModInv<_> = d.into();
172
173            let n: u32 = random();
174            let expect = if n % d == 0 { Some(n / d) } else { None };
175            assert_eq!(n.div_exact(d, &pre), expect, "{} / {}", n, d);
176            let n: u64 = random();
177            let expect = if n % (d as u64) == 0 {
178                Some(n / (d as u64))
179            } else {
180                None
181            };
182            assert_eq!(n.div_exact(d, &pre), expect, "{} / {}", n, d);
183
184            // u64 test
185            let d = random::<u64>() | 1;
186            let pre: PreModInv<_> = d.into();
187
188            let n: u64 = random();
189            let expect = if n % d == 0 { Some(n / d) } else { None };
190            assert_eq!(n.div_exact(d, &pre), expect, "{} / {}", n, d);
191            let n: u128 = random();
192            let expect = if n % (d as u128) == 0 {
193                Some(n / (d as u128))
194            } else {
195                None
196            };
197            assert_eq!(n.div_exact(d, &pre), expect, "{} / {}", n, d);
198        }
199    }
200
201    #[test]
202    #[allow(unstable_name_collisions)]
203    fn div_exact_assign_test() {
204        const N: u8 = 100;
205
206        // () precompute (native integer division)
207        for _ in 0..N {
208            let d = random::<u8>() | 1;
209            let n: u8 = random();
210            let mut m = n;
211            let expect = if n % d == 0 { Some(n / d) } else { None };
212            let exact = m.div_exact_assign(d, &());
213            assert_eq!(exact, expect.is_some(), "{} / {}", n, d);
214            assert_eq!(m, expect.unwrap_or(n), "{} / {}", n, d);
215        }
216
217        // PreModInv precompute
218        for _ in 0..N {
219            let d = random::<u8>() | 1;
220            let pre: PreModInv<_> = d.into();
221
222            // single word
223            let n: u8 = random();
224            let mut m = n;
225            let expect = if n % d == 0 { Some(n / d) } else { None };
226            let exact = m.div_exact_assign(d, &pre);
227            assert_eq!(exact, expect.is_some(), "{} / {}", n, d);
228            assert_eq!(m, expect.unwrap_or(n), "{} / {}", n, d);
229
230            // double word
231            let n: u16 = random();
232            let mut m = n;
233            let expect = if n % (d as u16) == 0 {
234                Some(n / (d as u16))
235            } else {
236                None
237            };
238            let exact = m.div_exact_assign(d, &pre);
239            assert_eq!(exact, expect.is_some(), "{} / {}", n, d);
240            assert_eq!(m, expect.unwrap_or(n), "{} / {}", n, d);
241        }
242    }
243}
num_modular/preinv.rs

num_modular/
preinv.rs
