const fn to_ordinal(x: u32) -> u32Examples found in repository?
crates/competitive/src/math/prime_list.rs (line 94)
92const fn sqrt_bits() -> [u64; SQRT_THRESHOLD as usize / 128] {
93 let mut result = [!0; SQRT_THRESHOLD as usize / 128];
94 let ordinal = to_ordinal(1) as usize;
95 result[ordinal / 64] &= !(1 << (ordinal % 64));
96 let mut i = RESIDUES[1] as u32;
97 while to_ordinal(i * i) < SQRT_THRESHOLD / 2 {
98 let ordinal = to_ordinal(i) as usize;
99 if result[ordinal / 64] >> (ordinal % 64) & 1 != 0 {
100 let mut k = i;
101 while to_ordinal(i * k) < SQRT_THRESHOLD / 2 {
102 let ordinal = to_ordinal(i * k) as usize;
103 result[ordinal / 64] &= !(1 << (ordinal % 64));
104 k += ADDITIONS[(k % PERIOD) as usize] as u32;
105 }
106 }
107 i += ADDITIONS[(i % PERIOD) as usize] as u32;
108 }
109 result
110}
111
112const SQRT_BITS: [u64; SQRT_THRESHOLD as usize / 128] = sqrt_bits();
113
114const fn count_sqrt_primes() -> usize {
115 let mut result = 0;
116 let mut i = RESIDUES[1] as u32;
117 while i < SQRT_THRESHOLD {
118 let ordinal = to_ordinal(i) as usize;
119 result += (SQRT_BITS[ordinal / 64] >> (ordinal % 64) & 1) as usize;
120 i += ADDITIONS[(i % PERIOD) as usize] as u32;
121 }
122 result
123}
124
125const SQRT_PRIME_COUNT: usize = count_sqrt_primes();
126
127const fn sqrt_primes() -> [u32; SQRT_PRIME_COUNT] {
128 let mut result = [0; SQRT_PRIME_COUNT];
129 let mut i = RESIDUES[1] as u32;
130 let mut j = 0;
131 while i < SQRT_THRESHOLD {
132 let ordinal = to_ordinal(i) as usize;
133 if SQRT_BITS[ordinal / 64] >> (ordinal % 64) & 1 != 0 {
134 result[j] = i;
135 j += 1;
136 }
137 i += ADDITIONS[(i % PERIOD) as usize] as u32;
138 }
139 result
140}
141
142static SQRT_PRIMES: [u32; SQRT_PRIME_COUNT] = sqrt_primes();
143
144struct Wheel {
145 mask: Vec<u64>,
146 product: u32,
147}
148
149impl Wheel {
150 fn new(primes: &[u32], product: u32) -> Self {
151 let mut mask = vec![!0; to_ordinal(product) as usize / 64];
152 for &p in primes {
153 let mut k = 1;
154 while p * k < product {
155 let ordinal = to_ordinal(p * k) as usize;
156 mask[ordinal / 64] &= !(1 << (ordinal % 64));
157 k += ADDITIONS[(k % PERIOD) as usize] as u32;
158 }
159 }
160 Self { mask, product }
161 }
162}
163
164fn make_wheels() -> (Vec<Wheel>, usize) {
165 const MAX_WHEEL_SIZE: u32 = 1 << 20;
166 const BASE: u32 = (PERIOD * 64) >> (WHEEL_PRIMES.len() - 2);
167 let mut product = BASE;
168 let mut current = vec![];
169 let mut wheels = vec![];
170 for (i, &p) in SQRT_PRIMES.iter().enumerate() {
171 if product * p > MAX_WHEEL_SIZE {
172 wheels.push(Wheel::new(¤t, product));
173 current.clear();
174 current.push(p);
175 product = BASE * p;
176 if product > MAX_WHEEL_SIZE {
177 return (wheels, i);
178 }
179 } else {
180 current.push(p);
181 product *= p;
182 }
183 }
184 unreachable!()
185}
186
187fn sieve_dense(bits: &mut [u64], l: u32, r: u32, wheel: &Wheel) {
188 let mut left = l as usize / 64;
189 let right = (r as usize).div_ceil(64);
190 while left + wheel.mask.len() <= right {
191 for (value, &mask) in bits[left..left + wheel.mask.len()]
192 .iter_mut()
193 .zip(&wheel.mask)
194 {
195 *value &= mask;
196 }
197 left += wheel.mask.len();
198 }
199 for (value, &mask) in bits[left..right].iter_mut().zip(&wheel.mask) {
200 *value &= mask;
201 }
202}
203
204fn ordinal_steps() -> Vec<[u32; COPRIME * 2]> {
205 SQRT_PRIMES
206 .iter()
207 .map(|&p| {
208 let mut result = [0; COPRIME * 2];
209 let mut last = to_ordinal(p);
210 for i in 0..COPRIME {
211 let next = to_ordinal(p * (RESIDUES[i] as u32 + GAPS[i] as u32));
212 result[i] = next - last;
213 result[i + COPRIME] = next - last;
214 last = next;
215 }
216 result
217 })
218 .collect()
219}
220
221fn sieve_sparse(
222 bits: &mut [u64],
223 mut left: u32,
224 right: u32,
225 prime_index: usize,
226 mut state: u8,
227 steps: &[[u32; COPRIME * 2]],
228) -> (u32, u8) {
229 let p = SQRT_PRIMES[prime_index];
230 while left + p * COPRIME as u32 <= right {
231 for _ in 0..COPRIME {
232 let ordinal = left as usize;
233 bits[ordinal / 64] &= !(1 << (ordinal % 64));
234 left += steps[prime_index][state as usize];
235 state += 1;
236 }
237 state -= COPRIME as u8;
238 }
239 while left < right {
240 let ordinal = left as usize;
241 bits[ordinal / 64] &= !(1 << (ordinal % 64));
242 left += steps[prime_index][state as usize];
243 state += 1;
244 }
245 if state >= COPRIME as u8 {
246 state -= COPRIME as u8;
247 }
248 (left, state)
249}
250
251#[derive(Debug, Clone)]
252pub struct PrimeList {
253 bits: Vec<u64>,
254 bit_len: usize,
255 max_n: u32,
256 prime_count: usize,
257}
258
259impl Default for PrimeList {
260 fn default() -> Self {
261 Self {
262 bits: vec![],
263 bit_len: 0,
264 max_n: 1,
265 prime_count: 0,
266 }
267 }
268}
269
270impl PrimeList {
271 pub fn new(max_n: u32) -> Self {
272 let mut self_: Self = Default::default();
273 self_.reserve(max_n);
274 self_
275 }
276 pub fn primes(&self) -> PrimeListIter<'_> {
277 self.primes_lte(self.max_n)
278 }
279 pub fn len(&self) -> usize {
280 self.prime_count
281 }
282 pub fn is_empty(&self) -> bool {
283 self.prime_count == 0
284 }
285 pub fn primes_lte(&self, n: u32) -> PrimeListIter<'_> {
286 assert!(n <= self.max_n, "expected `n={} <= {}`", n, self.max_n);
287 let bit_len = to_ordinal(n.saturating_add(1)) as usize;
288 let words = &self.bits[..self.bits.len().min(bit_len.div_ceil(64))];
289 let last_mask = if bit_len.is_multiple_of(64) {
290 !0
291 } else {
292 (1 << (bit_len % 64)) - 1
293 };
294 let (front_word, middle_words, back_word) = match words {
295 [] => (0, words, 0),
296 [word] => (*word & last_mask, &words[1..], 0),
297 [front, middle @ .., back] => (*front, middle, *back & last_mask),
298 };
299 let back_word_ordinal = words.len().saturating_sub(1) as u32 * 64;
300 PrimeListIter {
301 wheel_indices: 0..WHEEL_PRIMES.partition_point(|&p| p <= n) as u8,
302 middle_words,
303 front_word_base: 0,
304 front_word_phase: 0,
305 front_word,
306 back_word_base: back_word_ordinal / COPRIME as u32 * PERIOD,
307 back_word_phase: (back_word_ordinal % COPRIME as u32) as u8,
308 back_word,
309 }
310 }
311 pub fn is_prime(&self, n: u32) -> bool {
312 assert!(n <= self.max_n, "expected `n={} <= {}`", n, self.max_n);
313 if WHEEL_PRIMES.contains(&n) {
314 true
315 } else if !coprime_to_wheel(n) {
316 false
317 } else {
318 let ordinal = to_ordinal(n) as usize;
319 ordinal < self.bit_len && self.bits[ordinal / 64] >> (ordinal % 64) & 1 != 0
320 }
321 }
322 pub fn trial_division(&self, n: u64) -> PrimeListTrialDivision<'_> {
323 let bound = u64::from(self.max_n).pow(2);
324 assert!(n <= bound, "expected `n={} <= {}`", n, bound);
325 PrimeListTrialDivision {
326 primes: self.primes(),
327 n,
328 }
329 }
330 pub fn prime_factors(&self, n: u64) -> Vec<(u64, u32)> {
331 self.trial_division(n).collect()
332 }
333 pub fn count_divisors(&self, n: u64) -> u64 {
334 let mut divisor_cnt = 1u64;
335 for (_, cnt) in self.trial_division(n) {
336 divisor_cnt *= cnt as u64 + 1;
337 }
338 divisor_cnt
339 }
340 pub fn divisors(&self, n: u64) -> Vec<u64> {
341 let mut d = vec![1u64];
342 for (p, c) in self.trial_division(n) {
343 let k = d.len();
344 let mut acc = 1;
345 for _ in 0..c {
346 acc *= p;
347 for i in 0..k {
348 d.push(d[i] * acc);
349 }
350 }
351 }
352 d.sort_unstable();
353 d
354 }
355 /// Extends the prime list up to `max_n`.
356 pub fn reserve(&mut self, max_n: u32) {
357 if max_n <= self.max_n || max_n < 2 {
358 return;
359 }
360 let limit = max_n.saturating_add(1);
361 self.bit_len = to_ordinal(limit) as usize;
362 if limit <= SQRT_THRESHOLD {
363 self.bits = SQRT_BITS[..self.bit_len.div_ceil(64)].to_vec();
364 } else {
365 self.bits = vec![!0; self.bit_len.div_ceil(64)];
366 let (wheels, medium_primes_begin) = make_wheels();
367 const DENSE_BLOCK: u32 = 1 << 25;
368 for start in (0..limit).step_by(DENSE_BLOCK as usize) {
369 let right = start.saturating_add(DENSE_BLOCK).min(limit);
370 for wheel in &wheels {
371 let left = start / wheel.product * wheel.product;
372 sieve_dense(&mut self.bits, to_ordinal(left), to_ordinal(right), wheel);
373 }
374 }
375
376 let steps = ordinal_steps();
377 let mut positions: Vec<_> = SQRT_PRIMES.iter().map(|&p| to_ordinal(p * p)).collect();
378 let mut states: Vec<_> = SQRT_PRIMES
379 .iter()
380 .map(|&p| STATES[(p % PERIOD) as usize])
381 .collect();
382 const SPARSE_BLOCK: u32 = 1 << 22;
383 for start in (0..limit).step_by(SPARSE_BLOCK as usize) {
384 let right = to_ordinal(start.saturating_add(SPARSE_BLOCK).min(limit));
385 for i in medium_primes_begin..SQRT_PRIME_COUNT {
386 (positions[i], states[i]) =
387 sieve_sparse(&mut self.bits, positions[i], right, i, states[i], &steps);
388 }
389 }
390 for (value, &sqrt_bits) in self.bits.iter_mut().zip(&SQRT_BITS) {
391 *value = sqrt_bits;
392 }
393 }
394
395 self.prime_count = WHEEL_PRIMES.partition_point(|&p| p <= max_n);
396 if let Some((&last, rest)) = self.bits.split_last() {
397 self.prime_count += rest
398 .iter()
399 .map(|word| word.count_ones() as usize)
400 .sum::<usize>();
401 let last_mask = if self.bit_len.is_multiple_of(64) {
402 !0
403 } else {
404 (1 << (self.bit_len % 64)) - 1
405 };
406 self.prime_count += (last & last_mask).count_ones() as usize;
407 }
408 self.max_n = max_n;
409 }