pub fn modinv(a: u64, m: u64) -> u64Examples found in repository?
crates/competitive/src/math/discrete_logarithm.rs (line 26)
20fn solve_linear_congruence(a: u64, b: u64, m: u64) -> Option<(u64, u64)> {
21 let g = gcd(a, m);
22 if !b.is_multiple_of(g) {
23 return None;
24 }
25 let (a, b, m) = (a / g, b / g, m / g);
26 Some(((b as u128 * modinv(a, m) as u128 % m as u128) as _, m))
27}
28
29fn solve_linear_congruences<I>(abm: I) -> Option<(u64, u64)>
30where
31 I: IntoIterator<Item = (u64, u64, u64)>,
32{
33 let mut x = 0u64;
34 let mut m0 = 1u64;
35 for (a, b, m) in abm {
36 let mut b = b + m - a * x % m;
37 if b >= m {
38 b -= m;
39 }
40 let a = a * m0;
41 let g = gcd(a, m);
42 if !b.is_multiple_of(g) {
43 return None;
44 }
45 let (a, b, m) = (a / g, b / g, m / g);
46 x += (b as u128 * modinv(a, m) as u128 % m as u128 * m0 as u128) as u64;
47 m0 *= m;
48 }
49 Some((x, m0))
50}
51
52#[derive(Debug)]
53struct IndexCalculus {
54 primes: PrimeList,
55 br_primes: Vec<BarrettReduction<u64>>,
56 ic: HashMap<u64, IndexCalculusWithPrimitiveRoot>,
57}
58
59impl IndexCalculus {
60 fn new() -> Self {
61 Self {
62 primes: PrimeList::new(2),
63 br_primes: Default::default(),
64 ic: Default::default(),
65 }
66 }
67 fn discrete_logarithm(&mut self, a: u64, b: u64, p: u64) -> Option<(u64, u64)> {
68 let lim = ((((p as f64).log2() * (p as f64).log2().log2()).sqrt() / 2.0 + 1.).exp2() * 0.9)
69 as u32;
70 self.primes.reserve(lim);
71 let prime_count = self.primes.primes_lte(lim).count();
72 self.br_primes.extend(
73 self.primes
74 .primes_lte(lim)
75 .skip(self.br_primes.len())
76 .map(|p| BarrettReduction::<u64>::new(p.into())),
77 );
78 let br_primes = &self.br_primes[..prime_count];
79 self.ic
80 .entry(p)
81 .or_insert_with(|| IndexCalculusWithPrimitiveRoot::new(p, br_primes))
82 .discrete_logarithm(a, b, br_primes)
83 }
84}
85
86const A: [u32; 150] = [
87 62, 61, 60, 60, 59, 58, 58, 58, 57, 56, 56, 56, 56, 55, 55, 55, 54, 54, 54, 53, 53, 53, 53, 52,
88 52, 52, 52, 52, 52, 51, 50, 50, 50, 50, 49, 49, 49, 48, 48, 48, 48, 48, 47, 47, 47, 47, 47, 47,
89 47, 47, 47, 47, 47, 47, 47, 47, 45, 42, 42, 41, 41, 41, 41, 41, 41, 41, 40, 40, 40, 40, 40, 40,
90 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 38, 38, 38, 38, 38, 32, 32, 32, 32,
91 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 31, 31, 31, 31, 31,
92 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 22, 22, 22, 22,
93 22, 22, 22, 22, 22, 22,
94];
95
96fn factorize_smooth(mut x: u64, row: &mut [u64], br_primes: &[BarrettReduction<u64>]) -> bool {
97 for (j, (&br, r)) in br_primes.iter().zip(row).enumerate() {
98 *r = 0;
99 loop {
100 let (div, rem) = br.div_rem(x);
101 if rem != 0 {
102 break;
103 }
104 *r += 1;
105 x = div;
106 }
107 if j < 150 && x >= (1u64 << A[j]) {
108 break;
109 }
110 }
111 x == 1
112}
113
114#[derive(Debug)]
115struct QdrtPowPrec {
116 br_qdrt: BarrettReduction<u64>,
117 p0: Vec<u64>,
118 p1: Vec<u64>,
119 p2: Vec<u64>,
120 p3: Vec<u64>,
121}
122
123impl QdrtPowPrec {
124 fn new(a: u64, ord: u64, br: &BarrettReduction<u128>) -> Self {
125 let qdrt = (ord as f64).powf(0.25).ceil() as u64;
126 let br_qdrt = BarrettReduction::<u64>::new(qdrt);
127 let mut p0 = Vec::with_capacity(qdrt as usize);
128 let mut p1 = Vec::with_capacity(qdrt as usize);
129 let mut p2 = Vec::with_capacity(qdrt as usize);
130 let mut p3 = Vec::with_capacity(qdrt as usize);
131 let mut acc = 1u64;
132 for _ in 0..qdrt {
133 p0.push(acc);
134 acc = br.rem(acc as u128 * a as u128) as u64;
135 }
136 let a = acc;
137 acc = 1;
138 for _ in 0..qdrt {
139 p1.push(acc);
140 acc = br.rem(acc as u128 * a as u128) as u64;
141 }
142 let a = acc;
143 acc = 1;
144 for _ in 0..qdrt {
145 p2.push(acc);
146 acc = br.rem(acc as u128 * a as u128) as u64;
147 }
148 let a = acc;
149 acc = 1;
150 for _ in 0..qdrt {
151 p3.push(acc);
152 acc = br.rem(acc as u128 * a as u128) as u64;
153 }
154 Self {
155 br_qdrt,
156 p0,
157 p1,
158 p2,
159 p3,
160 }
161 }
162 fn pow(&self, mut k: u64, br: &BarrettReduction<u128>) -> u64 {
163 let (a, b) = self.br_qdrt.div_rem(k);
164 let mut x = self.p0[b as usize];
165 k = a;
166 if k > 0 {
167 let (a, b) = self.br_qdrt.div_rem(k);
168 x = br.rem(x as u128 * self.p1[b as usize] as u128) as u64;
169 k = a;
170 }
171 if k > 0 {
172 let (a, b) = self.br_qdrt.div_rem(k);
173 x = br.rem(x as u128 * self.p2[b as usize] as u128) as u64;
174 k = a;
175 }
176 if k > 0 {
177 let (_, b) = self.br_qdrt.div_rem(k);
178 x = br.rem(x as u128 * self.p3[b as usize] as u128) as u64;
179 }
180 x
181 }
182}
183
184fn index_calculus_for_primitive_root(
185 p: u64,
186 ord: u64,
187 br_primes: &[BarrettReduction<u64>],
188 prec: &QdrtPowPrec,
189) -> Vec<u64> {
190 let br_ord = BarrettReduction::<u128>::new(ord as u128);
191 let mul = |x: u64, y: u64| br_ord.rem(x as u128 * y as u128) as u64;
192 let sub = |x: u64, y: u64| if x < y { x + ord - y } else { x - y };
193
194 let pc = br_primes.len();
195 let mut mat: Vec<Vec<u64>> = vec![];
196 let mut rows: Vec<Vec<u64>> = vec![];
197
198 let mut rng = Xorshift::default();
199 let br = BarrettReduction::<u128>::new(p as u128);
200
201 for i in 0..pc {
202 for ri in 0usize.. {
203 let mut row = vec![0u64; pc + 1];
204 let mut kk = rng.rand(ord - 1) + 1;
205 let mut gkk = prec.pow(kk, &br);
206 let mut k = kk;
207 let mut gk = gkk;
208 while ri >= rows.len() {
209 row[pc] = k;
210 if factorize_smooth(gk, &mut row, br_primes) {
211 rows.push(row);
212 break;
213 }
214 if k + kk < ord {
215 k += kk;
216 gk = br.rem(gk as u128 * gkk as u128) as u64;
217 } else {
218 kk = rng.rand(ord - 1) + 1;
219 gkk = prec.pow(kk, &br);
220 k = kk;
221 gk = gkk;
222 }
223 }
224 let row = &mut rows[ri];
225 for j in 0..i {
226 if row[j] != 0 {
227 let b = mul(modinv(mat[j][j], ord), row[j]);
228 for (r, a) in row[j..].iter_mut().zip(&mat[j][j..]) {
229 *r = sub(*r, mul(*a, b));
230 }
231 }
232 assert_eq!(row[j], 0);
233 }
234 if gcd(row[i], ord) == 1 {
235 let last = rows.len() - 1;
236 rows.swap(ri, last);
237 mat.push(rows.pop().unwrap());
238 break;
239 }
240 }
241 }
242 for i in (0..pc).rev() {
243 for j in i + 1..pc {
244 mat[i][pc] = sub(mat[i][pc], mul(mat[i][j], mat[j][pc]));
245 }
246 mat[i][pc] = mul(mat[i][pc], modinv(mat[i][i], ord));
247 }
248 (0..pc).map(|i| mat[i][pc]).collect()
249}
250
251#[derive(Debug)]
252struct IndexCalculusWithPrimitiveRoot {
253 p: u64,
254 ord: u64,
255 prec: QdrtPowPrec,
256 coeff: Vec<u64>,
257}
258
259impl IndexCalculusWithPrimitiveRoot {
260 fn new(p: u64, br_primes: &[BarrettReduction<u64>]) -> Self {
261 let ord = p - 1;
262 let g = primitive_root(p);
263 let br = BarrettReduction::<u128>::new(p as u128);
264 let prec = QdrtPowPrec::new(g, ord, &br);
265 let coeff = index_calculus_for_primitive_root(p, ord, br_primes, &prec);
266 Self {
267 p,
268 ord,
269 prec,
270 coeff,
271 }
272 }
273 fn index_calculus(&self, a: u64, br_primes: &[BarrettReduction<u64>]) -> Option<u64> {
274 let p = self.p;
275 let ord = self.ord;
276 let br = BarrettReduction::<u128>::new(p as u128);
277 let a = br.rem(a as _) as u64;
278 if a == 1 {
279 return Some(0);
280 }
281 if p == 2 {
282 return None;
283 }
284
285 let mut rng = Xorshift::new();
286 let mut row = vec![0u64; br_primes.len()];
287 let mut kk = rng.rand(ord - 1) + 1;
288 let mut gkk = self.prec.pow(kk, &br);
289 let mut k = kk;
290 let mut gk = br.rem(gkk as u128 * a as u128) as u64;
291 loop {
292 if factorize_smooth(gk, &mut row, br_primes) {
293 let mut res = ord - k;
294 for (&c, &r) in self.coeff.iter().zip(&row) {
295 for _ in 0..r {
296 res += c;
297 if res >= ord {
298 res -= ord;
299 }
300 }
301 }
302 return Some(res);
303 }
304 if k + kk < ord {
305 k += kk;
306 gk = br.rem(gk as u128 * gkk as u128) as u64;
307 } else {
308 kk = rng.rand(ord - 1) + 1;
309 gkk = self.prec.pow(kk, &br);
310 k = kk;
311 gk = br.rem(gkk as u128 * a as u128) as u64;
312 }
313 }
314 }
315 fn discrete_logarithm(
316 &self,
317 a: u64,
318 b: u64,
319 br_primes: &[BarrettReduction<u64>],
320 ) -> Option<(u64, u64)> {
321 let p = self.p;
322 let ord = self.ord;
323 let br = BarrettReduction::<u128>::new(p as u128);
324 let a = br.rem(a as _) as u64;
325 let b = br.rem(b as _) as u64;
326 if a == 0 {
327 return if b == 0 { Some((1, 1)) } else { None };
328 }
329 if b == 0 {
330 return None;
331 }
332
333 let x = self.index_calculus(a, br_primes)?;
334 let y = self.index_calculus(b, br_primes)?;
335 solve_linear_congruence(x, y, ord)
336 }
337}
338
339thread_local!(
340 static IC: UnsafeCell<IndexCalculus> = UnsafeCell::new(IndexCalculus::new());
341);
342
343pub fn discrete_logarithm_prime_mod(a: u64, b: u64, p: u64) -> Option<u64> {
344 IC.with(|ic| unsafe { &mut *ic.get() }.discrete_logarithm(a, b, p))
345 .map(|t| t.0)
346}
347
348/// a^x ≡ b (mod n), a has order p^e
349fn pohlig_hellman_prime_power_order(a: u64, b: u64, n: u64, p: u64, e: u32) -> Option<u64> {
350 let br = BarrettReduction::<u128>::new(n as u128);
351 let mul = |x: u64, y: u64| br.rem(x as u128 * y as u128) as u64;
352 let block_size = (p as f64).sqrt().ceil() as u64;
353 let mut baby = HashMap::<u64, u64>::new();
354 let g = pow(a, p.pow(e - 1), &br);
355 let mut xj = 1;
356 for j in 0..block_size {
357 baby.entry(xj).or_insert(j);
358 xj = mul(xj, g);
359 }
360 let xi = modinv(xj, n);
361 let mut t = 0u64;
362 for k in 0..e {
363 let mut h = pow(mul(modinv(pow(a, t, &br), n), b), p.pow(e - 1 - k), &br);
364 let mut ok = false;
365 for i in (0..block_size * block_size).step_by(block_size as usize) {
366 if let Some(j) = baby.get(&h) {
367 t += (i + j) * p.pow(k);
368 ok = true;
369 break;
370 }
371 h = mul(h, xi);
372 }
373 if !ok {
374 return None;
375 }
376 }
377 Some(t)
378}
379
380/// a^x ≡ b (mod p^e)
381fn discrete_logarithm_prime_power(a: u64, b: u64, p: u64, e: u32) -> Option<(u64, u64)> {
382 assert_ne!(p, 0);
383 assert_ne!(e, 0);
384 let n = p.pow(e);
385 assert!(a < n);
386 assert!(b < n);
387 assert_eq!(gcd(a, p), 1);
388 if p == 1 {
389 return Some((0, 1));
390 }
391 if a == 0 {
392 return if b == 0 { Some((1, 1)) } else { None };
393 }
394 if b == 0 {
395 return None;
396 }
397 if e == 1 {
398 return IC.with(|ic| unsafe { &mut *ic.get() }.discrete_logarithm(a, b, p));
399 }
400 let br = BarrettReduction::<u128>::new(n as _);
401 if p == 2 {
402 if e >= 3 {
403 if a % 4 == 1 && b % 4 != 1 {
404 return None;
405 }
406 let aa = if a % 4 == 1 { a } else { n - a };
407 let bb = if b % 4 == 1 { b } else { n - b };
408 let g = 5;
409 let ord = n / 4;
410 let x = pohlig_hellman_prime_power_order(g, aa, n, p, e - 2)?;
411 let y = pohlig_hellman_prime_power_order(g, bb, n, p, e - 2)?;
412 let t = solve_linear_congruence(x, y, ord)?;
413 match (a % 4 == 1, b % 4 == 1) {
414 (true, true) => Some(t),
415 (false, true) if t.0 % 2 == 0 => Some((t.0, lcm(t.1, 2))),
416 (false, false) if t.0 % 2 == 1 => Some((t.0, lcm(t.1, 2))),
417 (false, false) if a == b => Some((1, lcm(t.1, 2))),
418 _ => None,
419 }
420 } else if a == 1 {
421 if b == 1 { Some((0, 1)) } else { None }
422 } else {
423 assert_eq!(a, 3);
424 if b == 1 {
425 Some((0, 2))
426 } else if b == 3 {
427 Some((1, 2))
428 } else {
429 None
430 }
431 }
432 } else {
433 let ord = n - n / p;
434 let pf_ord = prime_factors(ord);
435 let g = (2..)
436 .find(|&g| check_primitive_root(g, ord, &br, &pf_ord))
437 .unwrap();
438 let mut pf_p = prime_factors(p - 1);
439 pf_p.push((p, e - 1));
440 let mut abm = vec![];
441 for (q, c) in pf_p {
442 let m = q.pow(c);
443 let d = ord / m;
444 let gg = pow(g, d, &br);
445 let aa = pow(a, d, &br);
446 let bb = pow(b, d, &br);
447 let x = pohlig_hellman_prime_power_order(gg, aa, n, q, c)?;
448 let y = pohlig_hellman_prime_power_order(gg, bb, n, q, c)?;
449 abm.push((x, y, m));
450 }
451 solve_linear_congruences(abm)
452 }
453}
454
455/// a^x ≡ b (mod n)
456pub fn discrete_logarithm(a: u64, b: u64, n: u64) -> Option<u64> {
457 let a = a % n;
458 let b = b % n;
459 let d = 2.max(64 - n.leading_zeros() as u64);
460 let mut pw = 1 % n;
461 for i in 0..d {
462 if pw == b {
463 return Some(i);
464 }
465 pw = (pw as u128 * a as u128 % n as u128) as u64;
466 }
467 let g = gcd(pw, n);
468 if !b.is_multiple_of(g) {
469 return None;
470 }
471 let n = n / g;
472 let b = (b as u128 * modinv(pw, n) as u128 % n as u128) as u64;
473 let pf = prime_factors(n);
474 let mut abm = vec![];
475 for (p, e) in pf {
476 let q = p.pow(e);
477 let x = discrete_logarithm_prime_power(a % q, b % q, p, e)?;
478 abm.push((1, x.0, x.1));
479 }
480 solve_linear_congruences(abm).map(|x| x.0 + d)
481}