Skip to main content

ab_proof_of_space/chiapos/
table.rs

1#[cfg(any(feature = "alloc", test))]
2mod rmap;
3#[cfg(test)]
4mod tests;
5pub(super) mod types;
6
7use crate::chiapos::Seed;
8use crate::chiapos::constants::{PARAM_B, PARAM_BC, PARAM_C, PARAM_EXT, PARAM_M};
9#[cfg(feature = "alloc")]
10use crate::chiapos::table::rmap::Rmap;
11use crate::chiapos::table::types::{Metadata, X, Y};
12#[cfg(feature = "alloc")]
13use crate::chiapos::table::types::{Position, R};
14use ab_chacha8::{ChaCha8Block, ChaCha8State};
15#[cfg(feature = "alloc")]
16use ab_core_primitives::pieces::Record;
17#[cfg(feature = "alloc")]
18use alloc::boxed::Box;
19#[cfg(feature = "alloc")]
20use alloc::vec;
21#[cfg(feature = "alloc")]
22use alloc::vec::Vec;
23#[cfg(feature = "alloc")]
24use chacha20::cipher::{Iv, KeyIvInit, StreamCipher};
25#[cfg(feature = "alloc")]
26use chacha20::{ChaCha8, Key};
27use core::array;
28#[cfg(feature = "parallel")]
29use core::cell::SyncUnsafeCell;
30#[cfg(feature = "alloc")]
31use core::mem;
32#[cfg(feature = "alloc")]
33use core::mem::MaybeUninit;
34#[cfg(any(feature = "alloc", test))]
35use core::simd::prelude::*;
36#[cfg(feature = "parallel")]
37use core::sync::atomic::{AtomicUsize, Ordering};
38#[cfg(feature = "alloc")]
39use derive_more::Deref;
40#[cfg(feature = "alloc")]
41use rclite::Arc;
42#[cfg(any(feature = "alloc", test))]
43use seq_macro::seq;
44
45#[cfg(any(feature = "alloc", test))]
46const COMPUTE_F1_SIMD_FACTOR: usize = 8;
47#[cfg(any(feature = "alloc", test))]
48const COMPUTE_FN_SIMD_FACTOR: usize = 16;
49const MAX_BUCKET_SIZE: usize = 512;
50#[cfg(any(feature = "alloc", test))]
51const BUCKET_SIZE_UPPER_BOUND_SECURITY_BITS: u8 = 128;
52/// Reducing bucket size for better performance.
53///
54/// The number should be sufficient to produce enough proofs for sector encoding with high
55/// probability.
56const REDUCED_BUCKET_SIZE: usize = 272;
57/// Reducing matches count for better performance.
58///
59/// The number should be sufficient to produce enough proofs for sector encoding with high
60/// probability.
61const REDUCED_MATCHES_COUNT: usize = 288;
62#[cfg(feature = "parallel")]
63const CACHE_LINE_SIZE: usize = 64;
64
65const {
66    debug_assert!(REDUCED_BUCKET_SIZE <= MAX_BUCKET_SIZE);
67    debug_assert!(REDUCED_MATCHES_COUNT <= MAX_BUCKET_SIZE);
68}
69
70/// Compute the size of `y` in bits
71const fn y_size_bits(k: u8) -> usize {
72    usize::from(k) + usize::from(PARAM_EXT)
73}
74
75/// Metadata size in bits
76const fn metadata_size_bits(k: u8, table_number: u8) -> usize {
77    usize::from(k)
78        * match table_number {
79            1 => 1,
80            2 => 2,
81            3 | 4 => 4,
82            5 => 3,
83            6 => 2,
84            7 => 0,
85            _ => unreachable!(),
86        }
87}
88
89/// Number of buckets for a given `k`
90#[cfg(feature = "alloc")]
91const NUM_BUCKETS<const K: u8>: usize =
92    2_usize
93        .pow(y_size_bits(K) as u32)
94        .div_ceil(usize::from(PARAM_BC));
95#[cfg(feature = "parallel")]
96const NUM_BUCKET_PAIRS<const K: u8>: usize = NUM_BUCKETS::<K> - 1;
97
98/// Size of the first table and max size for other tables
99#[cfg(feature = "alloc")]
100const MAX_TABLE_SIZE<const K: u8>: usize = 1 << K;
101
102#[cfg(any(feature = "alloc", test))]
103const TABLE_1_YS_BATCH_SIMD<const K: u8>: usize =
104    usize::from(K) * COMPUTE_F1_SIMD_FACTOR / u8::BITS as usize;
105
106#[cfg(feature = "parallel")]
107#[inline(always)]
108fn strip_sync_unsafe_cell<const N: usize, T>(value: Box<[SyncUnsafeCell<T>; N]>) -> Box<[T; N]> {
109    // SAFETY: `SyncUnsafeCell` has the same layout as `T`
110    unsafe { Box::from_raw(Box::into_raw(value).cast()) }
111}
112
113/// ChaCha8 [`Vec`] sufficient for the whole first table for [`K`].
114/// Prefer [`partial_y`] if you need partial y just for a single `x`.
115#[cfg(feature = "alloc")]
116fn partial_ys<const K: u8>(seed: Seed) -> Vec<u8> {
117    let output_len_bits = usize::from(K) * (1 << K);
118    let mut output = vec![0; output_len_bits.div_ceil(u8::BITS as usize)];
119
120    let key = Key::from(seed);
121    let iv = Iv::<ChaCha8>::default();
122
123    let mut cipher = ChaCha8::new(&key, &iv);
124
125    cipher.write_keystream(&mut output);
126
127    output
128}
129
130/// Calculate a probabilistic upper bound on the Chia bucket size for a given `k` and
131/// `security_bits` (security level).
132///
133/// This is based on a Chernoff bound for the Poisson distribution with mean
134/// `lambda = PARAM_BC / 2^PARAM_EXT`, ensuring the probability that any bucket exceeds the bound is
135/// less than `2^{-security_bits}`.
136/// The bound is lambda + ceil(sqrt(3 * lambda * (k + security_bits) * ln(2))).
137#[cfg(feature = "alloc")]
138const fn bucket_size_upper_bound(k: u8, security_bits: u8) -> usize {
139    // Lambda is the expected number of entries in a bucket, approximated as
140    // `PARAM_BC / 2^PARAM_EXT`. It is independent of `k`.
141    const LAMBDA: u64 = PARAM_BC as u64 / 2u64.pow(PARAM_EXT as u32);
142    // Approximation of ln(2) as a fraction: ln(2) ≈ LN2_NUM / LN2_DEN.
143    // This allows integer-only computation of the square root term involving ln(2).
144    const LN2_NUM: u128 = 693_147;
145    const LN2_DEN: u128 = 1_000_000;
146
147    // `k + security_bits` for the union bound over ~2^k intervals
148    let ks = k as u128 + security_bits as u128;
149    // Compute numerator for the expression under the square root:
150    // `3 * lambda * (k + security_bits) * LN2_NUM`
151    let num = 3u128 * LAMBDA as u128 * ks * LN2_NUM;
152    // Denominator for ln(2): `LN2_DEN`
153    let den = LN2_DEN;
154
155    let ceil_div = num.div_ceil(den);
156
157    // Binary search to find the smallest `x` such that `x * x * den >= num`,
158    // which computes `ceil(sqrt(num / den))` without floating-point.
159    // We use a custom binary search over `u64` range because binary search in the standard library
160    // operates on sorted slices, not directly on integer ranges for solving inequalities like this.
161    let mut low = 0u64;
162    let mut high = u64::MAX;
163    while low < high {
164        let mid = low + (high - low) / 2;
165        let left = (mid as u128) * (mid as u128);
166        if left >= ceil_div {
167            high = mid;
168        } else {
169            low = mid + 1;
170        }
171    }
172    let add_term = low;
173
174    (LAMBDA + add_term) as usize
175}
176
177#[cfg(feature = "alloc")]
178fn group_by_buckets<const K: u8>(
179    ys: &[Y],
180) -> Box<[[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]> {
181    let mut bucket_offsets = [0_u16; NUM_BUCKETS::<K>];
182    // SAFETY: Contents is `MaybeUninit`
183    let mut buckets = unsafe {
184        Box::<[[MaybeUninit<(Position, Y)>; REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]>::new_uninit()
185            .assume_init()
186    };
187
188    for (&y, position) in ys.iter().zip(Position::ZERO..) {
189        let bucket_index = (u32::from(y) / u32::from(PARAM_BC)) as usize;
190
191        // SAFETY: Bucket is obtained using division by `PARAM_BC` and fits by definition
192        let bucket_offset = unsafe { bucket_offsets.get_unchecked_mut(bucket_index) };
193        // SAFETY: Bucket is obtained using division by `PARAM_BC` and fits by definition
194        let bucket = unsafe { buckets.get_unchecked_mut(bucket_index) };
195
196        if *bucket_offset < REDUCED_BUCKET_SIZE as u16 {
197            bucket[*bucket_offset as usize].write((position, y));
198            *bucket_offset += 1;
199        }
200    }
201
202    for (bucket, initialized) in buckets.iter_mut().zip(bucket_offsets) {
203        bucket[usize::from(initialized)..].write_filled((Position::SENTINEL, Y::SENTINEL));
204    }
205
206    // SAFETY: All entries are initialized
207    unsafe { Box::from_raw(Box::into_raw(buckets).cast()) }
208}
209
210/// Similar to [`group_by_buckets()`], but processes buckets instead of a flat list of `y`s.
211///
212/// # Safety
213/// Iterator item is a list of potentially uninitialized `y`s and a number of initialized `y`s. The
214/// number of initialized `y`s must be correct.
215#[cfg(feature = "parallel")]
216unsafe fn group_by_buckets_from_buckets<'a, const K: u8, I>(
217    iter: I,
218) -> Box<[[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]>
219where
220    I: Iterator<Item = (&'a [MaybeUninit<Y>; REDUCED_MATCHES_COUNT], usize)> + 'a,
221{
222    let mut bucket_offsets = [0_u16; NUM_BUCKETS::<K>];
223    // SAFETY: Contents is `MaybeUninit`
224    let mut buckets = unsafe {
225        Box::<[[MaybeUninit<(Position, Y)>; REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]>::new_uninit()
226            .assume_init()
227    };
228
229    for ((ys, count), batch_start) in iter.zip((Position::ZERO..).step_by(REDUCED_MATCHES_COUNT)) {
230        // SAFETY: Function contract guarantees that `y`s are initialized
231        let ys = unsafe { ys[..count].assume_init_ref() };
232        for (&y, position) in ys.iter().zip(batch_start..) {
233            let bucket_index = (u32::from(y) / u32::from(PARAM_BC)) as usize;
234
235            // SAFETY: Bucket is obtained using division by `PARAM_BC` and fits by definition
236            let bucket_offset = unsafe { bucket_offsets.get_unchecked_mut(bucket_index) };
237            // SAFETY: Bucket is obtained using division by `PARAM_BC` and fits by definition
238            let bucket = unsafe { buckets.get_unchecked_mut(bucket_index) };
239
240            if *bucket_offset < REDUCED_BUCKET_SIZE as u16 {
241                bucket[*bucket_offset as usize].write((position, y));
242                *bucket_offset += 1;
243            }
244        }
245    }
246
247    for (bucket, initialized) in buckets.iter_mut().zip(bucket_offsets) {
248        bucket[usize::from(initialized)..].write_filled((Position::SENTINEL, Y::SENTINEL));
249    }
250
251    // SAFETY: All entries are initialized
252    unsafe { Box::from_raw(Box::into_raw(buckets).cast()) }
253}
254
255#[cfg(feature = "alloc")]
256#[derive(Debug, Copy, Clone, Deref)]
257#[repr(align(64))]
258struct CacheLineAligned<T>(T);
259
260/// Mapping from `parity` to `r` to `m`
261#[cfg(feature = "alloc")]
262type LeftTargets =
263    [[CacheLineAligned<[R; const { usize::from(PARAM_M) }]>; const { usize::from(PARAM_BC) }]; 2];
264
265#[cfg(feature = "alloc")]
266fn calculate_left_targets() -> Arc<LeftTargets> {
267    let mut left_targets = Arc::<LeftTargets>::new_uninit();
268    // SAFETY: Same layout and uninitialized in both cases
269    let left_targets_slice = unsafe {
270        mem::transmute::<
271            &mut MaybeUninit<LeftTargets>,
272            &mut [[MaybeUninit<CacheLineAligned<[R; const { usize::from(PARAM_M) }]>>; const {
273                     usize::from(PARAM_BC)
274                 }]; 2],
275        >(Arc::get_mut_unchecked(&mut left_targets))
276    };
277
278    for parity in 0..=1 {
279        for r in 0..PARAM_BC {
280            let c = r / PARAM_C;
281
282            let arr = array::from_fn(|m| {
283                let m = m as u16;
284                R::from(
285                    ((c + m) % PARAM_B) * PARAM_C
286                        + (((2 * m + parity) * (2 * m + parity) + r) % PARAM_C),
287                )
288            });
289            left_targets_slice[usize::from(parity)][usize::from(r)].write(CacheLineAligned(arr));
290        }
291    }
292
293    // SAFETY: Initialized all entries
294    unsafe { left_targets.assume_init() }
295}
296
297fn calculate_left_target_on_demand(parity: u32, r: u32, m: u32) -> u32 {
298    let param_b = u32::from(PARAM_B);
299    let param_c = u32::from(PARAM_C);
300
301    ((r / param_c + m) % param_b) * param_c + (((2 * m + parity) * (2 * m + parity) + r) % param_c)
302}
303
304/// Caches that can be used to optimize the creation of multiple [`Tables`](super::Tables).
305#[cfg(feature = "alloc")]
306#[derive(Debug, Clone)]
307pub struct TablesCache {
308    left_targets: Arc<LeftTargets>,
309}
310
311#[cfg(feature = "alloc")]
312impl Default for TablesCache {
313    /// Create a new instance
314    fn default() -> Self {
315        Self {
316            left_targets: calculate_left_targets(),
317        }
318    }
319}
320
321#[cfg(feature = "alloc")]
322#[derive(Debug, Copy, Clone)]
323struct Match {
324    left_position: Position,
325    left_y: Y,
326    right_position: Position,
327}
328
329/// `partial_y_offset` is in bits within `partial_y`
330pub(super) fn compute_f1<const K: u8>(x: X, seed: &Seed) -> Y {
331    const U32S_PER_BLOCK: usize = size_of::<ChaCha8Block>() / size_of::<u32>();
332
333    let skip_bits = u32::from(K) * u32::from(x);
334    let skip_u32s = skip_bits / u32::BITS;
335    let partial_y_offset = skip_bits % u32::BITS;
336
337    let initial_state = ChaCha8State::init(seed, &[0; _]);
338    let first_block_counter = skip_u32s / U32S_PER_BLOCK as u32;
339    let u32_in_first_block = skip_u32s as usize % U32S_PER_BLOCK;
340
341    let first_block = initial_state.compute_block(first_block_counter);
342    let hi = first_block[u32_in_first_block].to_be();
343
344    // TODO: Is SIMD version of `compute_block()` that produces two blocks at once possible?
345    let lo = if u32_in_first_block + 1 == U32S_PER_BLOCK {
346        // Spilled over into the second block
347        let second_block = initial_state.compute_block(first_block_counter + 1);
348        second_block[0].to_be()
349    } else {
350        first_block[u32_in_first_block + 1].to_be()
351    };
352
353    let partial_y = (u64::from(hi) << u32::BITS) | u64::from(lo);
354
355    let pre_y = partial_y >> (u64::BITS - u32::from(K + PARAM_EXT) - partial_y_offset);
356    let pre_y = pre_y as u32;
357    // Mask for clearing the rest of bits of `pre_y`.
358    let pre_y_mask = (u32::MAX << PARAM_EXT) & (u32::MAX >> (u32::BITS - u32::from(K + PARAM_EXT)));
359
360    // Extract `PARAM_EXT` most significant bits from `x` and store in the final offset of
361    // eventual `y` with the rest of bits being zero (`x` is `0..2^K`)
362    let pre_ext = u32::from(x) >> (K - PARAM_EXT);
363
364    // Combine all of the bits together:
365    // [padding zero bits][`K` bits rom `partial_y`][`PARAM_EXT` bits from `x`]
366    Y::from((pre_y & pre_y_mask) | pre_ext)
367}
368
369#[cfg(any(feature = "alloc", test))]
370pub(super) fn compute_f1_simd<const K: u8>(
371    xs: Simd<u32, COMPUTE_F1_SIMD_FACTOR>,
372    partial_ys: &[u8; TABLE_1_YS_BATCH_SIMD::<K>],
373) -> [Y; COMPUTE_F1_SIMD_FACTOR] {
374    // Each element contains `K` desired bits of `partial_ys` in the final offset of eventual `ys`
375    // with the rest of bits being in undefined state
376    let pre_ys_bytes = array::from_fn(|i| {
377        let partial_y_offset = i * usize::from(K);
378        let partial_y_length =
379            (partial_y_offset % u8::BITS as usize + usize::from(K)).div_ceil(u8::BITS as usize);
380        let mut pre_y_bytes = 0u64.to_be_bytes();
381        pre_y_bytes[..partial_y_length].copy_from_slice(
382            &partial_ys[partial_y_offset / u8::BITS as usize..][..partial_y_length],
383        );
384
385        u64::from_be_bytes(pre_y_bytes)
386    });
387    let pre_ys_right_offset = array::from_fn(|i| {
388        let partial_y_offset = i as u32 * u32::from(K);
389        u64::from(u64::BITS - u32::from(K + PARAM_EXT) - partial_y_offset % u8::BITS)
390    });
391    let pre_ys = Simd::from_array(pre_ys_bytes) >> Simd::from_array(pre_ys_right_offset);
392
393    // Mask for clearing the rest of bits of `pre_ys`.
394    let pre_ys_mask = Simd::splat(
395        (u32::MAX << usize::from(PARAM_EXT))
396            & (u32::MAX >> (u32::BITS as usize - usize::from(K + PARAM_EXT))),
397    );
398
399    // Extract `PARAM_EXT` most significant bits from `xs` and store in the final offset of
400    // eventual `ys` with the rest of bits being in undefined state.
401    let pre_exts = xs >> Simd::splat(u32::from(K - PARAM_EXT));
402
403    // Combine all of the bits together:
404    // [padding zero bits][`K` bits rom `partial_y`][`PARAM_EXT` bits from `x`]
405    let ys = (pre_ys.cast() & pre_ys_mask) | pre_exts;
406
407    Y::array_from_repr(ys.to_array())
408}
409
410/// For verification use [`has_match`] instead.
411///
412/// # Safety
413/// Left and right bucket positions must correspond to the parent table.
414// TODO: Try to reduce the `matches` size further by processing `left_bucket` in chunks (like halves
415//  for example)
416#[cfg(feature = "alloc")]
417unsafe fn find_matches_in_buckets<'a>(
418    left_bucket_index: u32,
419    left_bucket: &[(Position, Y); REDUCED_BUCKET_SIZE],
420    right_bucket: &[(Position, Y); REDUCED_BUCKET_SIZE],
421    // `PARAM_M * 2` corresponds to the upper bound number of matches a single `y` in the
422    // left bucket might have here
423    matches: &'a mut [MaybeUninit<Match>; REDUCED_MATCHES_COUNT + usize::from(PARAM_M) * 2],
424    left_targets: &LeftTargets,
425) -> &'a [Match] {
426    let left_base = left_bucket_index * u32::from(PARAM_BC);
427    let right_base = left_base + u32::from(PARAM_BC);
428
429    let mut rmap = Rmap::new();
430    for &(right_position, y) in right_bucket {
431        // TODO: Wouldn't it make more sense to check the size here instead of sentinel?
432        if right_position == Position::SENTINEL {
433            break;
434        }
435        let r = R::from((u32::from(y) - right_base) as u16);
436        // SAFETY: `r` is within `0..PARAM_BC` range by definition, the right bucket is limited to
437        // `REDUCED_BUCKETS_SIZE`
438        unsafe {
439            rmap.add(r, right_position);
440        }
441    }
442
443    let parity = left_base % 2;
444    let left_targets_parity = &left_targets[parity as usize];
445    let mut next_match_index = 0;
446
447    // TODO: Simd read for left bucket? It might be more efficient in terms of memory access to
448    //  process chunks of the left bucket against one right value for each at a time
449    for &(left_position, y) in left_bucket {
450        // TODO: Wouldn't it make more sense to check the size here instead of sentinel?
451        // `next_match_index >= REDUCED_MATCHES_COUNT` is crucial to make sure
452        if left_position == Position::SENTINEL || next_match_index >= REDUCED_MATCHES_COUNT {
453            // Sentinel values are padded to the end of the bucket
454            break;
455        }
456
457        let r = R::from((u32::from(y) - left_base) as u16);
458        // SAFETY: `r` is within a bucket and exists by definition
459        let left_targets_r = unsafe { left_targets_parity.get_unchecked(usize::from(r)) };
460
461        for &r_target in left_targets_r.iter() {
462            // SAFETY: Targets are always limited to `PARAM_BC`
463            let [right_position_a, right_position_b] = unsafe { rmap.get(r_target) };
464
465            if right_position_a != Position::SENTINEL {
466                // SAFETY: Iteration will stop before `REDUCED_MATCHES_COUNT + PARAM_M * 2`
467                // elements is inserted
468                unsafe { matches.get_unchecked_mut(next_match_index) }.write(Match {
469                    left_position,
470                    left_y: y,
471                    right_position: right_position_a,
472                });
473                next_match_index += 1;
474
475                if right_position_b != Position::SENTINEL {
476                    // SAFETY: Iteration will stop before
477                    // `REDUCED_MATCHES_COUNT + PARAM_M * 2` elements is inserted
478                    unsafe { matches.get_unchecked_mut(next_match_index) }.write(Match {
479                        left_position,
480                        left_y: y,
481                        right_position: right_position_b,
482                    });
483                    next_match_index += 1;
484                }
485            }
486        }
487    }
488
489    // SAFETY: Initialized this many matches
490    unsafe { matches[..next_match_index].assume_init_ref() }
491}
492
493/// Simplified version of [`find_matches_in_buckets`] for verification purposes.
494pub(super) fn has_match(left_y: Y, right_y: Y) -> bool {
495    let right_r = u32::from(right_y) % u32::from(PARAM_BC);
496    let parity = (u32::from(left_y) / u32::from(PARAM_BC)) % 2;
497    let left_r = u32::from(left_y) % u32::from(PARAM_BC);
498
499    let r_targets = array::from_fn::<_, const { usize::from(PARAM_M) }, _>(|i| {
500        calculate_left_target_on_demand(parity, left_r, i as u32)
501    });
502
503    r_targets.contains(&right_r)
504}
505
506#[inline(always)]
507pub(super) fn compute_fn<const K: u8, const TABLE_NUMBER: u8, const PARENT_TABLE_NUMBER: u8>(
508    y: Y,
509    left_metadata: Metadata<K, PARENT_TABLE_NUMBER>,
510    right_metadata: Metadata<K, PARENT_TABLE_NUMBER>,
511) -> (Y, Metadata<K, TABLE_NUMBER>) {
512    let left_metadata = u128::from(left_metadata);
513    let right_metadata = u128::from(right_metadata);
514
515    let parent_metadata_bits = metadata_size_bits(K, PARENT_TABLE_NUMBER);
516
517    // Part of the `right_bits` at the final offset of eventual `input_a`
518    let y_and_left_bits = y_size_bits(K) + parent_metadata_bits;
519    let right_bits_start_offset = u128::BITS as usize - parent_metadata_bits;
520
521    // Take only bytes where bits were set
522    let num_bytes_with_data =
523        (y_size_bits(K) + parent_metadata_bits * 2).div_ceil(u8::BITS as usize);
524
525    // Only supports `K` from 15 to 25 (otherwise math will not be correct when concatenating y,
526    // left metadata and right metadata)
527    let hash = {
528        // Collect `K` most significant bits of `y` at the final offset of eventual `input_a`
529        let y_bits = u128::from(y) << (u128::BITS as usize - y_size_bits(K));
530
531        // Move bits of `left_metadata` at the final offset of eventual `input_a`
532        let left_metadata_bits =
533            left_metadata << (u128::BITS as usize - parent_metadata_bits - y_size_bits(K));
534
535        // If `right_metadata` bits start to the left of the desired position in `input_a` move
536        // bits right, else move left
537        if right_bits_start_offset < y_and_left_bits {
538            let right_bits_pushed_into_input_b = y_and_left_bits - right_bits_start_offset;
539            // Collect bits of `right_metadata` that will fit into `input_a` at the final offset in
540            // eventual `input_a`
541            let right_bits_a = right_metadata >> right_bits_pushed_into_input_b;
542            let input_a = y_bits | left_metadata_bits | right_bits_a;
543            // Collect bits of `right_metadata` that will spill over into `input_b`
544            let input_b = right_metadata << (u128::BITS as usize - right_bits_pushed_into_input_b);
545
546            let input = [input_a.to_be_bytes(), input_b.to_be_bytes()];
547            let input_len =
548                size_of::<u128>() + right_bits_pushed_into_input_b.div_ceil(u8::BITS as usize);
549            ab_blake3::single_block_hash(&input.as_flattened()[..input_len])
550                .expect("Exactly a single block worth of bytes; qed")
551        } else {
552            let right_bits_a = right_metadata << (right_bits_start_offset - y_and_left_bits);
553            let input_a = y_bits | left_metadata_bits | right_bits_a;
554
555            ab_blake3::single_block_hash(&input_a.to_be_bytes()[..num_bytes_with_data])
556                .expect("Less than a single block worth of bytes; qed")
557        }
558    };
559
560    let y_output = Y::from(
561        u32::from_be_bytes([hash[0], hash[1], hash[2], hash[3]])
562            >> (u32::BITS as usize - y_size_bits(K)),
563    );
564
565    let metadata_size_bits = metadata_size_bits(K, TABLE_NUMBER);
566
567    let metadata = if TABLE_NUMBER < 4 {
568        (left_metadata << parent_metadata_bits) | right_metadata
569    } else if metadata_size_bits > 0 {
570        // For K up to 25 it is guaranteed that metadata + bit offset will always fit into u128.
571        // We collect the bytes necessary, potentially with extra bits at the start and end of the
572        // bytes that will be taken care of later.
573        let metadata = u128::from_be_bytes(
574            hash[y_size_bits(K) / u8::BITS as usize..][..size_of::<u128>()]
575                .try_into()
576                .expect("Always enough bits for any K; qed"),
577        );
578        // Remove extra bits at the beginning
579        let metadata = metadata << (y_size_bits(K) % u8::BITS as usize);
580        // Move bits into the correct location
581        metadata >> (u128::BITS as usize - metadata_size_bits)
582    } else {
583        0
584    };
585
586    (y_output, Metadata::from(metadata))
587}
588
589// TODO: This is actually using only pipelining rather than real SIMD (at least explicitly) due to:
590//  * https://github.com/rust-lang/portable-simd/issues/108
591//  * https://github.com/BLAKE3-team/BLAKE3/issues/478#issuecomment-3200106103
592#[cfg(any(feature = "alloc", test))]
593fn compute_fn_simd<const K: u8, const TABLE_NUMBER: u8, const PARENT_TABLE_NUMBER: u8>(
594    left_ys: [Y; COMPUTE_FN_SIMD_FACTOR],
595    left_metadatas: [Metadata<K, PARENT_TABLE_NUMBER>; COMPUTE_FN_SIMD_FACTOR],
596    right_metadatas: [Metadata<K, PARENT_TABLE_NUMBER>; COMPUTE_FN_SIMD_FACTOR],
597) -> (
598    Simd<u32, COMPUTE_FN_SIMD_FACTOR>,
599    [Metadata<K, TABLE_NUMBER>; COMPUTE_FN_SIMD_FACTOR],
600) {
601    let parent_metadata_bits = metadata_size_bits(K, PARENT_TABLE_NUMBER);
602    let metadata_size_bits = metadata_size_bits(K, TABLE_NUMBER);
603
604    // TODO: `u128` is not supported as SIMD element yet, see
605    //  https://github.com/rust-lang/portable-simd/issues/108
606    let left_metadatas: [u128; COMPUTE_FN_SIMD_FACTOR] = seq!(N in 0..16 {
607        [
608        #(
609            u128::from(left_metadatas[N]),
610        )*
611        ]
612    });
613    let right_metadatas: [u128; COMPUTE_FN_SIMD_FACTOR] = seq!(N in 0..16 {
614        [
615        #(
616            u128::from(right_metadatas[N]),
617        )*
618        ]
619    });
620
621    // Part of the `right_bits` at the final offset of eventual `input_a`
622    let y_and_left_bits = y_size_bits(K) + parent_metadata_bits;
623    let right_bits_start_offset = u128::BITS as usize - parent_metadata_bits;
624
625    // Take only bytes where bits were set
626    let num_bytes_with_data =
627        (y_size_bits(K) + parent_metadata_bits * 2).div_ceil(u8::BITS as usize);
628
629    // Only supports `K` from 15 to 25 (otherwise math will not be correct when concatenating y,
630    // left metadata and right metadata)
631    // TODO: SIMD hashing once this is possible:
632    //  https://github.com/BLAKE3-team/BLAKE3/issues/478#issuecomment-3200106103
633    let hashes: [_; COMPUTE_FN_SIMD_FACTOR] = seq!(N in 0..16 {
634        [
635        #(
636        {
637            let y = left_ys[N];
638            let left_metadata = left_metadatas[N];
639            let right_metadata = right_metadatas[N];
640
641            // Collect `K` most significant bits of `y` at the final offset of eventual
642            // `input_a`
643            let y_bits = u128::from(y) << (u128::BITS as usize - y_size_bits(K));
644
645            // Move bits of `left_metadata` at the final offset of eventual `input_a`
646            let left_metadata_bits =
647                left_metadata << (u128::BITS as usize - parent_metadata_bits - y_size_bits(K));
648
649            // If `right_metadata` bits start to the left of the desired position in `input_a` move
650            // bits right, else move left
651            if right_bits_start_offset < y_and_left_bits {
652                let right_bits_pushed_into_input_b = y_and_left_bits - right_bits_start_offset;
653                // Collect bits of `right_metadata` that will fit into `input_a` at the final offset
654                // in eventual `input_a`
655                let right_bits_a = right_metadata >> right_bits_pushed_into_input_b;
656                let input_a = y_bits | left_metadata_bits | right_bits_a;
657                // Collect bits of `right_metadata` that will spill over into `input_b`
658                let input_b = right_metadata << (u128::BITS as usize - right_bits_pushed_into_input_b);
659
660                let input = [input_a.to_be_bytes(), input_b.to_be_bytes()];
661                let input_len =
662                    size_of::<u128>() + right_bits_pushed_into_input_b.div_ceil(u8::BITS as usize);
663                ab_blake3::single_block_hash(&input.as_flattened()[..input_len])
664                    .expect("Exactly a single block worth of bytes; qed")
665            } else {
666                let right_bits_a = right_metadata << (right_bits_start_offset - y_and_left_bits);
667                let input_a = y_bits | left_metadata_bits | right_bits_a;
668
669                ab_blake3::single_block_hash(&input_a.to_be_bytes()[..num_bytes_with_data])
670                    .expect("Exactly a single block worth of bytes; qed")
671            }
672        },
673        )*
674        ]
675    });
676
677    let y_outputs = Simd::from_array(
678        hashes.map(|hash| u32::from_be_bytes([hash[0], hash[1], hash[2], hash[3]])),
679    ) >> (u32::BITS - y_size_bits(K) as u32);
680
681    let metadatas = if TABLE_NUMBER < 4 {
682        seq!(N in 0..16 {
683            [
684            #(
685                Metadata::from((left_metadatas[N] << parent_metadata_bits) | right_metadatas[N]),
686            )*
687            ]
688        })
689    } else if metadata_size_bits > 0 {
690        // For K up to 25 it is guaranteed that metadata + bit offset will always fit into u128.
691        // We collect the bytes necessary, potentially with extra bits at the start and end of the
692        // bytes that will be taken care of later.
693        seq!(N in 0..16 {
694            [
695            #(
696            {
697                let metadata = u128::from_be_bytes(
698                    hashes[N][y_size_bits(K) / u8::BITS as usize..][..size_of::<u128>()]
699                        .try_into()
700                        .expect("Always enough bits for any K; qed"),
701                );
702                // Remove extra bits at the beginning
703                let metadata = metadata << (y_size_bits(K) % u8::BITS as usize);
704                // Move bits into the correct location
705                Metadata::from(metadata >> (u128::BITS as usize - metadata_size_bits))
706            },
707            )*
708            ]
709        })
710    } else {
711        [Metadata::default(); _]
712    };
713
714    (y_outputs, metadatas)
715}
716
717/// # Safety
718/// `m` must contain positions that correspond to the parent table
719#[cfg(feature = "alloc")]
720#[inline(always)]
721unsafe fn match_to_result<const K: u8, const TABLE_NUMBER: u8, const PARENT_TABLE_NUMBER: u8>(
722    parent_table: &Table<K, PARENT_TABLE_NUMBER>,
723    m: &Match,
724) -> (Y, [Position; 2], Metadata<K, TABLE_NUMBER>)
725where
726    Table<K, PARENT_TABLE_NUMBER>: private::NotLastTable,
727{
728    // SAFETY: Guaranteed by function contract
729    let left_metadata = unsafe { parent_table.metadata(m.left_position) };
730    // SAFETY: Guaranteed by function contract
731    let right_metadata = unsafe { parent_table.metadata(m.right_position) };
732
733    let (y, metadata) =
734        compute_fn::<K, TABLE_NUMBER, PARENT_TABLE_NUMBER>(m.left_y, left_metadata, right_metadata);
735
736    (y, [m.left_position, m.right_position], metadata)
737}
738
739/// # Safety
740/// `matches` must contain positions that correspond to the parent table
741#[cfg(feature = "alloc")]
742#[inline(always)]
743unsafe fn match_to_result_simd<const K: u8, const TABLE_NUMBER: u8, const PARENT_TABLE_NUMBER: u8>(
744    parent_table: &Table<K, PARENT_TABLE_NUMBER>,
745    matches: &[Match; COMPUTE_FN_SIMD_FACTOR],
746) -> (
747    Simd<u32, COMPUTE_FN_SIMD_FACTOR>,
748    [[Position; 2]; COMPUTE_FN_SIMD_FACTOR],
749    [Metadata<K, TABLE_NUMBER>; COMPUTE_FN_SIMD_FACTOR],
750)
751where
752    Table<K, PARENT_TABLE_NUMBER>: private::NotLastTable,
753{
754    let left_ys: [_; COMPUTE_FN_SIMD_FACTOR] = seq!(N in 0..16 {
755        [
756        #(
757            matches[N].left_y,
758        )*
759        ]
760    });
761    // SAFETY: Guaranteed by function contract
762    let left_metadatas: [_; COMPUTE_FN_SIMD_FACTOR] = unsafe {
763        seq!(N in 0..16 {
764            [
765            #(
766                parent_table.metadata(matches[N].left_position),
767            )*
768            ]
769        })
770    };
771    // SAFETY: Guaranteed by function contract
772    let right_metadatas: [_; COMPUTE_FN_SIMD_FACTOR] = unsafe {
773        seq!(N in 0..16 {
774            [
775            #(
776                parent_table.metadata(matches[N].right_position),
777            )*
778            ]
779        })
780    };
781
782    let (y_outputs, metadatas) = compute_fn_simd::<K, TABLE_NUMBER, PARENT_TABLE_NUMBER>(
783        left_ys,
784        left_metadatas,
785        right_metadatas,
786    );
787
788    let positions = seq!(N in 0..16 {
789        [
790        #(
791            [
792                matches[N].left_position,
793                matches[N].right_position,
794            ],
795        )*
796        ]
797    });
798
799    (y_outputs, positions, metadatas)
800}
801
802/// # Safety
803/// `matches` must contain positions that correspond to the parent table. `ys`, `position` and
804/// `metadatas` length must be at least the length of `matches`
805#[cfg(feature = "alloc")]
806#[inline(always)]
807unsafe fn matches_to_results<const K: u8, const TABLE_NUMBER: u8, const PARENT_TABLE_NUMBER: u8>(
808    parent_table: &Table<K, PARENT_TABLE_NUMBER>,
809    matches: &[Match],
810    ys: &mut [MaybeUninit<Y>],
811    positions: &mut [MaybeUninit<[Position; 2]>],
812    metadatas: &mut [MaybeUninit<Metadata<K, TABLE_NUMBER>>],
813) where
814    Table<K, PARENT_TABLE_NUMBER>: private::NotLastTable,
815{
816    let (grouped_matches, other_matches) = matches.as_chunks::<COMPUTE_FN_SIMD_FACTOR>();
817    let (grouped_ys, other_ys) = ys.split_at_mut(grouped_matches.as_flattened().len());
818    let grouped_ys = grouped_ys.as_chunks_mut::<COMPUTE_FN_SIMD_FACTOR>().0;
819    let (grouped_positions, other_positions) =
820        positions.split_at_mut(grouped_matches.as_flattened().len());
821    let grouped_positions = grouped_positions
822        .as_chunks_mut::<COMPUTE_FN_SIMD_FACTOR>()
823        .0;
824    let (grouped_metadatas, other_metadatas) =
825        metadatas.split_at_mut(grouped_matches.as_flattened().len());
826    let grouped_metadatas = grouped_metadatas
827        .as_chunks_mut::<COMPUTE_FN_SIMD_FACTOR>()
828        .0;
829
830    for (((grouped_matches, grouped_ys), grouped_positions), grouped_metadatas) in grouped_matches
831        .iter()
832        .zip(grouped_ys)
833        .zip(grouped_positions)
834        .zip(grouped_metadatas)
835    {
836        // SAFETY: Guaranteed by function contract
837        let (ys_group, positions_group, metadatas_group) =
838            unsafe { match_to_result_simd(parent_table, grouped_matches) };
839        let ys_group = Y::array_from_repr(ys_group.to_array());
840        grouped_ys.write_copy_of_slice(&ys_group);
841        grouped_positions.write_copy_of_slice(&positions_group);
842
843        // The last table doesn't have metadata
844        if metadata_size_bits(K, TABLE_NUMBER) > 0 {
845            grouped_metadatas.write_copy_of_slice(&metadatas_group);
846        }
847    }
848    for (((other_match, other_y), other_positions), other_metadata) in other_matches
849        .iter()
850        .zip(other_ys)
851        .zip(other_positions)
852        .zip(other_metadatas)
853    {
854        // SAFETY: Guaranteed by function contract
855        let (y, p, metadata) = unsafe { match_to_result(parent_table, other_match) };
856        other_y.write(y);
857        other_positions.write(p);
858        // The last table doesn't have metadata
859        if metadata_size_bits(K, TABLE_NUMBER) > 0 {
860            other_metadata.write(metadata);
861        }
862    }
863}
864
865/// Similar to [`Table`], but smaller size for later processing stages
866#[cfg(feature = "alloc")]
867#[derive(Debug)]
868pub(super) enum PrunedTable<const K: u8, const TABLE_NUMBER: u8> {
869    First,
870    /// Other tables
871    Other {
872        /// Left and right entry positions in a previous table encoded into bits
873        positions: Box<[MaybeUninit<[Position; 2]>; MAX_TABLE_SIZE::<K>]>,
874    },
875    /// Other tables
876    #[cfg(feature = "parallel")]
877    OtherBuckets {
878        /// Left and right entry positions in a previous table encoded into bits.
879        ///
880        /// Only positions from the `buckets` field are guaranteed to be initialized.
881        positions:
882            Box<[[MaybeUninit<[Position; 2]>; REDUCED_MATCHES_COUNT]; NUM_BUCKET_PAIRS::<K>]>,
883    },
884}
885
886#[cfg(feature = "alloc")]
887impl<const K: u8, const TABLE_NUMBER: u8> PrunedTable<K, TABLE_NUMBER> {
888    /// Get `[left_position, right_position]` of a previous table for a specified position in a
889    /// current table.
890    ///
891    /// # Safety
892    /// `position` must come from [`Table::buckets()`] or [`Self::position()`] and not be a sentinel
893    /// value.
894    #[inline(always)]
895    pub(super) unsafe fn position(&self, position: Position) -> [Position; 2] {
896        match self {
897            Self::First => {
898                unreachable!("Not the first table");
899            }
900            Self::Other { positions, .. } => {
901                // SAFETY: All non-sentinel positions returned by [`Self::buckets()`] are valid
902                unsafe { positions.get_unchecked(usize::from(position)).assume_init() }
903            }
904            #[cfg(feature = "parallel")]
905            Self::OtherBuckets { positions, .. } => {
906                // SAFETY: All non-sentinel positions returned by [`Self::buckets()`] are valid
907                unsafe {
908                    positions
909                        .as_flattened()
910                        .get_unchecked(usize::from(position))
911                        .assume_init()
912                }
913            }
914        }
915    }
916}
917
918#[cfg(feature = "alloc")]
919#[derive(Debug)]
920pub(super) enum Table<const K: u8, const TABLE_NUMBER: u8> {
921    /// First table
922    First {
923        /// Each bucket contains positions of `Y` values that belong to it and corresponding `y`.
924        ///
925        /// Buckets are padded with sentinel values to `REDUCED_BUCKETS_SIZE`.
926        buckets: Box<[[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]>,
927    },
928    /// Other tables
929    Other {
930        /// Left and right entry positions in a previous table encoded into bits
931        positions: Box<[MaybeUninit<[Position; 2]>; MAX_TABLE_SIZE::<K>]>,
932        /// Metadata corresponding to each entry
933        metadatas: Box<[MaybeUninit<Metadata<K, TABLE_NUMBER>>; MAX_TABLE_SIZE::<K>]>,
934        /// Each bucket contains positions of `Y` values that belong to it and corresponding `y`.
935        ///
936        /// Buckets are padded with sentinel values to `REDUCED_BUCKETS_SIZE`.
937        buckets: Box<[[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>]>,
938    },
939    /// Other tables
940    #[cfg(feature = "parallel")]
941    OtherBuckets {
942        /// Left and right entry positions in a previous table encoded into bits.
943        ///
944        /// Only positions from the `buckets` field are guaranteed to be initialized.
945        positions:
946            Box<[[MaybeUninit<[Position; 2]>; REDUCED_MATCHES_COUNT]; NUM_BUCKET_PAIRS::<K>]>,
947        /// Metadata corresponding to each entry.
948        ///
949        /// Only positions from the `buckets` field are guaranteed to be initialized.
950        metadatas: Box<
951            [[MaybeUninit<Metadata<K, TABLE_NUMBER>>; REDUCED_MATCHES_COUNT];
952                NUM_BUCKET_PAIRS::<K>],
953        >,
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}
960
961#[cfg(feature = "alloc")]
962impl<const K: u8> Table<K, 1> {
963    /// Create the table
964    pub(super) fn create(seed: Seed) -> Self {
965        // `MAX_BUCKET_SIZE` is not actively used, but is an upper-bound reference for the other
966        // parameters
967        debug_assert!(
968            MAX_BUCKET_SIZE >= bucket_size_upper_bound(K, BUCKET_SIZE_UPPER_BOUND_SECURITY_BITS),
969            "Max bucket size is not sufficiently large"
970        );
971
972        let partial_ys = partial_ys::<K>(seed);
973
974        // SAFETY: Contents is `MaybeUninit`
975        let mut ys =
976            unsafe { Box::<[MaybeUninit<Y>; MAX_TABLE_SIZE::<K>]>::new_uninit().assume_init() };
977
978        for ((ys, xs_batch_start), partial_ys) in ys
979            .as_chunks_mut::<COMPUTE_F1_SIMD_FACTOR>()
980            .0
981            .iter_mut()
982            .zip((X::ZERO..).step_by(COMPUTE_F1_SIMD_FACTOR))
983            .zip(partial_ys.as_chunks::<{ TABLE_1_YS_BATCH_SIMD::<K> }>().0)
984        {
985            let xs = Simd::splat(u32::from(xs_batch_start))
986                + Simd::from_array(array::from_fn(|i| i as u32));
987            let ys_batch = compute_f1_simd::<K>(xs, partial_ys);
988
989            ys.write_copy_of_slice(&ys_batch);
990        }
991
992        // SAFETY: All elements were initialized
993        let ys = unsafe { ys.assume_init_ref() };
994
995        // TODO: Try to group buckets in the process of collecting `y`s
996        let buckets = group_by_buckets::<K>(ys);
997
998        Self::First { buckets }
999    }
1000
1001    /// Create the table, leverages available parallelism
1002    #[cfg(feature = "parallel")]
1003    pub(super) fn create_parallel(seed: Seed) -> Self {
1004        // `MAX_BUCKET_SIZE` is not actively used, but is an upper-bound reference for the other
1005        // parameters
1006        debug_assert!(
1007            MAX_BUCKET_SIZE >= bucket_size_upper_bound(K, BUCKET_SIZE_UPPER_BOUND_SECURITY_BITS),
1008            "Max bucket size is not sufficiently large"
1009        );
1010
1011        let partial_ys = partial_ys::<K>(seed);
1012
1013        // SAFETY: Contents is `MaybeUninit`
1014        let mut ys =
1015            unsafe { Box::<[MaybeUninit<Y>; MAX_TABLE_SIZE::<K>]>::new_uninit().assume_init() };
1016
1017        // TODO: Try parallelism here?
1018        for ((ys, xs_batch_start), partial_ys) in ys
1019            .as_chunks_mut::<COMPUTE_F1_SIMD_FACTOR>()
1020            .0
1021            .iter_mut()
1022            .zip((X::ZERO..).step_by(COMPUTE_F1_SIMD_FACTOR))
1023            .zip(partial_ys.as_chunks::<{ TABLE_1_YS_BATCH_SIMD::<K> }>().0)
1024        {
1025            let xs = Simd::splat(u32::from(xs_batch_start))
1026                + Simd::from_array(array::from_fn(|i| i as u32));
1027            let ys_batch = compute_f1_simd::<K>(xs, partial_ys);
1028
1029            ys.write_copy_of_slice(&ys_batch);
1030        }
1031
1032        // SAFETY: All elements were initialized
1033        let ys = unsafe { ys.assume_init_ref() };
1034
1035        // TODO: Try to group buckets in the process of collecting `y`s
1036        let buckets = group_by_buckets::<K>(ys);
1037
1038        Self::First { buckets }
1039    }
1040}
1041
1042#[cfg(feature = "alloc")]
1043#[expect(clippy::inline_modules, reason = "Intentional tiny module inline")]
1044mod private {
1045    pub(in super::super) trait SupportedOtherTables {}
1046    pub(in super::super) trait NotLastTable {}
1047}
1048
1049#[cfg(feature = "alloc")]
1050impl<const K: u8> private::SupportedOtherTables for Table<K, 2> {}
1051#[cfg(feature = "alloc")]
1052impl<const K: u8> private::SupportedOtherTables for Table<K, 3> {}
1053#[cfg(feature = "alloc")]
1054impl<const K: u8> private::SupportedOtherTables for Table<K, 4> {}
1055#[cfg(feature = "alloc")]
1056impl<const K: u8> private::SupportedOtherTables for Table<K, 5> {}
1057#[cfg(feature = "alloc")]
1058impl<const K: u8> private::SupportedOtherTables for Table<K, 6> {}
1059#[cfg(feature = "alloc")]
1060impl<const K: u8> private::SupportedOtherTables for Table<K, 7> {}
1061
1062#[cfg(feature = "alloc")]
1063impl<const K: u8> private::NotLastTable for Table<K, 1> {}
1064#[cfg(feature = "alloc")]
1065impl<const K: u8> private::NotLastTable for Table<K, 2> {}
1066#[cfg(feature = "alloc")]
1067impl<const K: u8> private::NotLastTable for Table<K, 3> {}
1068#[cfg(feature = "alloc")]
1069impl<const K: u8> private::NotLastTable for Table<K, 4> {}
1070#[cfg(feature = "alloc")]
1071impl<const K: u8> private::NotLastTable for Table<K, 5> {}
1072#[cfg(feature = "alloc")]
1073impl<const K: u8> private::NotLastTable for Table<K, 6> {}
1074
1075#[cfg(feature = "alloc")]
1076impl<const K: u8, const TABLE_NUMBER: u8> Table<K, TABLE_NUMBER>
1077where
1078    Self: private::SupportedOtherTables,
1079{
1080    /// Creates a new [`TABLE_NUMBER`] table. There also exists [`Self::create_parallel()`] that
1081    /// trades CPU efficiency and memory usage for lower latency and with multiple parallel calls,
1082    /// better overall performance.
1083    pub(super) fn create<const PARENT_TABLE_NUMBER: u8>(
1084        parent_table: Table<K, PARENT_TABLE_NUMBER>,
1085        cache: &TablesCache,
1086    ) -> (Self, PrunedTable<K, PARENT_TABLE_NUMBER>)
1087    where
1088        Table<K, PARENT_TABLE_NUMBER>: private::NotLastTable,
1089    {
1090        let left_targets = &*cache.left_targets;
1091        let mut initialized_elements = 0_usize;
1092        // SAFETY: Contents is `MaybeUninit`
1093        let mut ys =
1094            unsafe { Box::<[MaybeUninit<Y>; MAX_TABLE_SIZE::<K>]>::new_uninit().assume_init() };
1095        // SAFETY: Contents is `MaybeUninit`
1096        let mut positions = unsafe {
1097            Box::<[MaybeUninit<[Position; 2]>; MAX_TABLE_SIZE::<K>]>::new_uninit().assume_init()
1098        };
1099        // SAFETY: Contents is `MaybeUninit`
1100        let mut metadatas = unsafe {
1101            Box::<[MaybeUninit<Metadata<K, TABLE_NUMBER>>; MAX_TABLE_SIZE::<K>]>::new_uninit()
1102                .assume_init()
1103        };
1104
1105        for ([left_bucket, right_bucket], left_bucket_index) in
1106            parent_table.buckets().array_windows().zip(0..)
1107        {
1108            let mut matches = [MaybeUninit::uninit(); _];
1109            // SAFETY: Positions are taken from `Table::buckets()` and correspond to initialized
1110            // values
1111            let matches = unsafe {
1112                find_matches_in_buckets(
1113                    left_bucket_index,
1114                    left_bucket,
1115                    right_bucket,
1116                    &mut matches,
1117                    left_targets,
1118                )
1119            };
1120            // Throw away some successful matches that are not that necessary
1121            let matches = &matches[..matches.len().min(REDUCED_MATCHES_COUNT)];
1122            // SAFETY: Already initialized this many elements
1123            let (ys, positions, metadatas) = unsafe {
1124                (
1125                    ys.get_unchecked_mut(initialized_elements..),
1126                    positions.get_unchecked_mut(initialized_elements..),
1127                    metadatas.get_unchecked_mut(initialized_elements..),
1128                )
1129            };
1130
1131            // SAFETY: Preallocated length is an upper bound and is always sufficient
1132            let (ys, positions, metadatas) = unsafe {
1133                (
1134                    ys.get_unchecked_mut(..matches.len()),
1135                    positions.get_unchecked_mut(..matches.len()),
1136                    metadatas.get_unchecked_mut(..matches.len()),
1137                )
1138            };
1139
1140            // SAFETY: Matches come from the parent table and the size of `ys`, `positions`
1141            // and `metadatas` is the same as the number of matches
1142            unsafe {
1143                matches_to_results(&parent_table, matches, ys, positions, metadatas);
1144            }
1145
1146            initialized_elements += matches.len();
1147        }
1148
1149        let parent_table = parent_table.prune();
1150
1151        // SAFETY: Converting a boxed array to a vector of the same size, which has the same memory
1152        // layout, the number of elements matches the number of elements that were initialized
1153        let ys = unsafe {
1154            let ys_len = ys.len();
1155            let ys = Box::into_raw(ys);
1156            Vec::from_raw_parts(ys.cast(), initialized_elements, ys_len)
1157        };
1158
1159        // TODO: Try to group buckets in the process of collecting `y`s
1160        let buckets = group_by_buckets::<K>(&ys);
1161
1162        let table = Self::Other {
1163            positions,
1164            metadatas,
1165            buckets,
1166        };
1167
1168        (table, parent_table)
1169    }
1170
1171    /// Almost the same as [`Self::create()`], but uses parallelism internally for better
1172    /// performance (though not efficiency of CPU and memory usage), if you create multiple tables
1173    /// in parallel, prefer this method for better overall performance.
1174    #[cfg(feature = "parallel")]
1175    pub(super) fn create_parallel<const PARENT_TABLE_NUMBER: u8>(
1176        parent_table: Table<K, PARENT_TABLE_NUMBER>,
1177        cache: &TablesCache,
1178    ) -> (Self, PrunedTable<K, PARENT_TABLE_NUMBER>)
1179    where
1180        Table<K, PARENT_TABLE_NUMBER>: private::NotLastTable,
1181    {
1182        // SAFETY: Contents is `MaybeUninit`
1183        let ys = unsafe {
1184            Box::<[SyncUnsafeCell<[MaybeUninit<_>; REDUCED_MATCHES_COUNT]>; NUM_BUCKET_PAIRS::<K>]>::new_uninit().assume_init()
1185        };
1186        // SAFETY: Contents is `MaybeUninit`
1187        let positions = unsafe {
1188            Box::<[SyncUnsafeCell<[MaybeUninit<_>; REDUCED_MATCHES_COUNT]>; NUM_BUCKET_PAIRS::<K>]>::new_uninit().assume_init()
1189        };
1190        // SAFETY: Contents is `MaybeUninit`
1191        let metadatas = unsafe {
1192            Box::<[SyncUnsafeCell<[MaybeUninit<_>; REDUCED_MATCHES_COUNT]>; NUM_BUCKET_PAIRS::<K>]>::new_uninit().assume_init()
1193        };
1194        let global_results_counts =
1195            array::from_fn::<_, { NUM_BUCKET_PAIRS::<K> }, _>(|_| SyncUnsafeCell::new(0u16));
1196
1197        let left_targets = &*cache.left_targets;
1198
1199        let buckets = parent_table.buckets();
1200        // Iterate over buckets in batches, such that a cache line worth of bytes is taken from
1201        // `global_results_counts` each time to avoid unnecessary false sharing
1202        let bucket_batch_size = CACHE_LINE_SIZE / size_of::<u16>();
1203        let bucket_batch_index = AtomicUsize::new(0);
1204
1205        rayon::broadcast(|_ctx| {
1206            loop {
1207                let bucket_batch_index = bucket_batch_index.fetch_add(1, Ordering::Relaxed);
1208
1209                let buckets_batch = buckets
1210                    .array_windows()
1211                    .enumerate()
1212                    .skip(bucket_batch_index * bucket_batch_size)
1213                    .take(bucket_batch_size);
1214
1215                if buckets_batch.is_empty() {
1216                    break;
1217                }
1218
1219                for (left_bucket_index, [left_bucket, right_bucket]) in buckets_batch {
1220                    let mut matches = [MaybeUninit::uninit(); _];
1221                    // SAFETY: Positions are taken from `Table::buckets()` and correspond to
1222                    // initialized values
1223                    let matches = unsafe {
1224                        find_matches_in_buckets(
1225                            left_bucket_index as u32,
1226                            left_bucket,
1227                            right_bucket,
1228                            &mut matches,
1229                            left_targets,
1230                        )
1231                    };
1232                    // Throw away some successful matches that are not that necessary
1233                    let matches = &matches[..matches.len().min(REDUCED_MATCHES_COUNT)];
1234
1235                    // SAFETY: This is the only place where `left_bucket_index`'s entry is accessed
1236                    // at this time, and it is guaranteed to be in range
1237                    let ys = unsafe { &mut *ys.get_unchecked(left_bucket_index).get() };
1238                    // SAFETY: This is the only place where `left_bucket_index`'s entry is accessed
1239                    // at this time, and it is guaranteed to be in range
1240                    let positions =
1241                        unsafe { &mut *positions.get_unchecked(left_bucket_index).get() };
1242                    // SAFETY: This is the only place where `left_bucket_index`'s entry is accessed
1243                    // at this time, and it is guaranteed to be in range
1244                    let metadatas =
1245                        unsafe { &mut *metadatas.get_unchecked(left_bucket_index).get() };
1246                    // SAFETY: This is the only place where `left_bucket_index`'s entry is accessed
1247                    // at this time, and it is guaranteed to be in range
1248                    let count = unsafe {
1249                        &mut *global_results_counts.get_unchecked(left_bucket_index).get()
1250                    };
1251
1252                    // SAFETY: Matches come from the parent table and the size of `ys`, `positions`
1253                    // and `metadatas` is larger or equal to the number of matches
1254                    unsafe {
1255                        matches_to_results::<_, TABLE_NUMBER, _>(
1256                            &parent_table,
1257                            matches,
1258                            ys,
1259                            positions,
1260                            metadatas,
1261                        );
1262                    }
1263                    *count = matches.len() as u16;
1264                }
1265            }
1266        });
1267
1268        let parent_table = parent_table.prune();
1269
1270        let ys = strip_sync_unsafe_cell(ys);
1271        let positions = strip_sync_unsafe_cell(positions);
1272        let metadatas = strip_sync_unsafe_cell(metadatas);
1273
1274        // TODO: Try to group buckets in the process of collecting `y`s
1275        // SAFETY: `global_results_counts` corresponds to the number of initialized `ys`
1276        let buckets = unsafe {
1277            group_by_buckets_from_buckets::<K, _>(
1278                ys.iter().zip(
1279                    global_results_counts
1280                        .into_iter()
1281                        .map(|count| usize::from(count.into_inner())),
1282                ),
1283            )
1284        };
1285
1286        let table = Self::OtherBuckets {
1287            positions,
1288            metadatas,
1289            buckets,
1290        };
1291
1292        (table, parent_table)
1293    }
1294
1295    /// Get `[left_position, right_position]` of a previous table for a specified position in a
1296    /// current table.
1297    ///
1298    /// # Safety
1299    /// `position` must come from [`Self::buckets()`] or [`Self::position()`] or
1300    /// [`PrunedTable::position()`] and not be a sentinel value.
1301    #[inline(always)]
1302    pub(super) unsafe fn position(&self, position: Position) -> [Position; 2] {
1303        match self {
1304            Self::First { .. } => {
1305                unreachable!("Not the first table");
1306            }
1307            Self::Other { positions, .. } => {
1308                // SAFETY: All non-sentinel positions returned by [`Self::buckets()`] are valid
1309                unsafe { positions.get_unchecked(usize::from(position)).assume_init() }
1310            }
1311            #[cfg(feature = "parallel")]
1312            Self::OtherBuckets { positions, .. } => {
1313                // SAFETY: All non-sentinel positions returned by [`Self::buckets()`] are valid
1314                unsafe {
1315                    positions
1316                        .as_flattened()
1317                        .get_unchecked(usize::from(position))
1318                        .assume_init()
1319                }
1320            }
1321        }
1322    }
1323}
1324
1325#[cfg(feature = "alloc")]
1326impl<const K: u8> Table<K, 7>
1327where
1328    Self: private::SupportedOtherTables,
1329{
1330    /// Proof targets from the last table into the previous table, one for each
1331    /// [`Record::NUM_S_BUCKETS`].
1332    pub(super) fn create_proof_targets(
1333        parent_table: Table<K, 6>,
1334        cache: &TablesCache,
1335    ) -> (
1336        Box<[[Position; 2]; const { Record::NUM_S_BUCKETS }]>,
1337        PrunedTable<K, 6>,
1338    )
1339    where
1340        Table<K, 6>: private::NotLastTable,
1341    {
1342        let left_targets = &*cache.left_targets;
1343        // SAFETY: Data structure filled with zeroes is a valid invariant
1344        let mut table_6_proof_targets = unsafe {
1345            Box::<[[Position; 2]; const { Record::NUM_S_BUCKETS }]>::new_zeroed().assume_init()
1346        };
1347
1348        for ([left_bucket, right_bucket], left_bucket_index) in
1349            parent_table.buckets().array_windows().zip(0..)
1350        {
1351            let mut matches = [MaybeUninit::uninit(); _];
1352            // SAFETY: Positions are taken from `Table::buckets()` and correspond to initialized
1353            // values
1354            let matches = unsafe {
1355                find_matches_in_buckets(
1356                    left_bucket_index,
1357                    left_bucket,
1358                    right_bucket,
1359                    &mut matches,
1360                    left_targets,
1361                )
1362            };
1363            // Throw away some successful matches that are not that necessary
1364            let matches = &matches[..matches.len().min(REDUCED_MATCHES_COUNT)];
1365
1366            let (grouped_matches, other_matches) = matches.as_chunks::<COMPUTE_FN_SIMD_FACTOR>();
1367
1368            for grouped_matches in grouped_matches {
1369                // SAFETY: Guaranteed by function contract
1370                let (ys_group, positions_group, _) =
1371                    unsafe { match_to_result_simd::<_, 7, _>(&parent_table, grouped_matches) };
1372
1373                let s_buckets = ys_group >> Simd::splat(u32::from(PARAM_EXT));
1374
1375                for (s_bucket, p) in s_buckets.to_array().into_iter().zip(positions_group) {
1376                    const {
1377                        assert!(Record::NUM_S_BUCKETS == usize::from(u16::MAX) + 1);
1378                    }
1379                    let Ok(s_bucket) = u16::try_from(s_bucket) else {
1380                        continue;
1381                    };
1382                    let positions = &mut table_6_proof_targets[usize::from(s_bucket)];
1383                    if positions == &[Position::ZERO; 2] {
1384                        *positions = p;
1385                    }
1386                }
1387            }
1388            for other_match in other_matches {
1389                // SAFETY: Guaranteed by function contract
1390                let (y, p, _) = unsafe { match_to_result::<_, 7, _>(&parent_table, other_match) };
1391
1392                let s_bucket = y.first_k_bits();
1393
1394                const {
1395                    assert!(Record::NUM_S_BUCKETS == usize::from(u16::MAX) + 1);
1396                }
1397                let Ok(s_bucket) = u16::try_from(s_bucket) else {
1398                    continue;
1399                };
1400
1401                let positions = &mut table_6_proof_targets[usize::from(s_bucket)];
1402                if positions == &[Position::ZERO; 2] {
1403                    *positions = p;
1404                }
1405            }
1406        }
1407
1408        let parent_table = parent_table.prune();
1409
1410        (table_6_proof_targets, parent_table)
1411    }
1412
1413    /// Almost the same as [`Self::create_proof_targets()`], but uses parallelism internally for
1414    /// better performance (though not efficiency of CPU and memory usage), if you create multiple
1415    /// tables in parallel, prefer this method for better overall performance.
1416    #[cfg(feature = "parallel")]
1417    pub(super) fn create_proof_targets_parallel(
1418        parent_table: Table<K, 6>,
1419        cache: &TablesCache,
1420    ) -> (
1421        Box<[[Position; 2]; const { Record::NUM_S_BUCKETS }]>,
1422        PrunedTable<K, 6>,
1423    )
1424    where
1425        Table<K, 6>: private::NotLastTable,
1426    {
1427        // SAFETY: Contents is `MaybeUninit`
1428        let buckets_positions = unsafe {
1429            Box::<[SyncUnsafeCell<[MaybeUninit<_>; REDUCED_MATCHES_COUNT]>; NUM_BUCKET_PAIRS::<K>]>::new_uninit().assume_init()
1430        };
1431        let global_results_counts =
1432            array::from_fn::<_, { NUM_BUCKET_PAIRS::<K> }, _>(|_| SyncUnsafeCell::new(0u16));
1433
1434        let left_targets = &*cache.left_targets;
1435
1436        let buckets = parent_table.buckets();
1437        // Iterate over buckets in batches, such that a cache line worth of bytes is taken from
1438        // `global_results_counts` each time to avoid unnecessary false sharing
1439        let bucket_batch_size = CACHE_LINE_SIZE / size_of::<u16>();
1440        let bucket_batch_index = AtomicUsize::new(0);
1441
1442        rayon::broadcast(|_ctx| {
1443            loop {
1444                let bucket_batch_index = bucket_batch_index.fetch_add(1, Ordering::Relaxed);
1445
1446                let buckets_batch = buckets
1447                    .array_windows()
1448                    .enumerate()
1449                    .skip(bucket_batch_index * bucket_batch_size)
1450                    .take(bucket_batch_size);
1451
1452                if buckets_batch.is_empty() {
1453                    break;
1454                }
1455
1456                for (left_bucket_index, [left_bucket, right_bucket]) in buckets_batch {
1457                    let mut matches = [MaybeUninit::uninit(); _];
1458                    // SAFETY: Positions are taken from `Table::buckets()` and correspond to
1459                    // initialized values
1460                    let matches = unsafe {
1461                        find_matches_in_buckets(
1462                            left_bucket_index as u32,
1463                            left_bucket,
1464                            right_bucket,
1465                            &mut matches,
1466                            left_targets,
1467                        )
1468                    };
1469                    // Throw away some successful matches that are not that necessary
1470                    let matches = &matches[..matches.len().min(REDUCED_MATCHES_COUNT)];
1471
1472                    // SAFETY: This is the only place where `left_bucket_index`'s entry is accessed
1473                    // at this time, and it is guaranteed to be in range
1474                    let buckets_positions =
1475                        unsafe { &mut *buckets_positions.get_unchecked(left_bucket_index).get() };
1476                    // SAFETY: This is the only place where `left_bucket_index`'s entry is accessed
1477                    // at this time, and it is guaranteed to be in range
1478                    let count = unsafe {
1479                        &mut *global_results_counts.get_unchecked(left_bucket_index).get()
1480                    };
1481
1482                    let (grouped_matches, other_matches) =
1483                        matches.as_chunks::<COMPUTE_FN_SIMD_FACTOR>();
1484
1485                    let mut reduced_count = 0_usize;
1486                    for grouped_matches in grouped_matches {
1487                        // SAFETY: Guaranteed by function contract
1488                        let (ys_group, positions_group, _) = unsafe {
1489                            match_to_result_simd::<_, 7, _>(&parent_table, grouped_matches)
1490                        };
1491
1492                        let s_buckets = ys_group >> Simd::splat(u32::from(PARAM_EXT));
1493                        let s_buckets = s_buckets.to_array();
1494
1495                        for (s_bucket, p) in s_buckets.into_iter().zip(positions_group) {
1496                            const {
1497                                assert!(Record::NUM_S_BUCKETS == usize::from(u16::MAX) + 1);
1498                            }
1499                            let Ok(s_bucket) = u16::try_from(s_bucket) else {
1500                                continue;
1501                            };
1502
1503                            buckets_positions[reduced_count].write((s_bucket, p));
1504                            reduced_count += 1;
1505                        }
1506                    }
1507                    for other_match in other_matches {
1508                        // SAFETY: Guaranteed by function contract
1509                        let (y, p, _) =
1510                            unsafe { match_to_result::<_, 7, _>(&parent_table, other_match) };
1511
1512                        let s_bucket = y.first_k_bits();
1513
1514                        const {
1515                            assert!(Record::NUM_S_BUCKETS == usize::from(u16::MAX) + 1);
1516                        }
1517                        let Ok(s_bucket) = u16::try_from(s_bucket) else {
1518                            continue;
1519                        };
1520
1521                        buckets_positions[reduced_count].write((s_bucket, p));
1522                        reduced_count += 1;
1523                    }
1524
1525                    *count = reduced_count as u16;
1526                }
1527            }
1528        });
1529
1530        let parent_table = parent_table.prune();
1531
1532        let buckets_positions = strip_sync_unsafe_cell(buckets_positions);
1533
1534        // SAFETY: Data structure filled with zeroes is a valid invariant
1535        let mut table_6_proof_targets = unsafe {
1536            Box::<[[Position; 2]; const { Record::NUM_S_BUCKETS }]>::new_zeroed().assume_init()
1537        };
1538
1539        for (bucket, results_count) in buckets_positions.iter().zip(
1540            global_results_counts
1541                .into_iter()
1542                .map(|count| usize::from(count.into_inner())),
1543        ) {
1544            // SAFETY: `results_count` corresponds to the number of initialized `bucket` elements
1545            for &(s_bucket, p) in unsafe { bucket[..results_count].assume_init_ref() } {
1546                let positions = &mut table_6_proof_targets[usize::from(s_bucket)];
1547                if positions == &[Position::ZERO; 2] {
1548                    *positions = p;
1549                }
1550            }
1551        }
1552
1553        (table_6_proof_targets, parent_table)
1554    }
1555}
1556
1557#[cfg(feature = "alloc")]
1558impl<const K: u8, const TABLE_NUMBER: u8> Table<K, TABLE_NUMBER>
1559where
1560    Self: private::NotLastTable,
1561{
1562    /// Returns `None` for an invalid position or for table number 7.
1563    ///
1564    /// # Safety
1565    /// `position` must come from [`Self::buckets()`] and not be a sentinel value.
1566    #[inline(always)]
1567    unsafe fn metadata(&self, position: Position) -> Metadata<K, TABLE_NUMBER> {
1568        match self {
1569            Self::First { .. } => {
1570                // X matches position
1571                Metadata::from(X::from(u32::from(position)))
1572            }
1573            Self::Other { metadatas, .. } => {
1574                // SAFETY: All non-sentinel positions returned by [`Self::buckets()`] are valid
1575                unsafe { metadatas.get_unchecked(usize::from(position)).assume_init() }
1576            }
1577            #[cfg(feature = "parallel")]
1578            Self::OtherBuckets { metadatas, .. } => {
1579                // SAFETY: All non-sentinel positions returned by [`Self::buckets()`] are valid
1580                unsafe {
1581                    metadatas
1582                        .as_flattened()
1583                        .get_unchecked(usize::from(position))
1584                        .assume_init()
1585                }
1586            }
1587        }
1588    }
1589}
1590
1591#[cfg(feature = "alloc")]
1592impl<const K: u8, const TABLE_NUMBER: u8> Table<K, TABLE_NUMBER> {
1593    #[inline(always)]
1594    fn prune(self) -> PrunedTable<K, TABLE_NUMBER> {
1595        match self {
1596            Self::First { .. } => PrunedTable::First,
1597            Self::Other { positions, .. } => PrunedTable::Other { positions },
1598            #[cfg(feature = "parallel")]
1599            Self::OtherBuckets { positions, .. } => PrunedTable::OtherBuckets { positions },
1600        }
1601    }
1602
1603    /// Positions of `y`s grouped by the bucket they belong to
1604    #[inline(always)]
1605    pub(super) fn buckets(&self) -> &[[(Position, Y); REDUCED_BUCKET_SIZE]; NUM_BUCKETS::<K>] {
1606        match self {
1607            Self::First { buckets, .. } => buckets,
1608            Self::Other { buckets, .. } => buckets,
1609            #[cfg(feature = "parallel")]
1610            Self::OtherBuckets { buckets, .. } => buckets,
1611        }
1612    }
1613}