1#[cfg(any(feature = "alloc", test))]
4mod rmap;
5#[cfg(test)]
6mod tests;
7pub(super) mod types;
8
9use crate::chiapos::Seed;
10use crate::chiapos::constants::{PARAM_B, PARAM_BC, PARAM_C, PARAM_EXT, PARAM_M};
11#[cfg(feature = "alloc")]
12use crate::chiapos::table::rmap::Rmap;
13use crate::chiapos::table::types::{Metadata, X, Y};
14#[cfg(feature = "alloc")]
15use crate::chiapos::table::types::{Position, R};
16use ab_chacha8::{ChaCha8Block, ChaCha8State};
17#[cfg(feature = "alloc")]
18use alloc::boxed::Box;
19#[cfg(feature = "alloc")]
20use alloc::sync::Arc;
21#[cfg(feature = "alloc")]
22use alloc::vec;
23#[cfg(feature = "alloc")]
24use alloc::vec::Vec;
25#[cfg(feature = "alloc")]
26use chacha20::cipher::{Iv, KeyIvInit, StreamCipher};
27#[cfg(feature = "alloc")]
28use chacha20::{ChaCha8, Key};
29use core::array;
30#[cfg(feature = "parallel")]
31use core::cell::SyncUnsafeCell;
32#[cfg(feature = "alloc")]
33use core::mem;
34#[cfg(feature = "alloc")]
35use core::mem::MaybeUninit;
36#[cfg(any(feature = "alloc", test))]
37use core::simd::prelude::*;
38#[cfg(feature = "parallel")]
39use core::sync::atomic::{AtomicUsize, Ordering};
40#[cfg(feature = "alloc")]
41use derive_more::Deref;
42#[cfg(any(feature = "alloc", test))]
43use seq_macro::seq;
44use subspace_core_primitives::hashes::blake3_hash;
45
46#[cfg(any(feature = "alloc", test))]
47const COMPUTE_F1_SIMD_FACTOR: usize = 8;
48#[cfg(any(feature = "alloc", test))]
49const COMPUTE_FN_SIMD_FACTOR: usize = 16;
50#[allow(dead_code, reason = "unused when crate is compiled separately")]
51const MAX_BUCKET_SIZE: usize = 512;
52#[cfg(feature = "alloc")]
53const BUCKET_SIZE_UPPER_BOUND_SECURITY_BITS: u8 = 128;
54#[allow(dead_code, reason = "unused when crate is compiled separately")]
61pub(crate) const REDUCED_BUCKET_SIZE: usize = MAX_BUCKET_SIZE;
62#[allow(dead_code, reason = "unused when crate is compiled separately")]
67const REDUCED_MATCHES_COUNT: usize = MAX_BUCKET_SIZE;
68#[cfg(feature = "parallel")]
69const CACHE_LINE_SIZE: usize = 64;
70
71const _: () = {
72 debug_assert!(REDUCED_BUCKET_SIZE <= MAX_BUCKET_SIZE);
73 debug_assert!(REDUCED_MATCHES_COUNT <= MAX_BUCKET_SIZE);
74};
75
76pub(super) const fn y_size_bits(k: u8) -> usize {
78 k as usize + PARAM_EXT as usize
79}
80
81pub(super) const fn metadata_size_bits(k: u8, table_number: u8) -> usize {
83 k as usize
84 * match table_number {
85 1 => 1,
86 2 => 2,
87 3 | 4 => 4,
88 5 => 3,
89 6 => 2,
90 7 => 0,
91 _ => unreachable!(),
92 }
93}
94
95#[cfg(feature = "alloc")]
97const fn num_buckets(k: u8) -> usize {
98 2_usize
99 .pow(y_size_bits(k) as u32)
100 .div_ceil(PARAM_BC as usize)
101}
102
103#[cfg(feature = "alloc")]
105pub(super) const NUM_BUCKETS<const K: u8>: usize = num_buckets(K);
106#[cfg(feature = "parallel")]
108const NUM_BUCKET_PAIRS<const K: u8>: usize = num_buckets(K) - 1;
109#[cfg(feature = "alloc")]
111const MAX_TABLE_SIZE<const K: u8>: usize = 1 << K;
112#[cfg(any(feature = "alloc", test))]
113pub(super) const TABLE_1_YS_BATCH_SIMD<const K: u8>: usize =
114 K as usize * COMPUTE_F1_SIMD_FACTOR / u8::BITS as usize;
115
116#[cfg(feature = "alloc")]
119const fn max_table_entries(k: u8) -> usize {
120 (num_buckets(k) - 1) * REDUCED_MATCHES_COUNT
121}
122
123#[cfg(feature = "parallel")]
124#[inline(always)]
125fn strip_sync_unsafe_cell<const N: usize, T>(value: Box<[SyncUnsafeCell<T>; N]>) -> Box<[T; N]> {
126 unsafe { Box::from_raw(Box::into_raw(value).cast()) }
128}
129
130#[cfg(feature = "alloc")]
133fn partial_ys<const K: u8>(seed: Seed) -> Vec<u8> {
134 let output_len_bits = usize::from(K) * (1 << K);
135 let mut output = vec![0; output_len_bits.div_ceil(u8::BITS as usize)];
136
137 let key = Key::from(seed);
138 let iv = Iv::<ChaCha8>::default();
139
140 let mut cipher = ChaCha8::new(&key, &iv);
141
142 cipher.apply_keystream(&mut output);
143
144 output
145}
146
147#[cfg(feature = "alloc")]
155const fn bucket_size_upper_bound(k: u8, security_bits: u8) -> usize {
156 const LAMBDA: u64 = PARAM_BC as u64 / 2u64.pow(PARAM_EXT as u32);
159 const LN2_NUM: u128 = 693147;
162 const LN2_DEN: u128 = 1000000;
163
164 let ks = k as u128 + security_bits as u128;
166 let num = 3u128 * LAMBDA as u128 * ks * LN2_NUM;
169 let den = LN2_DEN;
171
172 let ceil_div: u128 = num.div_ceil(den);
173
174 let mut low = 0u64;
179 let mut high = u64::MAX;
180 while low < high {
181 let mid = low + (high - low) / 2;
182 let left = (mid as u128) * (mid as u128);
183 if left >= ceil_div {
184 high = mid;
185 } else {
186 low = mid + 1;
187 }
188 }
189 let add_term = low;
190
191 (LAMBDA + add_term) as usize
192}
193
194#[cfg(feature = "alloc")]
195fn group_by_buckets<const K: u8>(
196 ys: &[Y],
197) -> Box<[[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]> {
198 let mut bucket_offsets = [0_u16; NUM_BUCKETS::<K>];
199 let mut buckets = unsafe {
201 Box::<[[MaybeUninit<(Position, Y)>; REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]>::new_uninit()
202 .assume_init()
203 };
204
205 for (&y, position) in ys.iter().zip(Position::ZERO..) {
206 let bucket_index = (u32::from(y) / u32::from(PARAM_BC)) as usize;
207
208 let bucket_offset = unsafe { bucket_offsets.get_unchecked_mut(bucket_index) };
210 let bucket = unsafe { buckets.get_unchecked_mut(bucket_index) };
212
213 if *bucket_offset < REDUCED_BUCKET_SIZE as u16 {
214 bucket[*bucket_offset as usize].write((position, y));
215 *bucket_offset += 1;
216 }
217 }
218
219 for (bucket, initialized) in buckets.iter_mut().zip(bucket_offsets) {
220 bucket[usize::from(initialized)..].write_filled((Position::SENTINEL, Y::SENTINEL));
221 }
222
223 unsafe { Box::from_raw(Box::into_raw(buckets).cast()) }
225}
226
227#[cfg(feature = "parallel")]
233unsafe fn group_by_buckets_from_buckets<'a, const K: u8, I>(
234 iter: I,
235) -> Box<[[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]>
236where
237 I: Iterator<Item = (&'a [MaybeUninit<Y>; REDUCED_MATCHES_COUNT], usize)> + 'a,
238{
239 let mut bucket_offsets = [0_u16; NUM_BUCKETS::<K>];
240 let mut buckets = unsafe {
242 Box::<[[MaybeUninit<(Position, Y)>; REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]>::new_uninit()
243 .assume_init()
244 };
245
246 for ((ys, count), batch_start) in iter.zip((Position::ZERO..).step_by(REDUCED_MATCHES_COUNT)) {
247 let ys = unsafe { ys[..count].assume_init_ref() };
249 for (&y, position) in ys.iter().zip(batch_start..) {
250 let bucket_index = (u32::from(y) / u32::from(PARAM_BC)) as usize;
251
252 let bucket_offset = unsafe { bucket_offsets.get_unchecked_mut(bucket_index) };
254 let bucket = unsafe { buckets.get_unchecked_mut(bucket_index) };
256
257 if *bucket_offset < REDUCED_BUCKET_SIZE as u16 {
258 bucket[*bucket_offset as usize].write((position, y));
259 *bucket_offset += 1;
260 }
261 }
262 }
263
264 for (bucket, initialized) in buckets.iter_mut().zip(bucket_offsets) {
265 bucket[usize::from(initialized)..].write_filled((Position::SENTINEL, Y::SENTINEL));
266 }
267
268 unsafe { Box::from_raw(Box::into_raw(buckets).cast()) }
270}
271
272#[cfg(feature = "alloc")]
275fn sort_buckets<const K: u8>(
276 buckets: &mut [[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>],
277) {
278 for bucket in buckets.iter_mut() {
279 let len = bucket
280 .iter()
281 .take_while(|&&(pos, _)| pos != Position::SENTINEL)
282 .count();
283 bucket[..len].sort_unstable_by_key(|&(pos, y)| (u32::from(y), u32::from(pos)));
284 }
285}
286
287#[cfg(feature = "alloc")]
288#[derive(Debug, Copy, Clone, Deref)]
289#[repr(align(64))]
290struct CacheLineAligned<T>(T);
291
292#[cfg(feature = "alloc")]
294type LeftTargets =
295 [[CacheLineAligned<[R; const { PARAM_M as usize }]>; const { PARAM_BC as usize }]; 2];
296
297#[cfg(feature = "alloc")]
298fn calculate_left_targets() -> Arc<LeftTargets> {
299 let mut left_targets = Arc::<LeftTargets>::new_uninit();
300 let left_targets_slice = unsafe {
302 mem::transmute::<
303 &mut MaybeUninit<
304 [[CacheLineAligned<[R; const { PARAM_M as usize }]>; const { PARAM_BC as usize }];
305 2],
306 >,
307 &mut [[MaybeUninit<CacheLineAligned<[R; const { PARAM_M as usize }]>>; const {
308 PARAM_BC as usize
309 }]; 2],
310 >(Arc::get_mut_unchecked(&mut left_targets))
311 };
312
313 for parity in 0..=1 {
314 for r in 0..PARAM_BC {
315 let c = r / PARAM_C;
316
317 let mut arr = array::from_fn(|m| {
318 let m = m as u16;
319 R::from(
320 ((c + m) % PARAM_B) * PARAM_C
321 + (((2 * m + parity) * (2 * m + parity) + r) % PARAM_C),
322 )
323 });
324 arr.sort_unstable();
325 left_targets_slice[parity as usize][r as usize].write(CacheLineAligned(arr));
326 }
327 }
328
329 unsafe { left_targets.assume_init() }
331}
332
333fn calculate_left_target_on_demand(parity: u32, r: u32, m: u32) -> u32 {
334 let param_b = u32::from(PARAM_B);
335 let param_c = u32::from(PARAM_C);
336
337 ((r / param_c + m) % param_b) * param_c + (((2 * m + parity) * (2 * m + parity) + r) % param_c)
338}
339
340#[cfg(feature = "alloc")]
342#[derive(Debug, Clone)]
343pub struct TablesCache {
344 left_targets: Arc<LeftTargets>,
345}
346
347#[cfg(feature = "alloc")]
348impl Default for TablesCache {
349 fn default() -> Self {
351 Self {
352 left_targets: calculate_left_targets(),
353 }
354 }
355}
356
357#[cfg(feature = "alloc")]
358#[derive(Debug, Copy, Clone)]
359struct Match {
360 left_position: Position,
361 left_y: Y,
362 right_position: Position,
363}
364
365pub(super) fn compute_f1<const K: u8>(x: X, seed: &Seed) -> Y {
367 let skip_bits = u32::from(K) * u32::from(x);
368 let skip_u32s = skip_bits / u32::BITS;
369 let partial_y_offset = skip_bits % u32::BITS;
370
371 const U32S_PER_BLOCK: usize = size_of::<ChaCha8Block>() / size_of::<u32>();
372
373 let initial_state = ChaCha8State::init(seed, &[0; _]);
374 let first_block_counter = skip_u32s / U32S_PER_BLOCK as u32;
375 let u32_in_first_block = skip_u32s as usize % U32S_PER_BLOCK;
376
377 let first_block = initial_state.compute_block(first_block_counter);
378 let hi = first_block[u32_in_first_block].to_be();
379
380 let lo = if u32_in_first_block + 1 == U32S_PER_BLOCK {
382 let second_block = initial_state.compute_block(first_block_counter + 1);
384 second_block[0].to_be()
385 } else {
386 first_block[u32_in_first_block + 1].to_be()
387 };
388
389 let partial_y = (u64::from(hi) << u32::BITS) | u64::from(lo);
390
391 let pre_y = partial_y >> (u64::BITS - u32::from(K + PARAM_EXT) - partial_y_offset);
392 let pre_y = pre_y as u32;
393 let pre_y_mask = (u32::MAX << PARAM_EXT) & (u32::MAX >> (u32::BITS - u32::from(K + PARAM_EXT)));
395
396 let pre_ext = u32::from(x) >> (K - PARAM_EXT);
399
400 Y::from((pre_y & pre_y_mask) | pre_ext)
403}
404
405#[cfg(any(feature = "alloc", test))]
406pub(super) fn compute_f1_simd<const K: u8>(
407 xs: Simd<u32, COMPUTE_F1_SIMD_FACTOR>,
408 partial_ys: &[u8; TABLE_1_YS_BATCH_SIMD::<K>],
409) -> [Y; COMPUTE_F1_SIMD_FACTOR] {
410 let pre_ys_bytes = array::from_fn(|i| {
413 let partial_y_offset = i * usize::from(K);
414 let partial_y_length =
415 (partial_y_offset % u8::BITS as usize + usize::from(K)).div_ceil(u8::BITS as usize);
416 let mut pre_y_bytes = 0u64.to_be_bytes();
417 pre_y_bytes[..partial_y_length].copy_from_slice(
418 &partial_ys[partial_y_offset / u8::BITS as usize..][..partial_y_length],
419 );
420
421 u64::from_be_bytes(pre_y_bytes)
422 });
423 let pre_ys_right_offset = array::from_fn(|i| {
424 let partial_y_offset = i as u32 * u32::from(K);
425 u64::from(u64::BITS - u32::from(K + PARAM_EXT) - partial_y_offset % u8::BITS)
426 });
427 let pre_ys = Simd::from_array(pre_ys_bytes) >> Simd::from_array(pre_ys_right_offset);
428
429 let pre_ys_mask = Simd::splat(
431 (u32::MAX << usize::from(PARAM_EXT))
432 & (u32::MAX >> (u32::BITS as usize - usize::from(K + PARAM_EXT))),
433 );
434
435 let pre_exts = xs >> Simd::splat(u32::from(K - PARAM_EXT));
438
439 let ys = (pre_ys.cast() & pre_ys_mask) | pre_exts;
442
443 Y::array_from_repr(ys.to_array())
444}
445
446#[cfg(feature = "alloc")]
453unsafe fn find_matches_in_buckets<'a>(
454 left_bucket_index: u32,
455 left_bucket: &[(Position, Y); REDUCED_BUCKET_SIZE],
456 right_bucket: &[(Position, Y); REDUCED_BUCKET_SIZE],
457 matches: &'a mut [MaybeUninit<Match>; REDUCED_MATCHES_COUNT + PARAM_M as usize * 2],
460 left_targets: &LeftTargets,
461) -> &'a [Match] {
462 let left_base = left_bucket_index * u32::from(PARAM_BC);
463 let right_base = left_base + u32::from(PARAM_BC);
464
465 let mut rmap = Rmap::new();
466 for &(right_position, y) in right_bucket {
467 if right_position == Position::SENTINEL {
468 break;
469 }
470 let r = R::from((u32::from(y) - right_base) as u16);
471 unsafe {
474 rmap.add(r, right_position);
475 }
476 }
477
478 let parity = left_base % 2;
479 let left_targets_parity = &left_targets[parity as usize];
480 let mut next_match_index = 0;
481
482 for &(left_position, y) in left_bucket {
485 if left_position == Position::SENTINEL || next_match_index >= REDUCED_MATCHES_COUNT {
487 break;
489 }
490
491 let r = R::from((u32::from(y) - left_base) as u16);
492 let left_targets_r = unsafe { left_targets_parity.get_unchecked(usize::from(r)) };
494
495 for &r_target in left_targets_r.iter() {
496 let right_positions = unsafe { rmap.get(r_target) };
498
499 for &right_position in right_positions {
500 if next_match_index >= matches.len() {
501 break;
502 }
503 unsafe { matches.get_unchecked_mut(next_match_index) }.write(Match {
505 left_position,
506 left_y: y,
507 right_position,
508 });
509 next_match_index += 1;
510 }
511 }
512 }
513
514 unsafe { matches[..next_match_index].assume_init_ref() }
516}
517
518pub(super) fn has_match(left_y: Y, right_y: Y) -> bool {
520 let right_r = u32::from(right_y) % u32::from(PARAM_BC);
521 let parity = (u32::from(left_y) / u32::from(PARAM_BC)) % 2;
522 let left_r = u32::from(left_y) % u32::from(PARAM_BC);
523
524 let r_targets = array::from_fn::<_, const { PARAM_M as usize }, _>(|i| {
525 calculate_left_target_on_demand(parity, left_r, i as u32)
526 });
527
528 r_targets.contains(&right_r)
529}
530
531#[inline(always)]
532pub(super) fn compute_fn<const K: u8, const TABLE_NUMBER: u8, const PARENT_TABLE_NUMBER: u8>(
533 y: Y,
534 left_metadata: Metadata<K, PARENT_TABLE_NUMBER>,
535 right_metadata: Metadata<K, PARENT_TABLE_NUMBER>,
536) -> (Y, Metadata<K, TABLE_NUMBER>) {
537 let left_metadata = u128::from(left_metadata);
538 let right_metadata = u128::from(right_metadata);
539
540 let parent_metadata_bits = metadata_size_bits(K, PARENT_TABLE_NUMBER);
541
542 let y_and_left_bits = y_size_bits(K) + parent_metadata_bits;
544 let right_bits_start_offset = u128::BITS as usize - parent_metadata_bits;
545
546 let num_bytes_with_data =
548 (y_size_bits(K) + parent_metadata_bits * 2).div_ceil(u8::BITS as usize);
549
550 let hash = {
553 let y_bits = u128::from(y) << (u128::BITS as usize - y_size_bits(K));
555
556 let left_metadata_bits =
558 left_metadata << (u128::BITS as usize - parent_metadata_bits - y_size_bits(K));
559
560 if right_bits_start_offset < y_and_left_bits {
563 let right_bits_pushed_into_input_b = y_and_left_bits - right_bits_start_offset;
564 let right_bits_a = right_metadata >> right_bits_pushed_into_input_b;
567 let input_a = y_bits | left_metadata_bits | right_bits_a;
568 let input_b = right_metadata << (u128::BITS as usize - right_bits_pushed_into_input_b);
570
571 let input = [input_a.to_be_bytes(), input_b.to_be_bytes()];
572 let input_len =
573 size_of::<u128>() + right_bits_pushed_into_input_b.div_ceil(u8::BITS as usize);
574 blake3_hash(&input.as_flattened()[..input_len])
575 } else {
576 let right_bits_a = right_metadata << (right_bits_start_offset - y_and_left_bits);
577 let input_a = y_bits | left_metadata_bits | right_bits_a;
578
579 blake3_hash(&input_a.to_be_bytes()[..num_bytes_with_data])
580 }
581 };
582
583 let y_output = Y::from(
584 u32::from_be_bytes([hash[0], hash[1], hash[2], hash[3]])
585 >> (u32::BITS as usize - y_size_bits(K)),
586 );
587
588 let metadata_size_bits = metadata_size_bits(K, TABLE_NUMBER);
589
590 let metadata = if TABLE_NUMBER < 4 {
591 (left_metadata << parent_metadata_bits) | right_metadata
592 } else if metadata_size_bits > 0 {
593 let metadata = u128::from_be_bytes(
597 hash[y_size_bits(K) / u8::BITS as usize..][..size_of::<u128>()]
598 .try_into()
599 .expect("Always enough bits for any K; qed"),
600 );
601 let metadata = metadata << (y_size_bits(K) % u8::BITS as usize);
603 metadata >> (u128::BITS as usize - metadata_size_bits)
605 } else {
606 0
607 };
608
609 (y_output, Metadata::from(metadata))
610}
611
612#[cfg(any(feature = "alloc", test))]
616fn compute_fn_simd<const K: u8, const TABLE_NUMBER: u8, const PARENT_TABLE_NUMBER: u8>(
617 left_ys: [Y; COMPUTE_FN_SIMD_FACTOR],
618 left_metadatas: [Metadata<K, PARENT_TABLE_NUMBER>; COMPUTE_FN_SIMD_FACTOR],
619 right_metadatas: [Metadata<K, PARENT_TABLE_NUMBER>; COMPUTE_FN_SIMD_FACTOR],
620) -> (
621 [Y; COMPUTE_FN_SIMD_FACTOR],
622 [Metadata<K, TABLE_NUMBER>; COMPUTE_FN_SIMD_FACTOR],
623) {
624 let parent_metadata_bits = metadata_size_bits(K, PARENT_TABLE_NUMBER);
625 let metadata_size_bits = metadata_size_bits(K, TABLE_NUMBER);
626
627 let left_metadatas: [u128; COMPUTE_FN_SIMD_FACTOR] = seq!(N in 0..16 {
630 [
631 #(
632 u128::from(left_metadatas[N]),
633 )*
634 ]
635 });
636 let right_metadatas: [u128; COMPUTE_FN_SIMD_FACTOR] = seq!(N in 0..16 {
637 [
638 #(
639 u128::from(right_metadatas[N]),
640 )*
641 ]
642 });
643
644 let y_and_left_bits = y_size_bits(K) + parent_metadata_bits;
646 let right_bits_start_offset = u128::BITS as usize - parent_metadata_bits;
647
648 let num_bytes_with_data =
650 (y_size_bits(K) + parent_metadata_bits * 2).div_ceil(u8::BITS as usize);
651
652 let hashes: [_; COMPUTE_FN_SIMD_FACTOR] = seq!(N in 0..16 {
657 [
658 #(
659 {
660 let y = left_ys[N];
661 let left_metadata = left_metadatas[N];
662 let right_metadata = right_metadatas[N];
663
664 let y_bits = u128::from(y) << (u128::BITS as usize - y_size_bits(K));
667
668 let left_metadata_bits =
670 left_metadata << (u128::BITS as usize - parent_metadata_bits - y_size_bits(K));
671
672 if right_bits_start_offset < y_and_left_bits {
675 let right_bits_pushed_into_input_b = y_and_left_bits - right_bits_start_offset;
676 let right_bits_a = right_metadata >> right_bits_pushed_into_input_b;
679 let input_a = y_bits | left_metadata_bits | right_bits_a;
680 let input_b = right_metadata << (u128::BITS as usize - right_bits_pushed_into_input_b);
682
683 let input = [input_a.to_be_bytes(), input_b.to_be_bytes()];
684 let input_len =
685 size_of::<u128>() + right_bits_pushed_into_input_b.div_ceil(u8::BITS as usize);
686 blake3_hash(&input.as_flattened()[..input_len])
687 } else {
688 let right_bits_a = right_metadata << (right_bits_start_offset - y_and_left_bits);
689 let input_a = y_bits | left_metadata_bits | right_bits_a;
690
691 blake3_hash(&input_a.to_be_bytes()[..num_bytes_with_data])
692 }
693 },
694 )*
695 ]
696 });
697
698 let y_outputs = Simd::from_array(
699 hashes.map(|hash| u32::from_be_bytes([hash[0], hash[1], hash[2], hash[3]])),
700 ) >> (u32::BITS - y_size_bits(K) as u32);
701 let y_outputs = Y::array_from_repr(y_outputs.to_array());
702
703 let metadatas = if TABLE_NUMBER < 4 {
704 seq!(N in 0..16 {
705 [
706 #(
707 Metadata::from((left_metadatas[N] << parent_metadata_bits) | right_metadatas[N]),
708 )*
709 ]
710 })
711 } else if metadata_size_bits > 0 {
712 seq!(N in 0..16 {
716 [
717 #(
718 {
719 let metadata = u128::from_be_bytes(
720 hashes[N][y_size_bits(K) / u8::BITS as usize..][..size_of::<u128>()]
721 .try_into()
722 .expect("Always enough bits for any K; qed"),
723 );
724 let metadata = metadata << (y_size_bits(K) % u8::BITS as usize);
726 Metadata::from(metadata >> (u128::BITS as usize - metadata_size_bits))
728 },
729 )*
730 ]
731 })
732 } else {
733 [Metadata::default(); _]
734 };
735
736 (y_outputs, metadatas)
737}
738
739#[cfg(feature = "alloc")]
742unsafe fn match_to_result<const K: u8, const TABLE_NUMBER: u8, const PARENT_TABLE_NUMBER: u8>(
743 parent_table: &Table<K, PARENT_TABLE_NUMBER>,
744 m: &Match,
745) -> (Y, [Position; 2], Metadata<K, TABLE_NUMBER>)
746where
747 Table<K, PARENT_TABLE_NUMBER>: private::NotLastTable,
748{
749 let left_metadata = unsafe { parent_table.metadata(m.left_position) };
751 let right_metadata = unsafe { parent_table.metadata(m.right_position) };
753
754 let (y, metadata) =
755 compute_fn::<K, TABLE_NUMBER, PARENT_TABLE_NUMBER>(m.left_y, left_metadata, right_metadata);
756
757 (y, [m.left_position, m.right_position], metadata)
758}
759
760#[cfg(feature = "alloc")]
763#[inline(always)]
764unsafe fn match_to_result_simd<const K: u8, const TABLE_NUMBER: u8, const PARENT_TABLE_NUMBER: u8>(
765 parent_table: &Table<K, PARENT_TABLE_NUMBER>,
766 matches: &[Match; COMPUTE_FN_SIMD_FACTOR],
767) -> (
768 [Y; COMPUTE_FN_SIMD_FACTOR],
769 [[Position; 2]; COMPUTE_FN_SIMD_FACTOR],
770 [Metadata<K, TABLE_NUMBER>; COMPUTE_FN_SIMD_FACTOR],
771)
772where
773 Table<K, PARENT_TABLE_NUMBER>: private::NotLastTable,
774{
775 let left_ys: [_; COMPUTE_FN_SIMD_FACTOR] = seq!(N in 0..16 {
776 [
777 #(
778 matches[N].left_y,
779 )*
780 ]
781 });
782 let left_metadatas: [_; COMPUTE_FN_SIMD_FACTOR] = unsafe {
784 seq!(N in 0..16 {
785 [
786 #(
787 parent_table.metadata(matches[N].left_position),
788 )*
789 ]
790 })
791 };
792 let right_metadatas: [_; COMPUTE_FN_SIMD_FACTOR] = unsafe {
794 seq!(N in 0..16 {
795 [
796 #(
797 parent_table.metadata(matches[N].right_position),
798 )*
799 ]
800 })
801 };
802
803 let (y_outputs, metadatas) = compute_fn_simd::<K, TABLE_NUMBER, PARENT_TABLE_NUMBER>(
804 left_ys,
805 left_metadatas,
806 right_metadatas,
807 );
808
809 let positions = seq!(N in 0..16 {
810 [
811 #(
812 [
813 matches[N].left_position,
814 matches[N].right_position,
815 ],
816 )*
817 ]
818 });
819
820 (y_outputs, positions, metadatas)
821}
822
823#[cfg(feature = "alloc")]
827#[inline(always)]
828unsafe fn matches_to_results<const K: u8, const TABLE_NUMBER: u8, const PARENT_TABLE_NUMBER: u8>(
829 parent_table: &Table<K, PARENT_TABLE_NUMBER>,
830 matches: &[Match],
831 ys: &mut [MaybeUninit<Y>],
832 positions: &mut [MaybeUninit<[Position; 2]>],
833 metadatas: &mut [MaybeUninit<Metadata<K, TABLE_NUMBER>>],
834) where
835 Table<K, PARENT_TABLE_NUMBER>: private::NotLastTable,
836{
837 let (grouped_matches, other_matches) = matches.as_chunks::<COMPUTE_FN_SIMD_FACTOR>();
838 let (grouped_ys, other_ys) = ys.split_at_mut(grouped_matches.as_flattened().len());
839 let grouped_ys = grouped_ys.as_chunks_mut::<COMPUTE_FN_SIMD_FACTOR>().0;
840 let (grouped_positions, other_positions) =
841 positions.split_at_mut(grouped_matches.as_flattened().len());
842 let grouped_positions = grouped_positions
843 .as_chunks_mut::<COMPUTE_FN_SIMD_FACTOR>()
844 .0;
845 let (grouped_metadatas, other_metadatas) =
846 metadatas.split_at_mut(grouped_matches.as_flattened().len());
847 let grouped_metadatas = grouped_metadatas
848 .as_chunks_mut::<COMPUTE_FN_SIMD_FACTOR>()
849 .0;
850
851 for (((grouped_matches, grouped_ys), grouped_positions), grouped_metadatas) in grouped_matches
852 .iter()
853 .zip(grouped_ys)
854 .zip(grouped_positions)
855 .zip(grouped_metadatas)
856 {
857 let (ys_group, positions_group, metadatas_group) =
859 unsafe { match_to_result_simd(parent_table, grouped_matches) };
860 grouped_ys.write_copy_of_slice(&ys_group);
861 grouped_positions.write_copy_of_slice(&positions_group);
862
863 if metadata_size_bits(K, TABLE_NUMBER) > 0 {
865 grouped_metadatas.write_copy_of_slice(&metadatas_group);
866 }
867 }
868 for (((other_match, other_y), other_positions), other_metadata) in other_matches
869 .iter()
870 .zip(other_ys)
871 .zip(other_positions)
872 .zip(other_metadatas)
873 {
874 let (y, p, metadata) = unsafe { match_to_result(parent_table, other_match) };
876 other_y.write(y);
877 other_positions.write(p);
878 if metadata_size_bits(K, TABLE_NUMBER) > 0 {
880 other_metadata.write(metadata);
881 }
882 }
883}
884
885#[cfg(feature = "alloc")]
887#[derive(Debug)]
888pub(super) enum PrunedTable<const K: u8, const TABLE_NUMBER: u8> {
889 First,
890 Other {
892 positions: Box<[MaybeUninit<[Position; 2]>]>,
894 },
895 #[cfg(feature = "parallel")]
897 OtherBuckets {
898 positions:
902 Box<[[MaybeUninit<[Position; 2]>; REDUCED_MATCHES_COUNT]; NUM_BUCKET_PAIRS::<K>]>,
903 },
904}
905
906#[cfg(feature = "alloc")]
907impl<const K: u8, const TABLE_NUMBER: u8> PrunedTable<K, TABLE_NUMBER> {
908 #[inline(always)]
915 pub(super) unsafe fn position(&self, position: Position) -> [Position; 2] {
916 match self {
917 Self::First => {
918 unreachable!("Not the first table");
919 }
920 Self::Other { positions, .. } => {
921 unsafe { positions.get_unchecked(usize::from(position)).assume_init() }
923 }
924 #[cfg(feature = "parallel")]
925 Self::OtherBuckets { positions, .. } => {
926 unsafe {
928 positions
929 .as_flattened()
930 .get_unchecked(usize::from(position))
931 .assume_init()
932 }
933 }
934 }
935 }
936}
937
938#[cfg(feature = "alloc")]
939#[derive(Debug)]
940pub(super) enum Table<const K: u8, const TABLE_NUMBER: u8> {
941 First {
943 buckets: Box<[[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]>,
947 },
948 Other {
950 positions: Box<[MaybeUninit<[Position; 2]>]>,
952 metadatas: Box<[MaybeUninit<Metadata<K, TABLE_NUMBER>>]>,
954 buckets: Box<[[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]>,
958 },
959 #[cfg(feature = "parallel")]
961 OtherBuckets {
962 positions:
966 Box<[[MaybeUninit<[Position; 2]>; REDUCED_MATCHES_COUNT]; NUM_BUCKET_PAIRS::<K>]>,
967 metadatas: Box<
971 [[MaybeUninit<Metadata<K, TABLE_NUMBER>>; REDUCED_MATCHES_COUNT];
972 NUM_BUCKET_PAIRS::<K>],
973 >,
974 buckets: Box<[[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]>,
978 },
979}
980
981#[cfg(feature = "alloc")]
982impl<const K: u8> Table<K, 1> {
983 pub(super) fn create(seed: Seed) -> Self {
985 debug_assert!(
988 MAX_BUCKET_SIZE >= bucket_size_upper_bound(K, BUCKET_SIZE_UPPER_BOUND_SECURITY_BITS),
989 "Max bucket size is not sufficiently large"
990 );
991
992 let partial_ys = partial_ys::<K>(seed);
993
994 let mut ys =
996 unsafe { Box::<[MaybeUninit<Y>; MAX_TABLE_SIZE::<K>]>::new_uninit().assume_init() };
997
998 for ((ys, xs_batch_start), partial_ys) in ys
999 .as_chunks_mut::<COMPUTE_F1_SIMD_FACTOR>()
1000 .0
1001 .iter_mut()
1002 .zip((X::ZERO..).step_by(COMPUTE_F1_SIMD_FACTOR))
1003 .zip(partial_ys.as_chunks::<{ TABLE_1_YS_BATCH_SIMD::<K> }>().0)
1004 {
1005 let xs = Simd::splat(u32::from(xs_batch_start))
1006 + Simd::from_array(array::from_fn(|i| i as u32));
1007 let ys_batch = compute_f1_simd::<K>(xs, partial_ys);
1008
1009 ys.write_copy_of_slice(&ys_batch);
1010 }
1011
1012 let ys = unsafe {
1015 let ys_len = ys.len();
1016 let ys = Box::into_raw(ys);
1017 Vec::from_raw_parts(ys.cast(), ys_len, ys_len)
1018 };
1019
1020 let mut buckets = group_by_buckets::<K>(&ys);
1022 sort_buckets::<K>(&mut buckets);
1023
1024 Self::First { buckets }
1025 }
1026
1027 #[cfg(feature = "parallel")]
1029 pub(super) fn create_parallel(seed: Seed) -> Self {
1030 debug_assert!(
1033 MAX_BUCKET_SIZE >= bucket_size_upper_bound(K, BUCKET_SIZE_UPPER_BOUND_SECURITY_BITS),
1034 "Max bucket size is not sufficiently large"
1035 );
1036
1037 let partial_ys = partial_ys::<K>(seed);
1038
1039 let mut ys =
1041 unsafe { Box::<[MaybeUninit<Y>; MAX_TABLE_SIZE::<K>]>::new_uninit().assume_init() };
1042
1043 for ((ys, xs_batch_start), partial_ys) in ys
1045 .as_chunks_mut::<COMPUTE_F1_SIMD_FACTOR>()
1046 .0
1047 .iter_mut()
1048 .zip((X::ZERO..).step_by(COMPUTE_F1_SIMD_FACTOR))
1049 .zip(partial_ys.as_chunks::<{ TABLE_1_YS_BATCH_SIMD::<K> }>().0)
1050 {
1051 let xs = Simd::splat(u32::from(xs_batch_start))
1052 + Simd::from_array(array::from_fn(|i| i as u32));
1053 let ys_batch = compute_f1_simd::<K>(xs, partial_ys);
1054
1055 ys.write_copy_of_slice(&ys_batch);
1056 }
1057
1058 let ys = unsafe {
1061 let ys_len = ys.len();
1062 let ys = Box::into_raw(ys);
1063 Vec::from_raw_parts(ys.cast(), ys_len, ys_len)
1064 };
1065
1066 let mut buckets = group_by_buckets::<K>(&ys);
1068 sort_buckets::<K>(&mut buckets);
1069
1070 Self::First { buckets }
1071 }
1072}
1073
1074#[cfg(feature = "alloc")]
1075mod private {
1076 pub(in super::super) trait SupportedOtherTables {}
1077 pub(in super::super) trait NotLastTable {}
1078}
1079
1080#[cfg(feature = "alloc")]
1081impl<const K: u8> private::SupportedOtherTables for Table<K, 2> {}
1082#[cfg(feature = "alloc")]
1083impl<const K: u8> private::SupportedOtherTables for Table<K, 3> {}
1084#[cfg(feature = "alloc")]
1085impl<const K: u8> private::SupportedOtherTables for Table<K, 4> {}
1086#[cfg(feature = "alloc")]
1087impl<const K: u8> private::SupportedOtherTables for Table<K, 5> {}
1088#[cfg(feature = "alloc")]
1089impl<const K: u8> private::SupportedOtherTables for Table<K, 6> {}
1090#[cfg(feature = "alloc")]
1091impl<const K: u8> private::SupportedOtherTables for Table<K, 7> {}
1092
1093#[cfg(feature = "alloc")]
1094impl<const K: u8> private::NotLastTable for Table<K, 1> {}
1095#[cfg(feature = "alloc")]
1096impl<const K: u8> private::NotLastTable for Table<K, 2> {}
1097#[cfg(feature = "alloc")]
1098impl<const K: u8> private::NotLastTable for Table<K, 3> {}
1099#[cfg(feature = "alloc")]
1100impl<const K: u8> private::NotLastTable for Table<K, 4> {}
1101#[cfg(feature = "alloc")]
1102impl<const K: u8> private::NotLastTable for Table<K, 5> {}
1103#[cfg(feature = "alloc")]
1104impl<const K: u8> private::NotLastTable for Table<K, 6> {}
1105
1106#[cfg(feature = "alloc")]
1107impl<const K: u8, const TABLE_NUMBER: u8> Table<K, TABLE_NUMBER>
1108where
1109 Self: private::SupportedOtherTables,
1110{
1111 pub(super) fn create<const PARENT_TABLE_NUMBER: u8>(
1115 parent_table: Table<K, PARENT_TABLE_NUMBER>,
1116 cache: &TablesCache,
1117 ) -> (Self, PrunedTable<K, PARENT_TABLE_NUMBER>)
1118 where
1119 Table<K, PARENT_TABLE_NUMBER>: private::NotLastTable,
1120 {
1121 let left_targets = &*cache.left_targets;
1122 let mut initialized_elements = 0_usize;
1123 let mut ys: Box<[MaybeUninit<Y>]> = Box::new_uninit_slice(max_table_entries(K));
1124 let mut positions: Box<[MaybeUninit<[Position; 2]>]> =
1125 Box::new_uninit_slice(max_table_entries(K));
1126 let mut metadatas: Box<[MaybeUninit<Metadata<K, TABLE_NUMBER>>]> =
1128 Box::new_uninit_slice(max_table_entries(K));
1129
1130 for ([left_bucket, right_bucket], left_bucket_index) in
1131 parent_table.buckets().array_windows().zip(0..)
1132 {
1133 let mut matches = [MaybeUninit::uninit(); _];
1134 let matches = unsafe {
1137 find_matches_in_buckets(
1138 left_bucket_index,
1139 left_bucket,
1140 right_bucket,
1141 &mut matches,
1142 left_targets,
1143 )
1144 };
1145 let matches = &matches[..matches.len().min(REDUCED_MATCHES_COUNT)];
1147 let (ys, positions, metadatas) = unsafe {
1149 (
1150 ys.split_at_mut_unchecked(initialized_elements).1,
1151 positions.split_at_mut_unchecked(initialized_elements).1,
1152 metadatas.split_at_mut_unchecked(initialized_elements).1,
1153 )
1154 };
1155
1156 let (ys, positions, metadatas) = unsafe {
1159 (
1160 ys.split_at_mut_unchecked(matches.len()).0,
1161 positions.split_at_mut_unchecked(matches.len()).0,
1162 metadatas.split_at_mut_unchecked(matches.len()).0,
1163 )
1164 };
1165
1166 unsafe {
1169 matches_to_results(&parent_table, matches, ys, positions, metadatas);
1170 }
1171
1172 initialized_elements += matches.len();
1173 }
1174
1175 let parent_table = parent_table.prune();
1176
1177 let ys = unsafe {
1180 let ys_len = ys.len();
1181 let ys = Box::into_raw(ys);
1182 Vec::from_raw_parts(ys.cast(), initialized_elements, ys_len)
1183 };
1184
1185 let mut buckets = group_by_buckets::<K>(&ys);
1187 sort_buckets::<K>(&mut buckets);
1188
1189 let table = Self::Other {
1190 positions,
1191 metadatas,
1192 buckets,
1193 };
1194
1195 (table, parent_table)
1196 }
1197
1198 #[cfg(feature = "parallel")]
1202 pub(super) fn create_parallel<const PARENT_TABLE_NUMBER: u8>(
1203 parent_table: Table<K, PARENT_TABLE_NUMBER>,
1204 cache: &TablesCache,
1205 ) -> (Self, PrunedTable<K, PARENT_TABLE_NUMBER>)
1206 where
1207 Table<K, PARENT_TABLE_NUMBER>: private::NotLastTable,
1208 {
1209 let ys = unsafe {
1211 Box::<[SyncUnsafeCell<[MaybeUninit<_>; REDUCED_MATCHES_COUNT]>; NUM_BUCKET_PAIRS::<K>]>::new_uninit().assume_init()
1212 };
1213 let positions = unsafe {
1215 Box::<[SyncUnsafeCell<[MaybeUninit<_>; REDUCED_MATCHES_COUNT]>; NUM_BUCKET_PAIRS::<K>]>::new_uninit().assume_init()
1216 };
1217 let metadatas = unsafe {
1219 Box::<[SyncUnsafeCell<[MaybeUninit<_>; REDUCED_MATCHES_COUNT]>; NUM_BUCKET_PAIRS::<K>]>::new_uninit().assume_init()
1220 };
1221 let global_results_counts =
1222 array::from_fn::<_, { NUM_BUCKET_PAIRS::<K> }, _>(|_| SyncUnsafeCell::new(0u16));
1223
1224 let left_targets = &*cache.left_targets;
1225
1226 let buckets = parent_table.buckets();
1227 let bucket_batch_size = CACHE_LINE_SIZE / size_of::<u16>();
1230 let bucket_batch_index = AtomicUsize::new(0);
1231
1232 rayon::broadcast(|_ctx| {
1233 loop {
1234 let bucket_batch_index = bucket_batch_index.fetch_add(1, Ordering::Relaxed);
1235
1236 let buckets_batch = buckets
1237 .array_windows::<2>()
1238 .enumerate()
1239 .skip(bucket_batch_index * bucket_batch_size)
1240 .take(bucket_batch_size);
1241
1242 if buckets_batch.is_empty() {
1243 break;
1244 }
1245
1246 for (left_bucket_index, [left_bucket, right_bucket]) in buckets_batch {
1247 let mut matches = [MaybeUninit::uninit(); _];
1248 let matches = unsafe {
1251 find_matches_in_buckets(
1252 left_bucket_index as u32,
1253 left_bucket,
1254 right_bucket,
1255 &mut matches,
1256 left_targets,
1257 )
1258 };
1259 let matches = &matches[..matches.len().min(REDUCED_MATCHES_COUNT)];
1261
1262 let ys = unsafe { &mut *ys.get_unchecked(left_bucket_index).get() };
1265 let positions =
1268 unsafe { &mut *positions.get_unchecked(left_bucket_index).get() };
1269 let metadatas =
1272 unsafe { &mut *metadatas.get_unchecked(left_bucket_index).get() };
1273 let count = unsafe {
1276 &mut *global_results_counts.get_unchecked(left_bucket_index).get()
1277 };
1278
1279 unsafe {
1282 matches_to_results::<_, TABLE_NUMBER, _>(
1283 &parent_table,
1284 matches,
1285 ys,
1286 positions,
1287 metadatas,
1288 )
1289 };
1290 *count = matches.len() as u16;
1291 }
1292 }
1293 });
1294
1295 let parent_table = parent_table.prune();
1296
1297 let ys = strip_sync_unsafe_cell(ys);
1298 let positions = strip_sync_unsafe_cell(positions);
1299 let metadatas = strip_sync_unsafe_cell(metadatas);
1300
1301 let mut buckets = unsafe {
1304 group_by_buckets_from_buckets::<K, _>(
1305 ys.iter().zip(
1306 global_results_counts
1307 .into_iter()
1308 .map(|count| usize::from(count.into_inner())),
1309 ),
1310 )
1311 };
1312 sort_buckets::<K>(&mut buckets);
1313
1314 let table = Self::OtherBuckets {
1315 positions,
1316 metadatas,
1317 buckets,
1318 };
1319
1320 (table, parent_table)
1321 }
1322
1323 #[inline(always)]
1330 pub(super) unsafe fn position(&self, position: Position) -> [Position; 2] {
1331 match self {
1332 Self::First { .. } => {
1333 unreachable!("Not the first table");
1334 }
1335 Self::Other { positions, .. } => {
1336 unsafe { positions.get_unchecked(usize::from(position)).assume_init() }
1338 }
1339 #[cfg(feature = "parallel")]
1340 Self::OtherBuckets { positions, .. } => {
1341 unsafe {
1343 positions
1344 .as_flattened()
1345 .get_unchecked(usize::from(position))
1346 .assume_init()
1347 }
1348 }
1349 }
1350 }
1351}
1352
1353#[cfg(feature = "alloc")]
1354impl<const K: u8, const TABLE_NUMBER: u8> Table<K, TABLE_NUMBER>
1355where
1356 Self: private::NotLastTable,
1357{
1358 #[inline(always)]
1363 unsafe fn metadata(&self, position: Position) -> Metadata<K, TABLE_NUMBER> {
1364 match self {
1365 Self::First { .. } => {
1366 Metadata::from(X::from(u32::from(position)))
1368 }
1369 Self::Other { metadatas, .. } => {
1370 unsafe { metadatas.get_unchecked(usize::from(position)).assume_init() }
1372 }
1373 #[cfg(feature = "parallel")]
1374 Self::OtherBuckets { metadatas, .. } => {
1375 unsafe {
1377 metadatas
1378 .as_flattened()
1379 .get_unchecked(usize::from(position))
1380 .assume_init()
1381 }
1382 }
1383 }
1384 }
1385}
1386
1387#[cfg(feature = "alloc")]
1388impl<const K: u8, const TABLE_NUMBER: u8> Table<K, TABLE_NUMBER> {
1389 #[inline(always)]
1390 fn prune(self) -> PrunedTable<K, TABLE_NUMBER> {
1391 match self {
1392 Self::First { .. } => PrunedTable::First,
1393 Self::Other { positions, .. } => PrunedTable::Other { positions },
1394 #[cfg(feature = "parallel")]
1395 Self::OtherBuckets { positions, .. } => PrunedTable::OtherBuckets { positions },
1396 }
1397 }
1398
1399 #[inline(always)]
1401 pub(super) fn buckets(&self) -> &[[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>] {
1402 match self {
1403 Self::First { buckets, .. } => buckets,
1404 Self::Other { buckets, .. } => buckets,
1405 #[cfg(feature = "parallel")]
1406 Self::OtherBuckets { buckets, .. } => buckets,
1407 }
1408 }
1409}