Skip to main content

subspace_proof_of_space/chiapos/
table.rs

1//! `chiapos` table implementation.
2
3#[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/// Bucket size used for fixed-size bucketed arrays.
55///
56/// Must be `MAX_BUCKET_SIZE` to avoid silently dropping entries when a bucket's Y values exceed
57/// the limit. With `PARAM_BC / 2^PARAM_EXT ≈ 236` expected entries per bucket and std dev ≈ 15.4,
58/// the original value of 272 (~2.3 sigma) caused overflow in production (K=20), dropping hundreds
59/// of proofs and causing "Missing PoS proof" farming errors.
60#[allow(dead_code, reason = "unused when crate is compiled separately")]
61pub(crate) const REDUCED_BUCKET_SIZE: usize = MAX_BUCKET_SIZE;
62/// Matches count limit per bucket pair.
63///
64/// Must be `MAX_BUCKET_SIZE` to avoid silently dropping matches. The original value of 288 could
65/// cause early termination in `find_matches_in_buckets`, losing valid proof paths.
66#[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
76/// Compute the size of `y` in bits
77pub(super) const fn y_size_bits(k: u8) -> usize {
78    k as usize + PARAM_EXT as usize
79}
80
81/// Metadata size in bits
82pub(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/// Number of buckets for a given `k`
96#[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/// Number of buckets for a given `K`
104#[cfg(feature = "alloc")]
105pub(super) const NUM_BUCKETS<const K: u8>: usize = num_buckets(K);
106/// Number of bucket pairs for a given `K`
107#[cfg(feature = "parallel")]
108const NUM_BUCKET_PAIRS<const K: u8>: usize = num_buckets(K) - 1;
109/// Size of the first table and max size for other tables
110#[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/// Upper bound on entries in a non-first table: `num_buckets(K) - 1` bucket pairs, each yielding at
117/// most `REDUCED_MATCHES_COUNT` matches. The actual count is data-dependent and can exceed `2^K`.
118#[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    // SAFETY: `SyncUnsafeCell` has the same layout as `T`
127    unsafe { Box::from_raw(Box::into_raw(value).cast()) }
128}
129
130/// ChaCha8 [`Vec`] sufficient for the whole first table for [`K`].
131/// Prefer [`partial_y`] if you need partial y just for a single `x`.
132#[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/// Calculate a probabilistic upper bound on the Chia bucket size for a given `k` and
148/// `security_bits` (security level).
149///
150/// This is based on a Chernoff bound for the Poisson distribution with mean
151/// `lambda = PARAM_BC / 2^PARAM_EXT`, ensuring the probability that any bucket exceeds the bound is
152/// less than `2^{-security_bits}`.
153/// The bound is lambda + ceil(sqrt(3 * lambda * (k + security_bits) * ln(2))).
154#[cfg(feature = "alloc")]
155const fn bucket_size_upper_bound(k: u8, security_bits: u8) -> usize {
156    // Lambda is the expected number of entries in a bucket, approximated as
157    // `PARAM_BC / 2^PARAM_EXT`. It is independent of `k`.
158    const LAMBDA: u64 = PARAM_BC as u64 / 2u64.pow(PARAM_EXT as u32);
159    // Approximation of ln(2) as a fraction: ln(2) ≈ LN2_NUM / LN2_DEN.
160    // This allows integer-only computation of the square root term involving ln(2).
161    const LN2_NUM: u128 = 693147;
162    const LN2_DEN: u128 = 1000000;
163
164    // `k + security_bits` for the union bound over ~2^k intervals
165    let ks = k as u128 + security_bits as u128;
166    // Compute numerator for the expression under the square root:
167    // `3 * lambda * (k + security_bits) * LN2_NUM`
168    let num = 3u128 * LAMBDA as u128 * ks * LN2_NUM;
169    // Denominator for ln(2): `LN2_DEN`
170    let den = LN2_DEN;
171
172    let ceil_div: u128 = num.div_ceil(den);
173
174    // Binary search to find the smallest `x` such that `x * x * den >= num`,
175    // which computes `ceil(sqrt(num / den))` without floating-point.
176    // We use a custom binary search over `u64` range because binary search in the standard library
177    // operates on sorted slices, not directly on integer ranges for solving inequalities like this.
178    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    // SAFETY: Contents is `MaybeUninit`
200    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        // SAFETY: Bucket is obtained using division by `PARAM_BC` and fits by definition
209        let bucket_offset = unsafe { bucket_offsets.get_unchecked_mut(bucket_index) };
210        // SAFETY: Bucket is obtained using division by `PARAM_BC` and fits by definition
211        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    // SAFETY: All entries are initialized
224    unsafe { Box::from_raw(Box::into_raw(buckets).cast()) }
225}
226
227/// Similar to [`group_by_buckets()`], but processes buckets instead of a flat list of `y`s.
228///
229/// # Safety
230/// Iterator item is a list of potentially uninitialized `y`s and a number of initialized `y`s. The
231/// number of initialized `y`s must be correct.
232#[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    // SAFETY: Contents is `MaybeUninit`
241    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        // SAFETY: Function contract guarantees that `y`s are initialized
248        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            // SAFETY: Bucket is obtained using division by `PARAM_BC` and fits by definition
253            let bucket_offset = unsafe { bucket_offsets.get_unchecked_mut(bucket_index) };
254            // SAFETY: Bucket is obtained using division by `PARAM_BC` and fits by definition
255            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    // SAFETY: All entries are initialized
269    unsafe { Box::from_raw(Box::into_raw(buckets).cast()) }
270}
271
272/// Sort entries within each bucket by Y value to ensure deterministic proof ordering.
273/// This matches the old code's behavior where entries were globally Y-sorted before bucketing.
274#[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/// Mapping from `parity` to `r` to `m`
293#[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    // SAFETY: Same layout and uninitialized in both cases
301    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    // SAFETY: Initialized all entries
330    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/// Caches that can be used to optimize the creation of multiple [`Tables`](super::Tables).
341#[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    /// Create a new instance
350    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
365/// `partial_y_offset` is in bits within `partial_y`
366pub(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    // TODO: Is SIMD version of `compute_block()` that produces two blocks at once possible?
381    let lo = if u32_in_first_block + 1 == U32S_PER_BLOCK {
382        // Spilled over into the second block
383        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    // Mask for clearing the rest of bits of `pre_y`.
394    let pre_y_mask = (u32::MAX << PARAM_EXT) & (u32::MAX >> (u32::BITS - u32::from(K + PARAM_EXT)));
395
396    // Extract `PARAM_EXT` most significant bits from `x` and store in the final offset of
397    // eventual `y` with the rest of bits being zero (`x` is `0..2^K`)
398    let pre_ext = u32::from(x) >> (K - PARAM_EXT);
399
400    // Combine all of the bits together:
401    // [padding zero bits][`K` bits rom `partial_y`][`PARAM_EXT` bits from `x`]
402    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    // Each element contains `K` desired bits of `partial_ys` in the final offset of eventual `ys`
411    // with the rest of bits being in undefined state
412    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    // Mask for clearing the rest of bits of `pre_ys`.
430    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    // Extract `PARAM_EXT` most significant bits from `xs` and store in the final offset of
436    // eventual `ys` with the rest of bits being in undefined state.
437    let pre_exts = xs >> Simd::splat(u32::from(K - PARAM_EXT));
438
439    // Combine all of the bits together:
440    // [padding zero bits][`K` bits rom `partial_y`][`PARAM_EXT` bits from `x`]
441    let ys = (pre_ys.cast() & pre_ys_mask) | pre_exts;
442
443    Y::array_from_repr(ys.to_array())
444}
445
446/// For verification use [`has_match`] instead.
447///
448/// # Safety
449/// Left and right bucket positions must correspond to the parent table.
450// TODO: Try to reduce the `matches` size further by processing `left_bucket` in chunks (like halves
451//  for example)
452#[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    // `PARAM_M as usize * 2` corresponds to the upper bound number of matches a single `y` in the
458    // left bucket might have here
459    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        // SAFETY: `r` is within `0..PARAM_BC` range by definition, the right bucket is limited to
472        // `REDUCED_BUCKETS_SIZE`
473        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    // TODO: Simd read for left bucket? It might be more efficient in terms of memory access to
483    //  process chunks of the left bucket against one right value for each at a time
484    for &(left_position, y) in left_bucket {
485        // `next_match_index >= REDUCED_MATCHES_COUNT` is crucial to make sure
486        if left_position == Position::SENTINEL || next_match_index >= REDUCED_MATCHES_COUNT {
487            // Sentinel values are padded to the end of the bucket
488            break;
489        }
490
491        let r = R::from((u32::from(y) - left_base) as u16);
492        // SAFETY: `r` is within a bucket and exists by definition
493        let left_targets_r = unsafe { left_targets_parity.get_unchecked(usize::from(r)) };
494
495        for &r_target in left_targets_r.iter() {
496            // SAFETY: Targets are always limited to `PARAM_BC`
497            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                // SAFETY: Bounds checked above
504                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    // SAFETY: Initialized this many matches
515    unsafe { matches[..next_match_index].assume_init_ref() }
516}
517
518/// Simplified version of [`find_matches_in_buckets`] for verification purposes.
519pub(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    // Part of the `right_bits` at the final offset of eventual `input_a`
543    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    // Take only bytes where bits were set
547    let num_bytes_with_data =
548        (y_size_bits(K) + parent_metadata_bits * 2).div_ceil(u8::BITS as usize);
549
550    // Only supports `K` from 15 to 25 (otherwise math will not be correct when concatenating y,
551    // left metadata and right metadata)
552    let hash = {
553        // Collect `K` most significant bits of `y` at the final offset of eventual `input_a`
554        let y_bits = u128::from(y) << (u128::BITS as usize - y_size_bits(K));
555
556        // Move bits of `left_metadata` at the final offset of eventual `input_a`
557        let left_metadata_bits =
558            left_metadata << (u128::BITS as usize - parent_metadata_bits - y_size_bits(K));
559
560        // If `right_metadata` bits start to the left of the desired position in `input_a` move
561        // bits right, else move left
562        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            // Collect bits of `right_metadata` that will fit into `input_a` at the final offset in
565            // eventual `input_a`
566            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            // Collect bits of `right_metadata` that will spill over into `input_b`
569            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        // For K up to 25 it is guaranteed that metadata + bit offset will always fit into u128.
594        // We collect the bytes necessary, potentially with extra bits at the start and end of the
595        // bytes that will be taken care of later.
596        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        // Remove extra bits at the beginning
602        let metadata = metadata << (y_size_bits(K) % u8::BITS as usize);
603        // Move bits into the correct location
604        metadata >> (u128::BITS as usize - metadata_size_bits)
605    } else {
606        0
607    };
608
609    (y_output, Metadata::from(metadata))
610}
611
612// TODO: This is actually using only pipelining rather than real SIMD (at least explicitly) due to:
613//  * https://github.com/rust-lang/portable-simd/issues/108
614//  * https://github.com/BLAKE3-team/BLAKE3/issues/478#issuecomment-3200106103
615#[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    // TODO: `u128` is not supported as SIMD element yet, see
628    //  https://github.com/rust-lang/portable-simd/issues/108
629    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    // Part of the `right_bits` at the final offset of eventual `input_a`
645    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    // Take only bytes where bits were set
649    let num_bytes_with_data =
650        (y_size_bits(K) + parent_metadata_bits * 2).div_ceil(u8::BITS as usize);
651
652    // Only supports `K` from 15 to 25 (otherwise math will not be correct when concatenating y,
653    // left metadata and right metadata)
654    // TODO: SIMD hashing once this is possible:
655    //  https://github.com/BLAKE3-team/BLAKE3/issues/478#issuecomment-3200106103
656    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            // Collect `K` most significant bits of `y` at the final offset of eventual
665            // `input_a`
666            let y_bits = u128::from(y) << (u128::BITS as usize - y_size_bits(K));
667
668            // Move bits of `left_metadata` at the final offset of eventual `input_a`
669            let left_metadata_bits =
670                left_metadata << (u128::BITS as usize - parent_metadata_bits - y_size_bits(K));
671
672            // If `right_metadata` bits start to the left of the desired position in `input_a` move
673            // bits right, else move left
674            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                // Collect bits of `right_metadata` that will fit into `input_a` at the final offset
677                // in eventual `input_a`
678                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                // Collect bits of `right_metadata` that will spill over into `input_b`
681                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        // For K up to 25 it is guaranteed that metadata + bit offset will always fit into u128.
713        // We collect the bytes necessary, potentially with extra bits at the start and end of the
714        // bytes that will be taken care of later.
715        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                // Remove extra bits at the beginning
725                let metadata = metadata << (y_size_bits(K) % u8::BITS as usize);
726                // Move bits into the correct location
727                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/// # Safety
740/// `m` must contain positions that correspond to the parent table
741#[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    // SAFETY: Guaranteed by function contract
750    let left_metadata = unsafe { parent_table.metadata(m.left_position) };
751    // SAFETY: Guaranteed by function contract
752    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/// # Safety
761/// `matches` must contain positions that correspond to the parent table
762#[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    // SAFETY: Guaranteed by function contract
783    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    // SAFETY: Guaranteed by function contract
793    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/// # Safety
824/// `matches` must contain positions that correspond to the parent table. `ys`, `position` and
825/// `metadatas` length must be at least the length of `matches`
826#[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        // SAFETY: Guaranteed by function contract
858        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        // The last table doesn't have metadata
864        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        // SAFETY: Guaranteed by function contract
875        let (y, p, metadata) = unsafe { match_to_result(parent_table, other_match) };
876        other_y.write(y);
877        other_positions.write(p);
878        // The last table doesn't have metadata
879        if metadata_size_bits(K, TABLE_NUMBER) > 0 {
880            other_metadata.write(metadata);
881        }
882    }
883}
884
885/// Similar to [`Table`], but smaller size for later processing stages
886#[cfg(feature = "alloc")]
887#[derive(Debug)]
888pub(super) enum PrunedTable<const K: u8, const TABLE_NUMBER: u8> {
889    First,
890    /// Other tables
891    Other {
892        /// Left and right entry positions in a previous table encoded into bits
893        positions: Box<[MaybeUninit<[Position; 2]>]>,
894    },
895    /// Other tables
896    #[cfg(feature = "parallel")]
897    OtherBuckets {
898        /// Left and right entry positions in a previous table encoded into bits.
899        ///
900        /// Only positions from the `buckets` field are guaranteed to be initialized.
901        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    /// Get `[left_position, right_position]` of a previous table for a specified position in a
909    /// current table.
910    ///
911    /// # Safety
912    /// `position` must come from [`Table::buckets()`] or [`Self::position()`] and not be a sentinel
913    /// value.
914    #[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                // SAFETY: All non-sentinel positions returned by [`Self::buckets()`] are valid
922                unsafe { positions.get_unchecked(usize::from(position)).assume_init() }
923            }
924            #[cfg(feature = "parallel")]
925            Self::OtherBuckets { positions, .. } => {
926                // SAFETY: All non-sentinel positions returned by [`Self::buckets()`] are valid
927                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 table
942    First {
943        /// Each bucket contains positions of `Y` values that belong to it and corresponding `y`.
944        ///
945        /// Buckets are padded with sentinel values to `REDUCED_BUCKETS_SIZE`.
946        buckets: Box<[[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]>,
947    },
948    /// Other tables
949    Other {
950        /// Left and right entry positions in a previous table encoded into bits
951        positions: Box<[MaybeUninit<[Position; 2]>]>,
952        /// Metadata corresponding to each entry
953        metadatas: Box<[MaybeUninit<Metadata<K, TABLE_NUMBER>>]>,
954        /// Each bucket contains positions of `Y` values that belong to it and corresponding `y`.
955        ///
956        /// Buckets are padded with sentinel values to `REDUCED_BUCKETS_SIZE`.
957        buckets: Box<[[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]>,
958    },
959    /// Other tables
960    #[cfg(feature = "parallel")]
961    OtherBuckets {
962        /// Left and right entry positions in a previous table encoded into bits.
963        ///
964        /// Only positions from the `buckets` field are guaranteed to be initialized.
965        positions:
966            Box<[[MaybeUninit<[Position; 2]>; REDUCED_MATCHES_COUNT]; NUM_BUCKET_PAIRS::<K>]>,
967        /// Metadata corresponding to each entry.
968        ///
969        /// Only positions from the `buckets` field are guaranteed to be initialized.
970        metadatas: Box<
971            [[MaybeUninit<Metadata<K, TABLE_NUMBER>>; REDUCED_MATCHES_COUNT];
972                NUM_BUCKET_PAIRS::<K>],
973        >,
974        /// Each bucket contains positions of `Y` values that belong to it and corresponding `y`.
975        ///
976        /// Buckets are padded with sentinel values to `REDUCED_BUCKETS_SIZE`.
977        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    /// Create the table
984    pub(super) fn create(seed: Seed) -> Self {
985        // `MAX_BUCKET_SIZE` is not actively used, but is an upper-bound reference for the other
986        // parameters
987        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        // SAFETY: Contents is `MaybeUninit`
995        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        // SAFETY: Converting a boxed array to a vector of the same size, which has the same memory
1013        // layout, all elements were initialized
1014        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        // TODO: Try to group buckets in the process of collecting `y`s
1021        let mut buckets = group_by_buckets::<K>(&ys);
1022        sort_buckets::<K>(&mut buckets);
1023
1024        Self::First { buckets }
1025    }
1026
1027    /// Create the table, leverages available parallelism
1028    #[cfg(feature = "parallel")]
1029    pub(super) fn create_parallel(seed: Seed) -> Self {
1030        // `MAX_BUCKET_SIZE` is not actively used, but is an upper-bound reference for the other
1031        // parameters
1032        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        // SAFETY: Contents is `MaybeUninit`
1040        let mut ys =
1041            unsafe { Box::<[MaybeUninit<Y>; MAX_TABLE_SIZE::<K>]>::new_uninit().assume_init() };
1042
1043        // TODO: Try parallelism here?
1044        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        // SAFETY: Converting a boxed array to a vector of the same size, which has the same memory
1059        // layout, all elements were initialized
1060        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        // TODO: Try to group buckets in the process of collecting `y`s
1067        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    /// Creates a new [`TABLE_NUMBER`] table. There also exists [`Self::create_parallel()`] that
1112    /// trades CPU efficiency and memory usage for lower latency and with multiple parallel calls,
1113    /// better overall performance.
1114    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        // The last table doesn't have metadata
1127        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            // SAFETY: Positions are taken from `Table::buckets()` and correspond to initialized
1135            // values
1136            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            // Throw away some successful matches that are not that necessary
1146            let matches = &matches[..matches.len().min(REDUCED_MATCHES_COUNT)];
1147            // SAFETY: Already initialized this many elements
1148            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            // SAFETY: arrays are sized to `max_table_entries(K)`, the exact upper bound on the total
1157            // number of matches, so `initialized_elements + matches.len()` never exceeds the length
1158            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            // SAFETY: Matches come from the parent table and the size of `ys`, `positions`
1167            // and `metadatas` is the same as the number of matches
1168            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        // SAFETY: Converting a boxed array to a vector of the same size, which has the same memory
1178        // layout, the number of elements matches the number of elements that were initialized
1179        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        // TODO: Try to group buckets in the process of collecting `y`s
1186        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    /// Almost the same as [`Self::create()`], but uses parallelism internally for better
1199    /// performance (though not efficiency of CPU and memory usage), if you create multiple tables
1200    /// in parallel, prefer [`Self::create_parallel()`] for better overall performance.
1201    #[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        // SAFETY: Contents is `MaybeUninit`
1210        let ys = unsafe {
1211            Box::<[SyncUnsafeCell<[MaybeUninit<_>; REDUCED_MATCHES_COUNT]>; NUM_BUCKET_PAIRS::<K>]>::new_uninit().assume_init()
1212        };
1213        // SAFETY: Contents is `MaybeUninit`
1214        let positions = unsafe {
1215            Box::<[SyncUnsafeCell<[MaybeUninit<_>; REDUCED_MATCHES_COUNT]>; NUM_BUCKET_PAIRS::<K>]>::new_uninit().assume_init()
1216        };
1217        // SAFETY: Contents is `MaybeUninit`
1218        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        // Iterate over buckets in batches, such that a cache line worth of bytes is taken from
1228        // `global_results_counts` each time to avoid unnecessary false sharing
1229        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                    // SAFETY: Positions are taken from `Table::buckets()` and correspond to
1249                    // initialized values
1250                    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                    // Throw away some successful matches that are not that necessary
1260                    let matches = &matches[..matches.len().min(REDUCED_MATCHES_COUNT)];
1261
1262                    // SAFETY: This is the only place where `left_bucket_index`'s entry is accessed
1263                    // at this time, and it is guaranteed to be in range
1264                    let ys = unsafe { &mut *ys.get_unchecked(left_bucket_index).get() };
1265                    // SAFETY: This is the only place where `left_bucket_index`'s entry is accessed
1266                    // at this time, and it is guaranteed to be in range
1267                    let positions =
1268                        unsafe { &mut *positions.get_unchecked(left_bucket_index).get() };
1269                    // SAFETY: This is the only place where `left_bucket_index`'s entry is accessed
1270                    // at this time, and it is guaranteed to be in range
1271                    let metadatas =
1272                        unsafe { &mut *metadatas.get_unchecked(left_bucket_index).get() };
1273                    // SAFETY: This is the only place where `left_bucket_index`'s entry is accessed
1274                    // at this time, and it is guaranteed to be in range
1275                    let count = unsafe {
1276                        &mut *global_results_counts.get_unchecked(left_bucket_index).get()
1277                    };
1278
1279                    // SAFETY: Matches come from the parent table and the size of `ys`, `positions`
1280                    // and `metadatas` is larger or equal to the number of matches
1281                    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        // TODO: Try to group buckets in the process of collecting `y`s
1302        // SAFETY: `global_results_counts` corresponds to the number of initialized `ys`
1303        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    /// Get `[left_position, right_position]` of a previous table for a specified position in a
1324    /// current table.
1325    ///
1326    /// # Safety
1327    /// `position` must come from [`Self::buckets()`] or [`Self::position()`] or
1328    /// [`PrunedTable::position()`] and not be a sentinel value.
1329    #[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                // SAFETY: All non-sentinel positions returned by [`Self::buckets()`] are valid
1337                unsafe { positions.get_unchecked(usize::from(position)).assume_init() }
1338            }
1339            #[cfg(feature = "parallel")]
1340            Self::OtherBuckets { positions, .. } => {
1341                // SAFETY: All non-sentinel positions returned by [`Self::buckets()`] are valid
1342                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    /// Returns `None` for an invalid position or for table number 7.
1359    ///
1360    /// # Safety
1361    /// `position` must come from [`Self::buckets()`] and not be a sentinel value.
1362    #[inline(always)]
1363    unsafe fn metadata(&self, position: Position) -> Metadata<K, TABLE_NUMBER> {
1364        match self {
1365            Self::First { .. } => {
1366                // X matches position
1367                Metadata::from(X::from(u32::from(position)))
1368            }
1369            Self::Other { metadatas, .. } => {
1370                // SAFETY: All non-sentinel positions returned by [`Self::buckets()`] are valid
1371                unsafe { metadatas.get_unchecked(usize::from(position)).assume_init() }
1372            }
1373            #[cfg(feature = "parallel")]
1374            Self::OtherBuckets { metadatas, .. } => {
1375                // SAFETY: All non-sentinel positions returned by [`Self::buckets()`] are valid
1376                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    /// Positions of `y`s grouped by the bucket they belong to
1400    #[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}