fn mul_mod_raw<const P: u32>(a: u32, b: u32) -> u32Examples found in repository?
crates/competitive/src/math/fast_prime_mod.rs (line 128)
125 pub fn inverse(&self, x: u32) -> u32 {
126 assert!(1 <= x && x < P);
127 let (i, b) = self.small_fraction(x);
128 mul_mod_raw::<P>(self.inv[i], b)
129 }
130}
131
132impl<const P: u32, const BUILD_INV: bool> FastPrimeMod<P, BUILD_INV, true> {
133 /// Returns the primitive root used by this table.
134 #[inline]
135 pub fn primitive_root(&self) -> u32 {
136 self.root
137 }
138
139 /// Returns `a^exp mod P`.
140 ///
141 /// `0^0` is defined as `1`.
142 ///
143 /// # Panics
144 ///
145 /// Panics unless `a < P`.
146 #[inline]
147 pub fn pow(&self, a: u32, exp: u64) -> u32 {
148 assert!(a < P);
149 if a == 0 {
150 return if exp == 0 { 1 } else { 0 };
151 }
152 let ord = (P - 1) as u64;
153 self.pow_nonzero_reduced(a, (exp % ord) as u32)
154 }
155
156 /// Returns `a^exp_mod mod P` for a non-zero base and a reduced exponent.
157 ///
158 /// # Panics
159 ///
160 /// Panics unless `1 <= a < P` and `exp_mod < P - 1`.
161 #[inline]
162 pub fn pow_nonzero_reduced(&self, a: u32, exp_mod: u32) -> u32 {
163 assert!(1 <= a && a < P);
164 assert!(exp_mod < P - 1);
165 let exp = (self.log_r(a) as u64 * exp_mod as u64 % (P - 1) as u64) as u32;
166 self.pow_root_reduced(exp)
167 }
168
169 /// Returns `r^exp_mod mod P`, where `r` is this table's primitive root.
170 ///
171 /// # Panics
172 ///
173 /// Panics unless `exp_mod < P - 1`.
174 #[inline]
175 pub fn pow_root_reduced(&self, exp_mod: u32) -> u32 {
176 assert!(exp_mod < P - 1);
177 pow_root_raw::<P>(exp_mod, &self.pow_lo, &self.pow_hi)
178 }
179
180 #[inline]
181 fn log_r(&self, x: u32) -> u32 {
182 let (i, b) = self.small_fraction(x);
183 let k = table_k::<P>();
184 let ord = P - 1;
185 self.log[i] + ord - self.log[k + b as usize]
186 }
187}
188
189impl<const P: u32, const BUILD_INV: bool, const BUILD_POW: bool> Default
190 for FastPrimeMod<P, BUILD_INV, BUILD_POW>
191{
192 fn default() -> Self {
193 Self::new()
194 }
195}
196
197fn build_pow<const P: u32>(root: u32) -> (Box<[u32]>, Box<[u32]>) {
198 let mut pow_lo = vec![0; POW_BLOCK + 1].into_boxed_slice();
199 let mut pow_hi = vec![0; POW_BLOCK + 1].into_boxed_slice();
200 pow_lo[0] = 1;
201 pow_hi[0] = 1;
202 for i in 0..POW_BLOCK {
203 pow_lo[i + 1] = mul_mod_raw::<P>(pow_lo[i], root);
204 }
205 let block_power = pow_lo[POW_BLOCK];
206 for i in 0..POW_BLOCK {
207 pow_hi[i + 1] = mul_mod_raw::<P>(pow_hi[i], block_power);
208 }
209 (pow_lo, pow_hi)
210}
211
212fn build_inv<const P: u32>() -> Box<[u32]> {
213 let k = table_k::<P>();
214 let mut inv = vec![0; table_len::<P>()].into_boxed_slice();
215 inv[k + 1] = 1;
216 for i in 2..=k {
217 let q = P.div_ceil(i as u32);
218 let r = i as u32 * q - P;
219 inv[k + i] = mul_mod_raw::<P>(inv[k + r as usize], q);
220 }
221 for i in 1..=k {
222 inv[k - i] = P - inv[k + i];
223 }
224 inv
225}
226
227fn build_log<const P: u32>(root: u32, pow_lo: &[u32], pow_hi: &[u32]) -> Box<[u32]> {
228 let k = table_k::<P>();
229 let ord = P - 1;
230 let mut lpf = vec![0; k + 1].into_boxed_slice();
231 let mut primes = vec![];
232 lpf[1] = 1;
233 for i in 2..=k {
234 if lpf[i] == 0 {
235 lpf[i] = i as u32;
236 primes.push(i as u32);
237 }
238 for &p in primes.iter() {
239 let p = p as usize;
240 if p > lpf[i] as usize || p > k / i {
241 break;
242 }
243 lpf[i * p] = p as u32;
244 }
245 }
246
247 let baby_size = (BSGS_SIZE as u32).min(ord);
248 let mut baby = U32Map::new(baby_size as usize);
249 let mut pw = 1;
250 for i in 0..baby_size {
251 baby.insert(pw, i);
252 pw = mul_mod_raw::<P>(pw, root);
253 }
254 let q = pow_root_raw::<P>(ord - baby_size, pow_lo, pow_hi);
255
256 let mut log = vec![0; table_len::<P>()].into_boxed_slice();
257 log[k + 1] = 0;
258 let mut rng = Xorshift::default();
259 let small_primes = [2, 3, 5, 7, 11, 13, 17, 19];
260 for i in 2..=k {
261 let p = lpf[i] as usize;
262 if p < i {
263 log[k + i] = add_mod(log[k + p], log[k + i / p], ord);
264 } else if i < 100 {
265 let mut x = i as u32;
266 let mut ans = 0;
267 loop {
268 if let Some(v) = baby.get(x) {
269 log[k + i] = ans + v;
270 break;
271 }
272 ans += baby_size;
273 x = mul_mod_raw::<P>(x, q);
274 }
275 } else if i > P as usize / i {
276 let j = (P as usize) / i;
277 let r = (P as usize) % i;
278 let x = add_mod(log[k + r], ord / 2, ord);
279 let y = log[k + j];
280 log[k + i] = if x >= y { x - y } else { x + ord - y };
281 } else {
282 loop {
283 let exp = rng.rand(ord as u64) as u32;
284 let mut ans = if exp == 0 { 0 } else { ord - exp };
285 let mut x = mul_mod_raw::<P>(i as u32, pow_root_raw::<P>(exp, pow_lo, pow_hi));
286 for q in small_primes {
287 while x.is_multiple_of(q) {
288 x /= q;
289 ans = add_mod(ans, log[k + q as usize], ord);
290 }
291 }
292 if x as usize >= k {
293 continue;
294 }
295 while (i as u32) < x && lpf[x as usize] < i as u32 {
296 let q = lpf[x as usize];
297 x /= q;
298 ans = add_mod(ans, log[k + q as usize], ord);
299 }
300 if 1 < x && x < i as u32 {
301 ans = add_mod(ans, log[k + x as usize], ord);
302 x = 1;
303 }
304 if x == 1 {
305 log[k + i] = ans;
306 break;
307 }
308 }
309 }
310 }
311 for i in 1..=k {
312 log[k - i] = add_mod(log[k + i], ord / 2, ord);
313 }
314 log
315}
316
317fn build_frac<const P: u32>() -> Box<[u32]> {
318 let mut frac = vec![0; FRAC_LEN].into_boxed_slice();
319 let mut stack = vec![(0u16, 1u16, 1u16, 1u16)];
320 while let Some((a, b, c, d)) = stack.pop() {
321 let nb = b + d;
322 if nb < FRAC_DEN_LIMIT {
323 stack.push((a + c, nb, c, d));
324 stack.push((a, b, a + c, nb));
325 } else {
326 let s = (a as u64 * P as u64 / ((1u64 << FRAC_SHIFT) * b as u64)) as usize;
327 let t = (c as u64 * P as u64 / ((1u64 << FRAC_SHIFT) * d as u64)) as usize;
328 frac[s] = pack_frac(a, b);
329 frac[t] = pack_frac(c, d);
330 let a = a.min(c);
331 let b = b.min(d);
332 if s + 1 < t {
333 for x in &mut frac[s + 1..t] {
334 *x = pack_frac(a, b);
335 }
336 }
337 }
338 }
339 frac
340}
341
342#[inline(always)]
343fn pow_root_raw<const P: u32>(exp: u32, pow_lo: &[u32], pow_hi: &[u32]) -> u32 {
344 let lo = exp as usize & (POW_BLOCK - 1);
345 let hi = exp as usize >> POW_BLOCK_BITS;
346 mul_mod_raw::<P>(pow_lo[lo], pow_hi[hi])
347}