vers_vecs/bit_vec/fast_rs_vec/
select.rs1use crate::bit_vec::fast_rs_vec::{BLOCK_SIZE, SELECT_BLOCK_SIZE, SUPER_BLOCK_SIZE};
4use crate::bit_vec::WORD_SIZE;
5use crate::util::pdep::Pdep;
6use crate::util::unroll_n;
7
8const BLOCKS_PER_SUPERBLOCK: usize = 16;
11
12impl super::RsVec {
13 #[must_use]
19 #[allow(clippy::assertions_on_constants)]
20 pub fn select0(&self, mut rank: usize) -> usize {
21 if rank >= self.rank0 {
22 return self.len;
23 }
24
25 let mut super_block = self.select_blocks[rank / SELECT_BLOCK_SIZE].index_0;
26
27 if self.super_blocks.len() > (super_block + 1)
28 && self.super_blocks[super_block + 1].zeros <= rank
29 {
30 super_block = self.search_super_block0(super_block, rank);
31 }
32
33 rank -= self.super_blocks[super_block].zeros;
34
35 let mut block_index = super_block * (SUPER_BLOCK_SIZE / BLOCK_SIZE);
36 self.search_block0(rank, &mut block_index);
37
38 rank -= self.blocks[block_index].zeros as usize;
39
40 self.search_word_in_block0(rank, block_index)
41 }
42
43 #[cfg(all(
52 feature = "simd",
53 target_arch = "x86_64",
54 target_feature = "avx",
55 target_feature = "avx512vl",
56 target_feature = "avx512bw",
57 ))]
58 #[inline(always)]
59 pub(super) fn search_block0(&self, rank: usize, block_index: &mut usize) {
60 use std::arch::x86_64::{_mm256_cmpgt_epu16_mask, _mm256_loadu_epi16, _mm256_set1_epi16};
61
62 if self.blocks.len() > *block_index + (SUPER_BLOCK_SIZE / BLOCK_SIZE) {
63 debug_assert!(
64 SUPER_BLOCK_SIZE / BLOCK_SIZE == BLOCKS_PER_SUPERBLOCK,
65 "change unroll constant to {}",
66 64 - (SUPER_BLOCK_SIZE / BLOCK_SIZE).leading_zeros() - 1
67 );
68
69 unsafe {
70 let blocks = _mm256_loadu_epi16(self.blocks[*block_index..].as_ptr() as *const i16);
71 let ranks = _mm256_set1_epi16(rank as i16);
72 let mask = _mm256_cmpgt_epu16_mask(blocks, ranks);
73
74 debug_assert!(
75 mask.count_zeros() > 0,
76 "first block should always be zero, but still claims to be greater than rank"
77 );
78 *block_index += mask.count_zeros() as usize - 1;
79 }
80 } else {
81 self.search_block0_naive(rank, block_index)
82 }
83 }
84
85 #[cfg(not(all(
89 feature = "simd",
90 target_arch = "x86_64",
91 target_feature = "avx",
92 target_feature = "avx512vl",
93 target_feature = "avx512bw",
94 )))]
95 #[inline(always)]
96 pub(super) fn search_block0(&self, rank: usize, block_index: &mut usize) {
97 self.search_block0_naive(rank, block_index);
98 }
99
100 #[inline(always)]
101 fn search_block0_naive(&self, rank: usize, block_index: &mut usize) {
102 debug_assert!(
107 SUPER_BLOCK_SIZE / BLOCK_SIZE == BLOCKS_PER_SUPERBLOCK,
108 "change unroll constant to {}",
109 64 - (SUPER_BLOCK_SIZE / BLOCK_SIZE).leading_zeros() - 1
110 );
111 unroll_n!(4,
112 |boundary = { (SUPER_BLOCK_SIZE / BLOCK_SIZE) / 2}|
113 if self.blocks.len() > *block_index + boundary && rank >= self.blocks[*block_index + boundary].zeros as usize {
115 *block_index += boundary;
116 },
117 boundary /= 2);
118 }
119
120 #[inline(always)]
129 pub(super) fn search_word_in_block0(&self, mut rank: usize, block_index: usize) -> usize {
130 let mut index_counter = 0;
135 debug_assert!(BLOCK_SIZE / WORD_SIZE == 8, "change unroll constant");
136 unroll_n!(7, |n = {0}| {
137 let word = self.data[block_index * BLOCK_SIZE / WORD_SIZE + n];
138 if (word.count_zeros() as usize) <= rank {
139 rank -= word.count_zeros() as usize;
140 index_counter += WORD_SIZE;
141 } else {
142 return block_index * BLOCK_SIZE
143 + index_counter
144 + (1 << rank).pdep(!word).trailing_zeros() as usize;
145 }
146 }, n += 1);
147
148 block_index * BLOCK_SIZE
151 + index_counter
152 + (1 << rank)
153 .pdep(!self.data[block_index * BLOCK_SIZE / WORD_SIZE + 7])
154 .trailing_zeros() as usize
155 }
156
157 #[inline(always)]
165 pub(super) fn search_super_block0(&self, mut super_block: usize, rank: usize) -> usize {
166 let mut upper_bound = self.select_blocks[rank / SELECT_BLOCK_SIZE + 1].index_0;
167
168 while upper_bound - super_block > 8 {
169 let middle = super_block + ((upper_bound - super_block) >> 1);
170 if self.super_blocks[middle].zeros <= rank {
172 super_block = middle;
173 } else {
174 upper_bound = middle;
175 }
176 }
177
178 while self.super_blocks.len() > (super_block + 1)
180 && self.super_blocks[super_block + 1].zeros <= rank
181 {
182 super_block += 1;
183 }
184
185 debug_assert!(super_block <= upper_bound, "calculated the upper bound to be {} (initially {}) but the super block was found at {}", upper_bound, self.select_blocks[rank / SELECT_BLOCK_SIZE + 1].index_0, super_block);
186
187 super_block
188 }
189
190 #[must_use]
196 #[allow(clippy::assertions_on_constants)]
197 pub fn select1(&self, mut rank: usize) -> usize {
198 if rank >= self.rank1 {
199 return self.len;
200 }
201
202 let mut super_block = self.select_blocks[rank / SELECT_BLOCK_SIZE].index_1;
203
204 if self.super_blocks.len() > (super_block + 1)
205 && ((super_block + 1) * SUPER_BLOCK_SIZE - self.super_blocks[super_block + 1].zeros)
206 <= rank
207 {
208 super_block = self.search_super_block1(super_block, rank);
209 }
210
211 rank -= (super_block) * SUPER_BLOCK_SIZE - self.super_blocks[super_block].zeros;
212
213 let block_at_super_block = super_block * (SUPER_BLOCK_SIZE / BLOCK_SIZE);
216 let mut block_index = block_at_super_block;
217 self.search_block1(rank, block_at_super_block, &mut block_index);
218
219 rank -= (block_index - block_at_super_block) * BLOCK_SIZE
220 - self.blocks[block_index].zeros as usize;
221
222 self.search_word_in_block1(rank, block_index)
223 }
224
225 #[cfg(all(
234 feature = "simd",
235 target_arch = "x86_64",
236 target_feature = "avx",
237 target_feature = "avx2",
238 target_feature = "avx512vl",
239 target_feature = "avx512bw",
240 ))]
241 #[inline(always)]
242 pub(super) fn search_block1(
243 &self,
244 rank: usize,
245 block_at_super_block: usize,
246 block_index: &mut usize,
247 ) {
248 use std::arch::x86_64::{
249 _mm256_cmpgt_epu16_mask, _mm256_loadu_epi16, _mm256_set1_epi16, _mm256_set_epi16,
250 _mm256_sub_epi16,
251 };
252
253 if self.blocks.len() > *block_index + BLOCKS_PER_SUPERBLOCK {
254 debug_assert!(
255 SUPER_BLOCK_SIZE / BLOCK_SIZE == BLOCKS_PER_SUPERBLOCK,
256 "change unroll constant to {}",
257 64 - (SUPER_BLOCK_SIZE / BLOCK_SIZE).leading_zeros() - 1
258 );
259
260 unsafe {
261 let bit_nums = _mm256_set_epi16(
262 (15 * BLOCK_SIZE) as i16,
263 (14 * BLOCK_SIZE) as i16,
264 (13 * BLOCK_SIZE) as i16,
265 (12 * BLOCK_SIZE) as i16,
266 (11 * BLOCK_SIZE) as i16,
267 (10 * BLOCK_SIZE) as i16,
268 (9 * BLOCK_SIZE) as i16,
269 (8 * BLOCK_SIZE) as i16,
270 (7 * BLOCK_SIZE) as i16,
271 (6 * BLOCK_SIZE) as i16,
272 (5 * BLOCK_SIZE) as i16,
273 (4 * BLOCK_SIZE) as i16,
274 (3 * BLOCK_SIZE) as i16,
275 (2 * BLOCK_SIZE) as i16,
276 (1 * BLOCK_SIZE) as i16,
277 (0 * BLOCK_SIZE) as i16,
278 );
279
280 let blocks = _mm256_loadu_epi16(self.blocks[*block_index..].as_ptr() as *const i16);
281 let ones = _mm256_sub_epi16(bit_nums, blocks);
282
283 let ranks = _mm256_set1_epi16(rank as i16);
284 let mask = _mm256_cmpgt_epu16_mask(ones, ranks);
285
286 debug_assert!(
287 mask.count_zeros() > 0,
288 "first block should always be zero, but still claims to be greater than rank"
289 );
290 *block_index += mask.count_zeros() as usize - 1;
291 }
292 } else {
293 self.search_block1_naive(rank, block_at_super_block, block_index)
294 }
295 }
296
297 #[cfg(not(all(
301 feature = "simd",
302 target_arch = "x86_64",
303 target_feature = "avx",
304 target_feature = "avx2",
305 target_feature = "avx512vl",
306 target_feature = "avx512bw",
307 )))]
308 #[inline(always)]
309 pub(super) fn search_block1(
310 &self,
311 rank: usize,
312 block_at_super_block: usize,
313 block_index: &mut usize,
314 ) {
315 self.search_block1_naive(rank, block_at_super_block, block_index);
316 }
317
318 #[inline(always)]
319 fn search_block1_naive(
320 &self,
321 rank: usize,
322 block_at_super_block: usize,
323 block_index: &mut usize,
324 ) {
325 debug_assert!(
330 SUPER_BLOCK_SIZE / BLOCK_SIZE == BLOCKS_PER_SUPERBLOCK,
331 "change unroll constant to {}",
332 64 - (SUPER_BLOCK_SIZE / BLOCK_SIZE).leading_zeros() - 1
333 );
334 unroll_n!(4,
335 |boundary = { (SUPER_BLOCK_SIZE / BLOCK_SIZE) / 2}|
336 if self.blocks.len() > *block_index + boundary && rank >= (*block_index + boundary - block_at_super_block) * BLOCK_SIZE - self.blocks[*block_index + boundary].zeros as usize {
338 *block_index += boundary;
339 },
340 boundary /= 2);
341 }
342
343 #[inline(always)]
352 pub(super) fn search_word_in_block1(&self, mut rank: usize, block_index: usize) -> usize {
353 let mut index_counter = 0;
358 debug_assert!(BLOCK_SIZE / WORD_SIZE == 8, "change unroll constant");
359 unroll_n!(7, |n = {0}| {
360 let word = self.data[block_index * BLOCK_SIZE / WORD_SIZE + n];
361 if (word.count_ones() as usize) <= rank {
362 rank -= word.count_ones() as usize;
363 index_counter += WORD_SIZE;
364 } else {
365 return block_index * BLOCK_SIZE
366 + index_counter
367 + (1 << rank).pdep(word).trailing_zeros() as usize;
368 }
369 }, n += 1);
370
371 block_index * BLOCK_SIZE
374 + index_counter
375 + (1 << rank)
376 .pdep(self.data[block_index * BLOCK_SIZE / WORD_SIZE + 7])
377 .trailing_zeros() as usize
378 }
379
380 #[inline(always)]
389 pub(super) fn search_super_block1(&self, mut super_block: usize, rank: usize) -> usize {
390 let mut upper_bound = self.select_blocks[rank / SELECT_BLOCK_SIZE + 1].index_1;
391
392 while upper_bound - super_block > 8 {
394 let middle = super_block + ((upper_bound - super_block) >> 1);
395 if (middle * SUPER_BLOCK_SIZE - self.super_blocks[middle].zeros) <= rank {
397 super_block = middle;
398 } else {
399 upper_bound = middle;
400 }
401 }
402 while self.super_blocks.len() > (super_block + 1)
404 && ((super_block + 1) * SUPER_BLOCK_SIZE - self.super_blocks[super_block + 1].zeros)
405 <= rank
406 {
407 super_block += 1;
408 }
409
410 debug_assert!(super_block <= upper_bound, "calculated the upper bound to be {} (initially {}) but the super block was found at {}", upper_bound, self.select_blocks[rank / SELECT_BLOCK_SIZE + 1].index_1, super_block);
411
412 super_block
413 }
414
415 #[must_use]
438 pub fn successor0(&self, pos: usize) -> Option<u64> {
439 if self.is_empty() {
440 return None;
441 }
442
443 let rank = self.rank0(pos);
444 let mut rank = if self.get(pos)? == 0 { rank + 1 } else { rank };
445
446 if rank >= self.rank0 {
447 return None;
448 }
449
450 let mut block_idx = pos / BLOCK_SIZE;
451 let super_block_idx = pos / SUPER_BLOCK_SIZE;
452
453 if self.super_blocks.len() > (super_block_idx + 1)
454 && self.super_blocks[super_block_idx + 1].zeros > rank
455 {
456 rank -= self.super_blocks[super_block_idx].zeros;
457
458 if block_idx % (BLOCKS_PER_SUPERBLOCK) == (BLOCKS_PER_SUPERBLOCK - 1)
460 || self.blocks[block_idx + 1].zeros as usize > rank
461 {
462 rank -= self.blocks[block_idx].zeros as usize;
463 return Some(self.search_word_in_block0(rank, block_idx) as u64);
464 }
465
466 block_idx = super_block_idx * (BLOCKS_PER_SUPERBLOCK);
467 self.search_block0(rank, &mut block_idx);
468
469 rank -= self.blocks[block_idx].zeros as usize;
470
471 Some(self.search_word_in_block0(rank, block_idx) as u64)
472 } else {
473 Some(self.select0(rank) as u64)
474 }
475 }
476
477 #[must_use]
500 pub fn successor1(&self, pos: usize) -> Option<u64> {
501 if self.is_empty() {
502 return None;
503 }
504
505 let rank = self.rank1(pos);
506 let mut rank = if self.get(pos)? == 1 { rank + 1 } else { rank };
507
508 if rank >= self.rank1 {
509 return None;
510 }
511
512 let mut block_idx = pos / BLOCK_SIZE;
513 let super_block_idx = pos / SUPER_BLOCK_SIZE;
514
515 if self.super_blocks.len() > (super_block_idx + 1)
516 && (super_block_idx + 1) * SUPER_BLOCK_SIZE
517 - self.super_blocks[super_block_idx + 1].zeros
518 > rank
519 {
520 let super_block_ones =
521 (super_block_idx * SUPER_BLOCK_SIZE) - self.super_blocks[super_block_idx].zeros;
522
523 rank -= super_block_ones;
524
525 let block_at_super_block = super_block_idx * (BLOCKS_PER_SUPERBLOCK);
526 if block_idx % (BLOCKS_PER_SUPERBLOCK) == BLOCKS_PER_SUPERBLOCK - 1
528 || (block_idx + 1 - block_at_super_block) * BLOCK_SIZE
529 - self.blocks[block_idx + 1].zeros as usize
530 > rank
531 {
532 let block_ones = (block_idx - block_at_super_block) * BLOCK_SIZE
533 - self.blocks[block_idx].zeros as usize;
534 rank -= block_ones;
535 return Some(self.search_word_in_block1(rank, block_idx) as u64);
536 }
537
538 block_idx = block_at_super_block;
539 self.search_block1(rank, block_at_super_block, &mut block_idx);
540 rank -= (block_idx - block_at_super_block) * BLOCK_SIZE
541 - self.blocks[block_idx].zeros as usize;
542
543 Some(self.search_word_in_block1(rank, block_idx) as u64)
544 } else {
545 Some(self.select1(rank) as u64)
546 }
547 }
548
549 #[must_use]
572 pub fn predecessor0(&self, pos: usize) -> Option<u64> {
573 if self.is_empty() {
574 return None;
575 }
576
577 let mut rank = self.rank0(pos).checked_sub(1)?;
578
579 let mut block_idx = pos / BLOCK_SIZE;
580 let super_block_idx = pos / SUPER_BLOCK_SIZE;
581
582 if self.super_blocks[super_block_idx].zeros < rank {
583 rank -= self.super_blocks[super_block_idx].zeros;
584
585 if (self.blocks[block_idx].zeros as usize) < rank {
587 rank -= self.blocks[block_idx].zeros as usize;
588 return Some(self.search_word_in_block0(rank, block_idx) as u64);
589 }
590
591 block_idx = super_block_idx * (BLOCKS_PER_SUPERBLOCK);
592 self.search_block0(rank, &mut block_idx);
593
594 rank -= self.blocks[block_idx].zeros as usize;
595
596 Some(self.search_word_in_block0(rank, block_idx) as u64)
597 } else {
598 Some(self.select0(rank) as u64)
599 }
600 }
601
602 #[must_use]
625 pub fn predecessor1(&self, pos: usize) -> Option<u64> {
626 if self.is_empty() {
627 return None;
628 }
629
630 let mut rank = self.rank1(pos).checked_sub(1)?;
631
632 let mut block_idx = pos / BLOCK_SIZE;
633 let super_block_idx = pos / SUPER_BLOCK_SIZE;
634
635 let super_block_ones =
636 (super_block_idx * SUPER_BLOCK_SIZE) - self.super_blocks[super_block_idx].zeros;
637
638 if super_block_ones < rank {
639 rank -= super_block_ones;
640
641 let block_at_super_block = super_block_idx * (BLOCKS_PER_SUPERBLOCK);
642 let block_ones = (block_idx - block_at_super_block) * BLOCK_SIZE
643 - self.blocks[block_idx].zeros as usize;
644 if block_ones < rank {
646 rank -= block_ones;
647 return Some(self.search_word_in_block1(rank, block_idx) as u64);
648 }
649
650 block_idx = block_at_super_block;
651 self.search_block1(rank, block_at_super_block, &mut block_idx);
652 rank -= (block_idx - block_at_super_block) * BLOCK_SIZE
653 - self.blocks[block_idx].zeros as usize;
654
655 Some(self.search_word_in_block1(rank, block_idx) as u64)
656 } else {
657 Some(self.select1(rank) as u64)
658 }
659 }
660}