Struct BarrettReduction

Source

pub struct BarrettReduction<T> { /* private fields */ }

Implementations§

Source §

pub const fn new(m: u64) -> Self

Examples found in repository ?

crates/competitive/src/num/mint/mint_basic.rs (line 132)

130    pub fn set_mod(m: u32) {
131        DYN_MODULUS_U32
132            .with(|cell| unsafe { *cell.get() = BarrettReduction::<u64>::new(m as u64) });
133    }

More examples

Hide additional examples

crates/competitive/src/math/floor_sum.rs (line 26)

22pub fn floor_sum(n: u64, a: u64, b: u64, m: u64) -> u64 {
23    let mut ans = Wrapping(0u64);
24    let (mut n, mut m, mut a, mut b) = (Wrapping(n), m, a, b);
25    loop {
26        let br = BarrettReduction::<u64>::new(m);
27        if a >= m {
28            let (q, r) = br.div_rem(a);
29            ans += choose2(n) * q;
30            a = r;
31        }
32        if b >= m {
33            let (q, r) = br.div_rem(b);
34            ans += n * q;
35            b = r;
36        }
37        let y_max = (n * a + b).0;
38        if y_max < m {
39            break;
40        }
41        let (q, r) = br.div_rem(y_max);
42        n = Wrapping(q);
43        b = r;
44        swap(&mut m, &mut a);
45    }
46    ans.0
47}

crates/competitive/src/math/discrete_logarithm.rs (line 73)

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 u64;
70        self.primes.reserve(lim);
71        let primes = self.primes.primes_lte(lim);
72        while self.br_primes.len() < primes.len() {
73            let br = BarrettReduction::<u64>::new(primes[self.br_primes.len()]);
74            self.br_primes.push(br);
75        }
76        let br_primes = &self.br_primes[..primes.len()];
77        self.ic
78            .entry(p)
79            .or_insert_with(|| IndexCalculusWithPrimitiveRoot::new(p, br_primes))
80            .discrete_logarithm(a, b, br_primes)
81    }
82}
83
84const A: [u32; 150] = [
85    62, 61, 60, 60, 59, 58, 58, 58, 57, 56, 56, 56, 56, 55, 55, 55, 54, 54, 54, 53, 53, 53, 53, 52,
86    52, 52, 52, 52, 52, 51, 50, 50, 50, 50, 49, 49, 49, 48, 48, 48, 48, 48, 47, 47, 47, 47, 47, 47,
87    47, 47, 47, 47, 47, 47, 47, 47, 45, 42, 42, 41, 41, 41, 41, 41, 41, 41, 40, 40, 40, 40, 40, 40,
88    40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 38, 38, 38, 38, 38, 32, 32, 32, 32,
89    32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 31, 31, 31, 31, 31,
90    31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 22, 22, 22, 22,
91    22, 22, 22, 22, 22, 22,
92];
93
94fn factorize_smooth(mut x: u64, row: &mut [u64], br_primes: &[BarrettReduction<u64>]) -> bool {
95    for (j, (&br, r)) in br_primes.iter().zip(row).enumerate() {
96        *r = 0;
97        loop {
98            let (div, rem) = br.div_rem(x);
99            if rem != 0 {
100                break;
101            }
102            *r += 1;
103            x = div;
104        }
105        if j < 150 && x >= (1u64 << A[j]) {
106            break;
107        }
108    }
109    x == 1
110}
111
112#[derive(Debug)]
113struct QdrtPowPrec {
114    br_qdrt: BarrettReduction<u64>,
115    p0: Vec<u64>,
116    p1: Vec<u64>,
117    p2: Vec<u64>,
118    p3: Vec<u64>,
119}
120
121impl QdrtPowPrec {
122    fn new(a: u64, ord: u64, br: &BarrettReduction<u128>) -> Self {
123        let qdrt = (ord as f64).powf(0.25).ceil() as u64;
124        let br_qdrt = BarrettReduction::<u64>::new(qdrt);
125        let mut p0 = Vec::with_capacity(qdrt as usize);
126        let mut p1 = Vec::with_capacity(qdrt as usize);
127        let mut p2 = Vec::with_capacity(qdrt as usize);
128        let mut p3 = Vec::with_capacity(qdrt as usize);
129        let mut acc = 1u64;
130        for _ in 0..qdrt {
131            p0.push(acc);
132            acc = br.rem(acc as u128 * a as u128) as u64;
133        }
134        let a = acc;
135        acc = 1;
136        for _ in 0..qdrt {
137            p1.push(acc);
138            acc = br.rem(acc as u128 * a as u128) as u64;
139        }
140        let a = acc;
141        acc = 1;
142        for _ in 0..qdrt {
143            p2.push(acc);
144            acc = br.rem(acc as u128 * a as u128) as u64;
145        }
146        let a = acc;
147        acc = 1;
148        for _ in 0..qdrt {
149            p3.push(acc);
150            acc = br.rem(acc as u128 * a as u128) as u64;
151        }
152        Self {
153            br_qdrt,
154            p0,
155            p1,
156            p2,
157            p3,
158        }
159    }

Source

pub const fn get_mod(&self) -> u64

Source

pub const fn div_rem(&self, a: u64) -> (u64, u64)

Examples found in repository ?

crates/competitive/src/math/discrete_logarithm.rs (line 98)

94fn factorize_smooth(mut x: u64, row: &mut [u64], br_primes: &[BarrettReduction<u64>]) -> bool {
95    for (j, (&br, r)) in br_primes.iter().zip(row).enumerate() {
96        *r = 0;
97        loop {
98            let (div, rem) = br.div_rem(x);
99            if rem != 0 {
100                break;
101            }
102            *r += 1;
103            x = div;
104        }
105        if j < 150 && x >= (1u64 << A[j]) {
106            break;
107        }
108    }
109    x == 1
110}
111
112#[derive(Debug)]
113struct QdrtPowPrec {
114    br_qdrt: BarrettReduction<u64>,
115    p0: Vec<u64>,
116    p1: Vec<u64>,
117    p2: Vec<u64>,
118    p3: Vec<u64>,
119}
120
121impl QdrtPowPrec {
122    fn new(a: u64, ord: u64, br: &BarrettReduction<u128>) -> Self {
123        let qdrt = (ord as f64).powf(0.25).ceil() as u64;
124        let br_qdrt = BarrettReduction::<u64>::new(qdrt);
125        let mut p0 = Vec::with_capacity(qdrt as usize);
126        let mut p1 = Vec::with_capacity(qdrt as usize);
127        let mut p2 = Vec::with_capacity(qdrt as usize);
128        let mut p3 = Vec::with_capacity(qdrt as usize);
129        let mut acc = 1u64;
130        for _ in 0..qdrt {
131            p0.push(acc);
132            acc = br.rem(acc as u128 * a as u128) as u64;
133        }
134        let a = acc;
135        acc = 1;
136        for _ in 0..qdrt {
137            p1.push(acc);
138            acc = br.rem(acc as u128 * a as u128) as u64;
139        }
140        let a = acc;
141        acc = 1;
142        for _ in 0..qdrt {
143            p2.push(acc);
144            acc = br.rem(acc as u128 * a as u128) as u64;
145        }
146        let a = acc;
147        acc = 1;
148        for _ in 0..qdrt {
149            p3.push(acc);
150            acc = br.rem(acc as u128 * a as u128) as u64;
151        }
152        Self {
153            br_qdrt,
154            p0,
155            p1,
156            p2,
157            p3,
158        }
159    }
160    fn pow(&self, mut k: u64, br: &BarrettReduction<u128>) -> u64 {
161        let (a, b) = self.br_qdrt.div_rem(k);
162        let mut x = self.p0[b as usize];
163        k = a;
164        if k > 0 {
165            let (a, b) = self.br_qdrt.div_rem(k);
166            x = br.rem(x as u128 * self.p1[b as usize] as u128) as u64;
167            k = a;
168        }
169        if k > 0 {
170            let (a, b) = self.br_qdrt.div_rem(k);
171            x = br.rem(x as u128 * self.p2[b as usize] as u128) as u64;
172            k = a;
173        }
174        if k > 0 {
175            let (_, b) = self.br_qdrt.div_rem(k);
176            x = br.rem(x as u128 * self.p3[b as usize] as u128) as u64;
177        }
178        x
179    }

More examples

Hide additional examples

crates/competitive/src/math/floor_sum.rs (line 28)

22pub fn floor_sum(n: u64, a: u64, b: u64, m: u64) -> u64 {
23    let mut ans = Wrapping(0u64);
24    let (mut n, mut m, mut a, mut b) = (Wrapping(n), m, a, b);
25    loop {
26        let br = BarrettReduction::<u64>::new(m);
27        if a >= m {
28            let (q, r) = br.div_rem(a);
29            ans += choose2(n) * q;
30            a = r;
31        }
32        if b >= m {
33            let (q, r) = br.div_rem(b);
34            ans += n * q;
35            b = r;
36        }
37        let y_max = (n * a + b).0;
38        if y_max < m {
39            break;
40        }
41        let (q, r) = br.div_rem(y_max);
42        n = Wrapping(q);
43        b = r;
44        swap(&mut m, &mut a);
45    }
46    ans.0
47}

Source

pub const fn div(&self, a: u64) -> u64

Source

pub const fn rem(&self, a: u64) -> u64

Examples found in repository ?

crates/competitive/src/num/mint/mint_basic.rs (line 135)

134    fn rem(x: u64) -> u64 {
135        DYN_MODULUS_U32.with(|cell| unsafe { (*cell.get()).rem(x) })
136    }

Source §

impl BarrettReduction<u128>

Source

pub const fn new(m: u128) -> Self

Examples found in repository ?

crates/competitive/src/math/miller_rabin.rs (line 122)

121pub fn miller_rabin(n: u64) -> bool {
122    miller_rabin_with_br(n, &BarrettReduction::<u128>::new(n as u128))
123}

More examples

Hide additional examples

crates/competitive/src/num/mint/mint_basic.rs (line 148)

146    pub fn set_mod(m: u64) {
147        DYN_MODULUS_U64
148            .with(|cell| unsafe { *cell.get() = BarrettReduction::<u128>::new(m as u128) })
149    }

crates/competitive/src/math/primitive_root.rs (line 9)

3pub fn primitive_root(p: u64) -> u64 {
4    if p == 2 {
5        return 1;
6    }
7    let phi = p - 1;
8    let pf = prime_factors(phi);
9    let br = BarrettReduction::<u128>::new(p as _);
10    (2..)
11        .find(|&g| check_primitive_root(g, phi, &br, &pf))
12        .unwrap()
13}

crates/competitive/src/math/prime_factors.rs (line 4)

3fn find_factor(n: u64) -> Option<u64> {
4    let br = BarrettReduction::<u128>::new(n as u128);
5    if miller_rabin_with_br(n, &br) {
6        return None;
7    }
8    let m = 1u64 << ((63 - n.leading_zeros()) / 5);
9    let sub = |x: u64, y: u64| if x > y { x - y } else { y - x };
10    let mul = |x: u64, y: u64| br.rem(x as u128 * y as u128) as u64;
11    for c in 12.. {
12        let f = |x: u64| (br.rem(x as u128 * x as u128) + c) as u64;
13        let (mut x, mut y, mut r, mut g, mut k, mut ys) = (0, 2, 1, 1, 0, 0);
14        while g == 1 {
15            x = y;
16            for _ in 0..r {
17                y = f(y);
18            }
19            while r > k && g == 1 {
20                ys = y;
21                let mut q = 1;
22                for _ in 0..m.min(r - k) {
23                    y = f(y);
24                    q = mul(q, sub(x, y));
25                }
26                g = gcd(q, n);
27                k += m;
28            }
29            r <<= 1;
30        }
31        if g == n {
32            g = 1;
33            while g == 1 {
34                ys = f(ys);
35                g = gcd(sub(x, ys), n);
36            }
37        }
38        if g < n {
39            return Some(g);
40        }
41    }
42    unreachable!();
43}

crates/competitive/src/math/discrete_logarithm.rs (line 188)

182fn index_calculus_for_primitive_root(
183    p: u64,
184    ord: u64,
185    br_primes: &[BarrettReduction<u64>],
186    prec: &QdrtPowPrec,
187) -> Vec<u64> {
188    let br_ord = BarrettReduction::<u128>::new(ord as u128);
189    let mul = |x: u64, y: u64| br_ord.rem(x as u128 * y as u128) as u64;
190    let sub = |x: u64, y: u64| if x < y { x + ord - y } else { x - y };
191
192    let pc = br_primes.len();
193    let mut mat: Vec<Vec<u64>> = vec![];
194    let mut rows: Vec<Vec<u64>> = vec![];
195
196    let mut rng = Xorshift::default();
197    let br = BarrettReduction::<u128>::new(p as u128);
198
199    for i in 0..pc {
200        for ri in 0usize.. {
201            let mut row = vec![0u64; pc + 1];
202            let mut kk = rng.rand(ord - 1) + 1;
203            let mut gkk = prec.pow(kk, &br);
204            let mut k = kk;
205            let mut gk = gkk;
206            while ri >= rows.len() {
207                row[pc] = k;
208                if factorize_smooth(gk, &mut row, br_primes) {
209                    rows.push(row);
210                    break;
211                }
212                if k + kk < ord {
213                    k += kk;
214                    gk = br.rem(gk as u128 * gkk as u128) as u64;
215                } else {
216                    kk = rng.rand(ord - 1) + 1;
217                    gkk = prec.pow(kk, &br);
218                    k = kk;
219                    gk = gkk;
220                }
221            }
222            let row = &mut rows[ri];
223            for j in 0..i {
224                if row[j] != 0 {
225                    let b = mul(modinv(mat[j][j], ord), row[j]);
226                    for (r, a) in row[j..].iter_mut().zip(&mat[j][j..]) {
227                        *r = sub(*r, mul(*a, b));
228                    }
229                }
230                assert_eq!(row[j], 0);
231            }
232            if gcd(row[i], ord) == 1 {
233                let last = rows.len() - 1;
234                rows.swap(ri, last);
235                mat.push(rows.pop().unwrap());
236                break;
237            }
238        }
239    }
240    for i in (0..pc).rev() {
241        for j in i + 1..pc {
242            mat[i][pc] = sub(mat[i][pc], mul(mat[i][j], mat[j][pc]));
243        }
244        mat[i][pc] = mul(mat[i][pc], modinv(mat[i][i], ord));
245    }
246    (0..pc).map(|i| (mat[i][pc])).collect()
247}
248
249#[derive(Debug)]
250struct IndexCalculusWithPrimitiveRoot {
251    p: u64,
252    ord: u64,
253    prec: QdrtPowPrec,
254    coeff: Vec<u64>,
255}
256
257impl IndexCalculusWithPrimitiveRoot {
258    fn new(p: u64, br_primes: &[BarrettReduction<u64>]) -> Self {
259        let ord = p - 1;
260        let g = primitive_root(p);
261        let br = BarrettReduction::<u128>::new(p as u128);
262        let prec = QdrtPowPrec::new(g, ord, &br);
263        let coeff = index_calculus_for_primitive_root(p, ord, br_primes, &prec);
264        Self {
265            p,
266            ord,
267            prec,
268            coeff,
269        }
270    }
271    fn index_calculus(&self, a: u64, br_primes: &[BarrettReduction<u64>]) -> Option<u64> {
272        let p = self.p;
273        let ord = self.ord;
274        let br = BarrettReduction::<u128>::new(p as u128);
275        let a = br.rem(a as _) as u64;
276        if a == 1 {
277            return Some(0);
278        }
279        if p == 2 {
280            return None;
281        }
282
283        let mut rng = Xorshift::new();
284        let mut row = vec![0u64; br_primes.len()];
285        let mut kk = rng.rand(ord - 1) + 1;
286        let mut gkk = self.prec.pow(kk, &br);
287        let mut k = kk;
288        let mut gk = br.rem(gkk as u128 * a as u128) as u64;
289        loop {
290            if factorize_smooth(gk, &mut row, br_primes) {
291                let mut res = ord - k;
292                for (&c, &r) in self.coeff.iter().zip(&row) {
293                    for _ in 0..r {
294                        res += c;
295                        if res >= ord {
296                            res -= ord;
297                        }
298                    }
299                }
300                return Some(res);
301            }
302            if k + kk < ord {
303                k += kk;
304                gk = br.rem(gk as u128 * gkk as u128) as u64;
305            } else {
306                kk = rng.rand(ord - 1) + 1;
307                gkk = self.prec.pow(kk, &br);
308                k = kk;
309                gk = br.rem(gkk as u128 * a as u128) as u64;
310            }
311        }
312    }
313    fn discrete_logarithm(
314        &self,
315        a: u64,
316        b: u64,
317        br_primes: &[BarrettReduction<u64>],
318    ) -> Option<(u64, u64)> {
319        let p = self.p;
320        let ord = self.ord;
321        let br = BarrettReduction::<u128>::new(p as u128);
322        let a = br.rem(a as _) as u64;
323        let b = br.rem(b as _) as u64;
324        if a == 0 {
325            return if b == 0 { Some((1, 1)) } else { None };
326        }
327        if b == 0 {
328            return None;
329        }
330
331        let x = self.index_calculus(a, br_primes)?;
332        let y = self.index_calculus(b, br_primes)?;
333        solve_linear_congruence(x, y, ord)
334    }
335}
336
337thread_local!(
338    static IC: UnsafeCell<IndexCalculus> = UnsafeCell::new(IndexCalculus::new());
339);
340
341pub fn discrete_logarithm_prime_mod(a: u64, b: u64, p: u64) -> Option<u64> {
342    IC.with(|ic| unsafe { &mut *ic.get() }.discrete_logarithm(a, b, p))
343        .map(|t| t.0)
344}
345
346/// a^x ≡ b (mod n), a has order p^e
347fn pohlig_hellman_prime_power_order(a: u64, b: u64, n: u64, p: u64, e: u32) -> Option<u64> {
348    let br = BarrettReduction::<u128>::new(n as u128);
349    let mul = |x: u64, y: u64| br.rem(x as u128 * y as u128) as u64;
350    let block_size = (p as f64).sqrt().ceil() as u64;
351    let mut baby = HashMap::<u64, u64>::new();
352    let g = pow(a, p.pow(e - 1), &br);
353    let mut xj = 1;
354    for j in 0..block_size {
355        baby.entry(xj).or_insert(j);
356        xj = mul(xj, g);
357    }
358    let xi = modinv(xj, n);
359    let mut t = 0u64;
360    for k in 0..e {
361        let mut h = pow(mul(modinv(pow(a, t, &br), n), b), p.pow(e - 1 - k), &br);
362        let mut ok = false;
363        for i in (0..block_size * block_size).step_by(block_size as usize) {
364            if let Some(j) = baby.get(&h) {
365                t += (i + j) * p.pow(k);
366                ok = true;
367                break;
368            }
369            h = mul(h, xi);
370        }
371        if !ok {
372            return None;
373        }
374    }
375    Some(t)
376}
377
378/// a^x ≡ b (mod p^e)
379fn discrete_logarithm_prime_power(a: u64, b: u64, p: u64, e: u32) -> Option<(u64, u64)> {
380    assert_ne!(p, 0);
381    assert_ne!(e, 0);
382    let n = p.pow(e);
383    assert!(a < n);
384    assert!(b < n);
385    assert_eq!(gcd(a, p), 1);
386    if p == 1 {
387        return Some((0, 1));
388    }
389    if a == 0 {
390        return if b == 0 { Some((1, 1)) } else { None };
391    }
392    if b == 0 {
393        return None;
394    }
395    if e == 1 {
396        return IC.with(|ic| unsafe { &mut *ic.get() }.discrete_logarithm(a, b, p));
397    }
398    let br = BarrettReduction::<u128>::new(n as _);
399    if p == 2 {
400        if e >= 3 {
401            if a % 4 == 1 && b % 4 != 1 {
402                return None;
403            }
404            let aa = if a % 4 == 1 { a } else { n - a };
405            let bb = if b % 4 == 1 { b } else { n - b };
406            let g = 5;
407            let ord = n / 4;
408            let x = pohlig_hellman_prime_power_order(g, aa, n, p, e - 2)?;
409            let y = pohlig_hellman_prime_power_order(g, bb, n, p, e - 2)?;
410            let t = solve_linear_congruence(x, y, ord)?;
411            match (a % 4 == 1, b % 4 == 1) {
412                (true, true) => Some(t),
413                (false, true) if t.0 % 2 == 0 => Some((t.0, lcm(t.1, 2))),
414                (false, false) if t.0 % 2 == 1 => Some((t.0, lcm(t.1, 2))),
415                (false, false) if a == b => Some((1, lcm(t.1, 2))),
416                _ => None,
417            }
418        } else if a == 1 {
419            if b == 1 { Some((0, 1)) } else { None }
420        } else {
421            assert_eq!(a, 3);
422            if b == 1 {
423                Some((0, 2))
424            } else if b == 3 {
425                Some((1, 2))
426            } else {
427                None
428            }
429        }
430    } else {
431        let ord = n - n / p;
432        let pf_ord = prime_factors(ord);
433        let g = (2..)
434            .find(|&g| check_primitive_root(g, ord, &br, &pf_ord))
435            .unwrap();
436        let mut pf_p = prime_factors(p - 1);
437        pf_p.push((p, e - 1));
438        let mut abm = vec![];
439        for (q, c) in pf_p {
440            let m = q.pow(c);
441            let d = ord / m;
442            let gg = pow(g, d, &br);
443            let aa = pow(a, d, &br);
444            let bb = pow(b, d, &br);
445            let x = pohlig_hellman_prime_power_order(gg, aa, n, q, c)?;
446            let y = pohlig_hellman_prime_power_order(gg, bb, n, q, c)?;
447            abm.push((x, y, m));
448        }
449        solve_linear_congruences(abm)
450    }
451}

Source

pub const fn get_mod(&self) -> u128

Source

pub const fn div_rem(&self, a: u128) -> (u128, u128)

Source

pub const fn div(&self, a: u128) -> u128

Source

pub const fn rem(&self, a: u128) -> u128

Examples found in repository ?

crates/competitive/src/num/mint/mint_basic.rs (line 151)

150    fn rem(x: u128) -> u128 {
151        DYN_MODULUS_U64.with(|cell| unsafe { (*cell.get()).rem(x) })
152    }

More examples

Hide additional examples

crates/competitive/src/math/discrete_logarithm.rs (line 12)

7fn pow(x: u64, mut y: u64, br: &BarrettReduction<u128>) -> u64 {
8    let mut x = x as u128;
9    let mut z: u128 = 1;
10    while y > 0 {
11        if y & 1 == 1 {
12            z = br.rem(z * x);
13        }
14        x = br.rem(x * x);
15        y >>= 1;
16    }
17    z as u64
18}
19
20fn solve_linear_congruence(a: u64, b: u64, m: u64) -> Option<(u64, u64)> {
21    let g = gcd(a, m);
22    if b % g != 0 {
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 % g != 0 {
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 u64;
70        self.primes.reserve(lim);
71        let primes = self.primes.primes_lte(lim);
72        while self.br_primes.len() < primes.len() {
73            let br = BarrettReduction::<u64>::new(primes[self.br_primes.len()]);
74            self.br_primes.push(br);
75        }
76        let br_primes = &self.br_primes[..primes.len()];
77        self.ic
78            .entry(p)
79            .or_insert_with(|| IndexCalculusWithPrimitiveRoot::new(p, br_primes))
80            .discrete_logarithm(a, b, br_primes)
81    }
82}
83
84const A: [u32; 150] = [
85    62, 61, 60, 60, 59, 58, 58, 58, 57, 56, 56, 56, 56, 55, 55, 55, 54, 54, 54, 53, 53, 53, 53, 52,
86    52, 52, 52, 52, 52, 51, 50, 50, 50, 50, 49, 49, 49, 48, 48, 48, 48, 48, 47, 47, 47, 47, 47, 47,
87    47, 47, 47, 47, 47, 47, 47, 47, 45, 42, 42, 41, 41, 41, 41, 41, 41, 41, 40, 40, 40, 40, 40, 40,
88    40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 40, 38, 38, 38, 38, 38, 32, 32, 32, 32,
89    32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 32, 31, 31, 31, 31, 31,
90    31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 22, 22, 22, 22,
91    22, 22, 22, 22, 22, 22,
92];
93
94fn factorize_smooth(mut x: u64, row: &mut [u64], br_primes: &[BarrettReduction<u64>]) -> bool {
95    for (j, (&br, r)) in br_primes.iter().zip(row).enumerate() {
96        *r = 0;
97        loop {
98            let (div, rem) = br.div_rem(x);
99            if rem != 0 {
100                break;
101            }
102            *r += 1;
103            x = div;
104        }
105        if j < 150 && x >= (1u64 << A[j]) {
106            break;
107        }
108    }
109    x == 1
110}
111
112#[derive(Debug)]
113struct QdrtPowPrec {
114    br_qdrt: BarrettReduction<u64>,
115    p0: Vec<u64>,
116    p1: Vec<u64>,
117    p2: Vec<u64>,
118    p3: Vec<u64>,
119}
120
121impl QdrtPowPrec {
122    fn new(a: u64, ord: u64, br: &BarrettReduction<u128>) -> Self {
123        let qdrt = (ord as f64).powf(0.25).ceil() as u64;
124        let br_qdrt = BarrettReduction::<u64>::new(qdrt);
125        let mut p0 = Vec::with_capacity(qdrt as usize);
126        let mut p1 = Vec::with_capacity(qdrt as usize);
127        let mut p2 = Vec::with_capacity(qdrt as usize);
128        let mut p3 = Vec::with_capacity(qdrt as usize);
129        let mut acc = 1u64;
130        for _ in 0..qdrt {
131            p0.push(acc);
132            acc = br.rem(acc as u128 * a as u128) as u64;
133        }
134        let a = acc;
135        acc = 1;
136        for _ in 0..qdrt {
137            p1.push(acc);
138            acc = br.rem(acc as u128 * a as u128) as u64;
139        }
140        let a = acc;
141        acc = 1;
142        for _ in 0..qdrt {
143            p2.push(acc);
144            acc = br.rem(acc as u128 * a as u128) as u64;
145        }
146        let a = acc;
147        acc = 1;
148        for _ in 0..qdrt {
149            p3.push(acc);
150            acc = br.rem(acc as u128 * a as u128) as u64;
151        }
152        Self {
153            br_qdrt,
154            p0,
155            p1,
156            p2,
157            p3,
158        }
159    }
160    fn pow(&self, mut k: u64, br: &BarrettReduction<u128>) -> u64 {
161        let (a, b) = self.br_qdrt.div_rem(k);
162        let mut x = self.p0[b as usize];
163        k = a;
164        if k > 0 {
165            let (a, b) = self.br_qdrt.div_rem(k);
166            x = br.rem(x as u128 * self.p1[b as usize] as u128) as u64;
167            k = a;
168        }
169        if k > 0 {
170            let (a, b) = self.br_qdrt.div_rem(k);
171            x = br.rem(x as u128 * self.p2[b as usize] as u128) as u64;
172            k = a;
173        }
174        if k > 0 {
175            let (_, b) = self.br_qdrt.div_rem(k);
176            x = br.rem(x as u128 * self.p3[b as usize] as u128) as u64;
177        }
178        x
179    }
180}
181
182fn index_calculus_for_primitive_root(
183    p: u64,
184    ord: u64,
185    br_primes: &[BarrettReduction<u64>],
186    prec: &QdrtPowPrec,
187) -> Vec<u64> {
188    let br_ord = BarrettReduction::<u128>::new(ord as u128);
189    let mul = |x: u64, y: u64| br_ord.rem(x as u128 * y as u128) as u64;
190    let sub = |x: u64, y: u64| if x < y { x + ord - y } else { x - y };
191
192    let pc = br_primes.len();
193    let mut mat: Vec<Vec<u64>> = vec![];
194    let mut rows: Vec<Vec<u64>> = vec![];
195
196    let mut rng = Xorshift::default();
197    let br = BarrettReduction::<u128>::new(p as u128);
198
199    for i in 0..pc {
200        for ri in 0usize.. {
201            let mut row = vec![0u64; pc + 1];
202            let mut kk = rng.rand(ord - 1) + 1;
203            let mut gkk = prec.pow(kk, &br);
204            let mut k = kk;
205            let mut gk = gkk;
206            while ri >= rows.len() {
207                row[pc] = k;
208                if factorize_smooth(gk, &mut row, br_primes) {
209                    rows.push(row);
210                    break;
211                }
212                if k + kk < ord {
213                    k += kk;
214                    gk = br.rem(gk as u128 * gkk as u128) as u64;
215                } else {
216                    kk = rng.rand(ord - 1) + 1;
217                    gkk = prec.pow(kk, &br);
218                    k = kk;
219                    gk = gkk;
220                }
221            }
222            let row = &mut rows[ri];
223            for j in 0..i {
224                if row[j] != 0 {
225                    let b = mul(modinv(mat[j][j], ord), row[j]);
226                    for (r, a) in row[j..].iter_mut().zip(&mat[j][j..]) {
227                        *r = sub(*r, mul(*a, b));
228                    }
229                }
230                assert_eq!(row[j], 0);
231            }
232            if gcd(row[i], ord) == 1 {
233                let last = rows.len() - 1;
234                rows.swap(ri, last);
235                mat.push(rows.pop().unwrap());
236                break;
237            }
238        }
239    }
240    for i in (0..pc).rev() {
241        for j in i + 1..pc {
242            mat[i][pc] = sub(mat[i][pc], mul(mat[i][j], mat[j][pc]));
243        }
244        mat[i][pc] = mul(mat[i][pc], modinv(mat[i][i], ord));
245    }
246    (0..pc).map(|i| (mat[i][pc])).collect()
247}
248
249#[derive(Debug)]
250struct IndexCalculusWithPrimitiveRoot {
251    p: u64,
252    ord: u64,
253    prec: QdrtPowPrec,
254    coeff: Vec<u64>,
255}
256
257impl IndexCalculusWithPrimitiveRoot {
258    fn new(p: u64, br_primes: &[BarrettReduction<u64>]) -> Self {
259        let ord = p - 1;
260        let g = primitive_root(p);
261        let br = BarrettReduction::<u128>::new(p as u128);
262        let prec = QdrtPowPrec::new(g, ord, &br);
263        let coeff = index_calculus_for_primitive_root(p, ord, br_primes, &prec);
264        Self {
265            p,
266            ord,
267            prec,
268            coeff,
269        }
270    }
271    fn index_calculus(&self, a: u64, br_primes: &[BarrettReduction<u64>]) -> Option<u64> {
272        let p = self.p;
273        let ord = self.ord;
274        let br = BarrettReduction::<u128>::new(p as u128);
275        let a = br.rem(a as _) as u64;
276        if a == 1 {
277            return Some(0);
278        }
279        if p == 2 {
280            return None;
281        }
282
283        let mut rng = Xorshift::new();
284        let mut row = vec![0u64; br_primes.len()];
285        let mut kk = rng.rand(ord - 1) + 1;
286        let mut gkk = self.prec.pow(kk, &br);
287        let mut k = kk;
288        let mut gk = br.rem(gkk as u128 * a as u128) as u64;
289        loop {
290            if factorize_smooth(gk, &mut row, br_primes) {
291                let mut res = ord - k;
292                for (&c, &r) in self.coeff.iter().zip(&row) {
293                    for _ in 0..r {
294                        res += c;
295                        if res >= ord {
296                            res -= ord;
297                        }
298                    }
299                }
300                return Some(res);
301            }
302            if k + kk < ord {
303                k += kk;
304                gk = br.rem(gk as u128 * gkk as u128) as u64;
305            } else {
306                kk = rng.rand(ord - 1) + 1;
307                gkk = self.prec.pow(kk, &br);
308                k = kk;
309                gk = br.rem(gkk as u128 * a as u128) as u64;
310            }
311        }
312    }
313    fn discrete_logarithm(
314        &self,
315        a: u64,
316        b: u64,
317        br_primes: &[BarrettReduction<u64>],
318    ) -> Option<(u64, u64)> {
319        let p = self.p;
320        let ord = self.ord;
321        let br = BarrettReduction::<u128>::new(p as u128);
322        let a = br.rem(a as _) as u64;
323        let b = br.rem(b as _) as u64;
324        if a == 0 {
325            return if b == 0 { Some((1, 1)) } else { None };
326        }
327        if b == 0 {
328            return None;
329        }
330
331        let x = self.index_calculus(a, br_primes)?;
332        let y = self.index_calculus(b, br_primes)?;
333        solve_linear_congruence(x, y, ord)
334    }
335}
336
337thread_local!(
338    static IC: UnsafeCell<IndexCalculus> = UnsafeCell::new(IndexCalculus::new());
339);
340
341pub fn discrete_logarithm_prime_mod(a: u64, b: u64, p: u64) -> Option<u64> {
342    IC.with(|ic| unsafe { &mut *ic.get() }.discrete_logarithm(a, b, p))
343        .map(|t| t.0)
344}
345
346/// a^x ≡ b (mod n), a has order p^e
347fn pohlig_hellman_prime_power_order(a: u64, b: u64, n: u64, p: u64, e: u32) -> Option<u64> {
348    let br = BarrettReduction::<u128>::new(n as u128);
349    let mul = |x: u64, y: u64| br.rem(x as u128 * y as u128) as u64;
350    let block_size = (p as f64).sqrt().ceil() as u64;
351    let mut baby = HashMap::<u64, u64>::new();
352    let g = pow(a, p.pow(e - 1), &br);
353    let mut xj = 1;
354    for j in 0..block_size {
355        baby.entry(xj).or_insert(j);
356        xj = mul(xj, g);
357    }
358    let xi = modinv(xj, n);
359    let mut t = 0u64;
360    for k in 0..e {
361        let mut h = pow(mul(modinv(pow(a, t, &br), n), b), p.pow(e - 1 - k), &br);
362        let mut ok = false;
363        for i in (0..block_size * block_size).step_by(block_size as usize) {
364            if let Some(j) = baby.get(&h) {
365                t += (i + j) * p.pow(k);
366                ok = true;
367                break;
368            }
369            h = mul(h, xi);
370        }
371        if !ok {
372            return None;
373        }
374    }
375    Some(t)
376}

crates/competitive/src/math/primitive_root.rs (line 27)

15pub fn check_primitive_root(
16    g: u64,
17    phi: u64,
18    br: &BarrettReduction<u128>,
19    pf: &[(u64, u32)],
20) -> bool {
21    pf.iter().all(|&(q, _)| {
22        let mut g = g as u128;
23        let mut k = phi / q;
24        let mut r: u128 = 1;
25        while k > 0 {
26            if k & 1 == 1 {
27                r = br.rem(r * g);
28            }
29            g = br.rem(g * g);
30            k >>= 1;
31        }
32        r > 1
33    })
34}

crates/competitive/src/math/prime_factors.rs (line 10)

3fn find_factor(n: u64) -> Option<u64> {
4    let br = BarrettReduction::<u128>::new(n as u128);
5    if miller_rabin_with_br(n, &br) {
6        return None;
7    }
8    let m = 1u64 << ((63 - n.leading_zeros()) / 5);
9    let sub = |x: u64, y: u64| if x > y { x - y } else { y - x };
10    let mul = |x: u64, y: u64| br.rem(x as u128 * y as u128) as u64;
11    for c in 12.. {
12        let f = |x: u64| (br.rem(x as u128 * x as u128) + c) as u64;
13        let (mut x, mut y, mut r, mut g, mut k, mut ys) = (0, 2, 1, 1, 0, 0);
14        while g == 1 {
15            x = y;
16            for _ in 0..r {
17                y = f(y);
18            }
19            while r > k && g == 1 {
20                ys = y;
21                let mut q = 1;
22                for _ in 0..m.min(r - k) {
23                    y = f(y);
24                    q = mul(q, sub(x, y));
25                }
26                g = gcd(q, n);
27                k += m;
28            }
29            r <<= 1;
30        }
31        if g == n {
32            g = 1;
33            while g == 1 {
34                ys = f(ys);
35                g = gcd(sub(x, ys), n);
36            }
37        }
38        if g < n {
39            return Some(g);
40        }
41    }
42    unreachable!();
43}