subspace_proof_of_space/chiapos/
table.rs

1#[cfg(test)]
2mod tests;
3pub(super) mod types;
4
5#[cfg(not(feature = "std"))]
6extern crate alloc;
7
8use crate::chiapos::Seed;
9use crate::chiapos::constants::{PARAM_B, PARAM_BC, PARAM_C, PARAM_EXT, PARAM_M};
10use crate::chiapos::table::types::{Metadata, Position, X, Y};
11use crate::chiapos::utils::EvaluatableUsize;
12#[cfg(not(feature = "std"))]
13use alloc::vec;
14#[cfg(not(feature = "std"))]
15use alloc::vec::Vec;
16use chacha20::cipher::{KeyIvInit, StreamCipher, StreamCipherSeek};
17use chacha20::{ChaCha8, Key, Nonce};
18use core::array;
19use core::simd::Simd;
20use core::simd::num::SimdUint;
21#[cfg(all(feature = "std", any(feature = "parallel", test)))]
22use parking_lot::Mutex;
23#[cfg(any(feature = "parallel", test))]
24use rayon::prelude::*;
25use seq_macro::seq;
26#[cfg(all(not(feature = "std"), any(feature = "parallel", test)))]
27use spin::Mutex;
28
29pub(super) const COMPUTE_F1_SIMD_FACTOR: usize = 8;
30pub(super) const FIND_MATCHES_AND_COMPUTE_UNROLL_FACTOR: usize = 8;
31
32/// Compute the size of `y` in bits
33pub(super) const fn y_size_bits(k: u8) -> usize {
34    k as usize + PARAM_EXT as usize
35}
36
37/// Metadata size in bytes
38pub const fn metadata_size_bytes(k: u8, table_number: u8) -> usize {
39    metadata_size_bits(k, table_number).div_ceil(u8::BITS as usize)
40}
41
42/// Metadata size in bits
43pub(super) const fn metadata_size_bits(k: u8, table_number: u8) -> usize {
44    k as usize
45        * match table_number {
46            1 => 1,
47            2 => 2,
48            3 | 4 => 4,
49            5 => 3,
50            6 => 2,
51            7 => 0,
52            _ => unreachable!(),
53        }
54}
55
56/// ChaCha8 [`Vec`] sufficient for the whole first table for [`K`].
57/// Prefer [`partial_y`] if you need partial y just for a single `x`.
58fn partial_ys<const K: u8>(seed: Seed) -> Vec<u8> {
59    let output_len_bits = usize::from(K) * (1 << K);
60    let mut output = vec![0; output_len_bits.div_ceil(u8::BITS as usize)];
61
62    let key = Key::from(seed);
63    let nonce = Nonce::default();
64
65    let mut cipher = ChaCha8::new(&key, &nonce);
66
67    cipher.apply_keystream(&mut output);
68
69    output
70}
71
72/// ChaCha8 byte for a single `y` at `x` in the first table for [`K`], returns bytes and offset (in
73/// bits) within those bytes at which data start.
74/// Prefer [`partial_ys`] if you process the whole first table.
75pub(super) fn partial_y<const K: u8>(
76    seed: Seed,
77    x: X,
78) -> ([u8; (K as usize * 2).div_ceil(u8::BITS as usize)], usize) {
79    let skip_bits = usize::from(K) * usize::from(x);
80    let skip_bytes = skip_bits / u8::BITS as usize;
81    let skip_bits = skip_bits % u8::BITS as usize;
82
83    let mut output = [0; (K as usize * 2).div_ceil(u8::BITS as usize)];
84
85    let key = Key::from(seed);
86    let nonce = Nonce::default();
87
88    let mut cipher = ChaCha8::new(&key, &nonce);
89
90    cipher.seek(skip_bytes);
91    cipher.apply_keystream(&mut output);
92
93    (output, skip_bits)
94}
95
96#[derive(Debug, Clone)]
97struct LeftTargets {
98    left_targets: Vec<Position>,
99}
100
101fn calculate_left_targets() -> LeftTargets {
102    let mut left_targets = Vec::with_capacity(2 * usize::from(PARAM_BC) * usize::from(PARAM_M));
103
104    let param_b = u32::from(PARAM_B);
105    let param_c = u32::from(PARAM_C);
106
107    for parity in 0..=1u32 {
108        for r in 0..u32::from(PARAM_BC) {
109            let c = r / param_c;
110
111            for m in 0..u32::from(PARAM_M) {
112                let target = ((c + m) % param_b) * param_c
113                    + (((2 * m + parity) * (2 * m + parity) + r) % param_c);
114                left_targets.push(Position::from(target));
115            }
116        }
117    }
118
119    LeftTargets { left_targets }
120}
121
122fn calculate_left_target_on_demand(parity: u32, r: u32, m: u32) -> u32 {
123    let param_b = u32::from(PARAM_B);
124    let param_c = u32::from(PARAM_C);
125
126    let c = r / param_c;
127
128    ((c + m) % param_b) * param_c + (((2 * m + parity) * (2 * m + parity) + r) % param_c)
129}
130
131/// Caches that can be used to optimize creation of multiple [`Tables`](super::Tables).
132#[derive(Debug, Clone)]
133pub struct TablesCache<const K: u8> {
134    buckets: Vec<Bucket>,
135    rmap_scratch: Vec<RmapItem>,
136    left_targets: LeftTargets,
137}
138
139impl<const K: u8> Default for TablesCache<K> {
140    /// Create new instance
141    fn default() -> Self {
142        Self {
143            buckets: Vec::new(),
144            rmap_scratch: Vec::new(),
145            left_targets: calculate_left_targets(),
146        }
147    }
148}
149
150#[derive(Debug)]
151struct Match {
152    left_position: Position,
153    left_y: Y,
154    right_position: Position,
155}
156
157#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
158struct Bucket {
159    /// Bucket index
160    bucket_index: u32,
161    /// Start position of this bucket in the table
162    start_position: Position,
163    /// Size of this bucket
164    size: Position,
165}
166
167#[derive(Debug, Default, Copy, Clone)]
168pub(super) struct RmapItem {
169    count: Position,
170    start_position: Position,
171}
172
173/// `partial_y_offset` is in bits
174pub(super) fn compute_f1<const K: u8>(x: X, partial_y: &[u8], partial_y_offset: usize) -> Y {
175    let partial_y_length =
176        (partial_y_offset % u8::BITS as usize + usize::from(K)).div_ceil(u8::BITS as usize);
177    let mut pre_y_bytes = 0u64.to_be_bytes();
178    pre_y_bytes[..partial_y_length]
179        .copy_from_slice(&partial_y[partial_y_offset / u8::BITS as usize..][..partial_y_length]);
180    // Contains `K` desired bits of `partial_y` in the final offset of eventual `y` with the rest
181    // of bits being in undefined state
182    let pre_y = u64::from_be_bytes(pre_y_bytes)
183        >> (u64::BITS as usize - usize::from(K + PARAM_EXT) - partial_y_offset % u8::BITS as usize);
184    let pre_y = pre_y as u32;
185    // Mask for clearing the rest of bits of `pre_y`.
186    let pre_y_mask = (u32::MAX << usize::from(PARAM_EXT))
187        & (u32::MAX >> (u32::BITS as usize - usize::from(K + PARAM_EXT)));
188
189    // Extract `PARAM_EXT` most significant bits from `x` and store in the final offset of
190    // eventual `y` with the rest of bits being in undefined state.
191    let pre_ext = u32::from(x) >> (usize::from(K - PARAM_EXT));
192    // Mask for clearing the rest of bits of `pre_ext`.
193    let pre_ext_mask = u32::MAX >> (u32::BITS as usize - usize::from(PARAM_EXT));
194
195    // Combine all of the bits together:
196    // [padding zero bits][`K` bits rom `partial_y`][`PARAM_EXT` bits from `x`]
197    Y::from((pre_y & pre_y_mask) | (pre_ext & pre_ext_mask))
198}
199
200pub(super) fn compute_f1_simd<const K: u8>(
201    xs: [u32; COMPUTE_F1_SIMD_FACTOR],
202    partial_ys: &[u8; K as usize * COMPUTE_F1_SIMD_FACTOR / u8::BITS as usize],
203) -> [Y; COMPUTE_F1_SIMD_FACTOR] {
204    // Each element contains `K` desired bits of `partial_ys` in the final offset of eventual `ys`
205    // with the rest of bits being in undefined state
206    let pre_ys_bytes = array::from_fn(|i| {
207        let partial_y_offset = i * usize::from(K);
208        let partial_y_length =
209            (partial_y_offset % u8::BITS as usize + usize::from(K)).div_ceil(u8::BITS as usize);
210        let mut pre_y_bytes = 0u64.to_be_bytes();
211        pre_y_bytes[..partial_y_length].copy_from_slice(
212            &partial_ys[partial_y_offset / u8::BITS as usize..][..partial_y_length],
213        );
214
215        u64::from_be_bytes(pre_y_bytes)
216    });
217    let pre_ys_right_offset = array::from_fn(|i| {
218        let partial_y_offset = i as u32 * u32::from(K);
219        u64::from(u64::BITS - u32::from(K + PARAM_EXT) - partial_y_offset % u8::BITS)
220    });
221    let pre_ys = Simd::from_array(pre_ys_bytes) >> Simd::from_array(pre_ys_right_offset);
222
223    // Mask for clearing the rest of bits of `pre_ys`.
224    let pre_ys_mask = Simd::splat(
225        (u32::MAX << usize::from(PARAM_EXT))
226            & (u32::MAX >> (u32::BITS as usize - usize::from(K + PARAM_EXT))),
227    );
228
229    // Extract `PARAM_EXT` most significant bits from `xs` and store in the final offset of
230    // eventual `ys` with the rest of bits being in undefined state.
231    let pre_exts = Simd::from_array(xs) >> Simd::splat(u32::from(K - PARAM_EXT));
232
233    // Combine all of the bits together:
234    // [padding zero bits][`K` bits rom `partial_y`][`PARAM_EXT` bits from `x`]
235    let ys = (pre_ys.cast() & pre_ys_mask) | pre_exts;
236
237    Y::array_from_repr(ys.to_array())
238}
239
240/// `rmap_scratch` is just an optimization to reuse allocations between calls.
241///
242/// For verification purposes use [`has_match`] instead.
243///
244/// Returns `None` if either of buckets is empty.
245#[allow(clippy::too_many_arguments)]
246fn find_matches<T, Map>(
247    left_bucket_ys: &[Y],
248    left_bucket_start_position: Position,
249    right_bucket_ys: &[Y],
250    right_bucket_start_position: Position,
251    rmap_scratch: &mut Vec<RmapItem>,
252    left_targets: &LeftTargets,
253    map: Map,
254    output: &mut Vec<T>,
255) where
256    Map: Fn(Match) -> T,
257{
258    // Clear and set to correct size with zero values
259    rmap_scratch.clear();
260    rmap_scratch.resize_with(usize::from(PARAM_BC), RmapItem::default);
261    let rmap = rmap_scratch;
262
263    // Both left and right buckets can be empty
264    let Some(&first_left_bucket_y) = left_bucket_ys.first() else {
265        return;
266    };
267    let Some(&first_right_bucket_y) = right_bucket_ys.first() else {
268        return;
269    };
270    // Since all entries in a bucket are obtained after division by `PARAM_BC`, we can compute
271    // quotient more efficiently by subtracting base value rather than computing remainder of
272    // division
273    let base = (usize::from(first_right_bucket_y) / usize::from(PARAM_BC)) * usize::from(PARAM_BC);
274    for (&y, right_position) in right_bucket_ys.iter().zip(right_bucket_start_position..) {
275        let r = usize::from(y) - base;
276
277        // Same `y` and as the result `r` can appear in the table multiple times, in which case
278        // they'll all occupy consecutive slots in `right_bucket` and all we need to store is just
279        // the first position and number of elements.
280        if rmap[r].count == Position::ZERO {
281            rmap[r].start_position = right_position;
282        }
283        rmap[r].count += Position::ONE;
284    }
285    let rmap = rmap.as_slice();
286
287    // Same idea as above, but avoids division by leveraging the fact that each bucket is exactly
288    // `PARAM_BC` away from the previous one in terms of divisor by `PARAM_BC`
289    let base = base - usize::from(PARAM_BC);
290    let parity = (usize::from(first_left_bucket_y) / usize::from(PARAM_BC)) % 2;
291    let left_targets_parity = {
292        let (a, b) = left_targets
293            .left_targets
294            .split_at(left_targets.left_targets.len() / 2);
295        if parity == 0 { a } else { b }
296    };
297
298    for (&y, left_position) in left_bucket_ys.iter().zip(left_bucket_start_position..) {
299        let r = usize::from(y) - base;
300        let left_targets_r = left_targets_parity
301            .chunks_exact(left_targets_parity.len() / usize::from(PARAM_BC))
302            .nth(r)
303            .expect("r is valid");
304
305        const _: () = {
306            assert!((PARAM_M as usize).is_multiple_of(FIND_MATCHES_AND_COMPUTE_UNROLL_FACTOR));
307        };
308
309        for r_targets in left_targets_r
310            .as_chunks::<{ FIND_MATCHES_AND_COMPUTE_UNROLL_FACTOR }>()
311            .0
312            .iter()
313            .take(usize::from(PARAM_M) / FIND_MATCHES_AND_COMPUTE_UNROLL_FACTOR)
314        {
315            let _: [(); FIND_MATCHES_AND_COMPUTE_UNROLL_FACTOR] = seq!(N in 0..8 {
316                [
317                #(
318                {
319                    let rmap_item = rmap[usize::from(r_targets[N])];
320
321                    for right_position in
322                        rmap_item.start_position..rmap_item.start_position + rmap_item.count
323                    {
324                        let m = Match {
325                            left_position,
326                            left_y: y,
327                            right_position,
328                        };
329                        output.push(map(m));
330                    }
331                },
332                )*
333                ]
334            });
335        }
336    }
337}
338
339/// Simplified version of [`find_matches`] for verification purposes.
340pub(super) fn has_match(left_y: Y, right_y: Y) -> bool {
341    let right_r = u32::from(right_y) % u32::from(PARAM_BC);
342    let parity = (u32::from(left_y) / u32::from(PARAM_BC)) % 2;
343    let left_r = u32::from(left_y) % u32::from(PARAM_BC);
344
345    let r_targets = array::from_fn::<_, { PARAM_M as usize }, _>(|i| {
346        calculate_left_target_on_demand(parity, left_r, i as u32)
347    });
348
349    r_targets.contains(&right_r)
350}
351
352pub(super) fn compute_fn<const K: u8, const TABLE_NUMBER: u8, const PARENT_TABLE_NUMBER: u8>(
353    y: Y,
354    left_metadata: Metadata<K, PARENT_TABLE_NUMBER>,
355    right_metadata: Metadata<K, PARENT_TABLE_NUMBER>,
356) -> (Y, Metadata<K, TABLE_NUMBER>)
357where
358    EvaluatableUsize<{ metadata_size_bytes(K, PARENT_TABLE_NUMBER) }>: Sized,
359    EvaluatableUsize<{ metadata_size_bytes(K, TABLE_NUMBER) }>: Sized,
360{
361    let left_metadata = u128::from(left_metadata);
362    let right_metadata = u128::from(right_metadata);
363
364    let parent_metadata_bits = metadata_size_bits(K, PARENT_TABLE_NUMBER);
365
366    // Only supports `K` from 15 to 25 (otherwise math will not be correct when concatenating y,
367    // left metadata and right metadata)
368    let hash = {
369        // Take only bytes where bits were set
370        let num_bytes_with_data = (y_size_bits(K) + metadata_size_bits(K, PARENT_TABLE_NUMBER) * 2)
371            .div_ceil(u8::BITS as usize);
372
373        // Collect `K` most significant bits of `y` at the final offset of eventual `input_a`
374        let y_bits = u128::from(y) << (u128::BITS as usize - y_size_bits(K));
375
376        // Move bits of `left_metadata` at the final offset of eventual `input_a`
377        let left_metadata_bits =
378            left_metadata << (u128::BITS as usize - parent_metadata_bits - y_size_bits(K));
379
380        // Part of the `right_bits` at the final offset of eventual `input_a`
381        let y_and_left_bits = y_size_bits(K) + parent_metadata_bits;
382        let right_bits_start_offset = u128::BITS as usize - parent_metadata_bits;
383
384        // If `right_metadata` bits start to the left of the desired position in `input_a` move
385        // bits right, else move left
386        if right_bits_start_offset < y_and_left_bits {
387            let right_bits_pushed_into_input_b = y_and_left_bits - right_bits_start_offset;
388            // Collect bits of `right_metadata` that will fit into `input_a` at the final offset in
389            // eventual `input_a`
390            let right_bits_a = right_metadata >> right_bits_pushed_into_input_b;
391            let input_a = y_bits | left_metadata_bits | right_bits_a;
392            // Collect bits of `right_metadata` that will spill over into `input_b`
393            let input_b = right_metadata << (u128::BITS as usize - right_bits_pushed_into_input_b);
394
395            let input = [input_a.to_be_bytes(), input_b.to_be_bytes()];
396            let input_len =
397                size_of::<u128>() + right_bits_pushed_into_input_b.div_ceil(u8::BITS as usize);
398            blake3::hash(&input.as_flattened()[..input_len])
399        } else {
400            let right_bits_a = right_metadata << (right_bits_start_offset - y_and_left_bits);
401            let input_a = y_bits | left_metadata_bits | right_bits_a;
402
403            blake3::hash(&input_a.to_be_bytes()[..num_bytes_with_data])
404        }
405    };
406    let hash = <[u8; 32]>::from(hash);
407
408    let y_output = Y::from(
409        u32::from_be_bytes([hash[0], hash[1], hash[2], hash[3]])
410            >> (u32::BITS as usize - y_size_bits(K)),
411    );
412
413    let metadata_size_bits = metadata_size_bits(K, TABLE_NUMBER);
414
415    let metadata = if TABLE_NUMBER < 4 {
416        (left_metadata << parent_metadata_bits) | right_metadata
417    } else if metadata_size_bits > 0 {
418        // For K under 25 it is guaranteed that metadata + bit offset will always fit into u128.
419        // We collect bytes necessary, potentially with extra bits at the start and end of the bytes
420        // that will be taken care of later.
421        let metadata = u128::from_be_bytes(
422            hash[y_size_bits(K) / u8::BITS as usize..][..size_of::<u128>()]
423                .try_into()
424                .expect("Always enough bits for any K; qed"),
425        );
426        // Remove extra bits at the beginning
427        let metadata = metadata << (y_size_bits(K) % u8::BITS as usize);
428        // Move bits into correct location
429        metadata >> (u128::BITS as usize - metadata_size_bits)
430    } else {
431        0
432    };
433
434    (y_output, Metadata::from(metadata))
435}
436
437fn match_to_result<const K: u8, const TABLE_NUMBER: u8, const PARENT_TABLE_NUMBER: u8>(
438    last_table: &Table<K, PARENT_TABLE_NUMBER>,
439    m: Match,
440) -> (Y, [Position; 2], Metadata<K, TABLE_NUMBER>)
441where
442    EvaluatableUsize<{ metadata_size_bytes(K, PARENT_TABLE_NUMBER) }>: Sized,
443    EvaluatableUsize<{ metadata_size_bytes(K, TABLE_NUMBER) }>: Sized,
444{
445    let left_metadata = last_table
446        .metadata(m.left_position)
447        .expect("Position resulted from matching is correct; qed");
448    let right_metadata = last_table
449        .metadata(m.right_position)
450        .expect("Position resulted from matching is correct; qed");
451
452    let (y, metadata) =
453        compute_fn::<K, TABLE_NUMBER, PARENT_TABLE_NUMBER>(m.left_y, left_metadata, right_metadata);
454
455    (y, [m.left_position, m.right_position], metadata)
456}
457
458fn match_and_compute_fn<'a, const K: u8, const TABLE_NUMBER: u8, const PARENT_TABLE_NUMBER: u8>(
459    last_table: &'a Table<K, PARENT_TABLE_NUMBER>,
460    left_bucket: Bucket,
461    right_bucket: Bucket,
462    rmap_scratch: &'a mut Vec<RmapItem>,
463    left_targets: &'a LeftTargets,
464    results_table: &mut Vec<(Y, [Position; 2], Metadata<K, TABLE_NUMBER>)>,
465) where
466    EvaluatableUsize<{ metadata_size_bytes(K, PARENT_TABLE_NUMBER) }>: Sized,
467    EvaluatableUsize<{ metadata_size_bytes(K, TABLE_NUMBER) }>: Sized,
468{
469    find_matches(
470        &last_table.ys()[usize::from(left_bucket.start_position)..]
471            [..usize::from(left_bucket.size)],
472        left_bucket.start_position,
473        &last_table.ys()[usize::from(right_bucket.start_position)..]
474            [..usize::from(right_bucket.size)],
475        right_bucket.start_position,
476        rmap_scratch,
477        left_targets,
478        |m| match_to_result(last_table, m),
479        results_table,
480    )
481}
482
483#[derive(Debug)]
484pub(super) enum Table<const K: u8, const TABLE_NUMBER: u8>
485where
486    EvaluatableUsize<{ metadata_size_bytes(K, TABLE_NUMBER) }>: Sized,
487{
488    /// First table with contents of entries split into separate vectors for more efficient access
489    First {
490        /// Derived values computed from `x`
491        ys: Vec<Y>,
492        /// X values
493        xs: Vec<X>,
494    },
495    /// Other tables
496    Other {
497        /// Derived values computed from previous table
498        ys: Vec<Y>,
499        /// Left and right entry positions in a previous table encoded into bits
500        positions: Vec<[Position; 2]>,
501        /// Metadata corresponding to each entry
502        metadatas: Vec<Metadata<K, TABLE_NUMBER>>,
503    },
504}
505
506impl<const K: u8> Table<K, 1>
507where
508    EvaluatableUsize<{ metadata_size_bytes(K, 1) }>: Sized,
509{
510    /// Create the table
511    pub(super) fn create(seed: Seed) -> Self
512    where
513        EvaluatableUsize<{ K as usize * COMPUTE_F1_SIMD_FACTOR / u8::BITS as usize }>: Sized,
514    {
515        let partial_ys = partial_ys::<K>(seed);
516
517        let mut t_1 = Vec::with_capacity(1_usize << K);
518        for (x_batch, partial_ys) in partial_ys
519            .as_chunks::<{ K as usize * COMPUTE_F1_SIMD_FACTOR / u8::BITS as usize }>()
520            .0
521            .iter()
522            .copied()
523            .enumerate()
524        {
525            let xs = array::from_fn::<_, COMPUTE_F1_SIMD_FACTOR, _>(|i| {
526                (x_batch * COMPUTE_F1_SIMD_FACTOR + i) as u32
527            });
528            let ys = compute_f1_simd::<K>(xs, &partial_ys);
529            t_1.extend(ys.into_iter().zip(X::array_from_repr(xs)));
530        }
531
532        t_1.sort_unstable();
533
534        let (ys, xs) = t_1.into_iter().unzip();
535
536        Self::First { ys, xs }
537    }
538
539    /// Create the table, leverages available parallelism
540    #[cfg(any(feature = "parallel", test))]
541    pub(super) fn create_parallel(seed: Seed) -> Self
542    where
543        EvaluatableUsize<{ K as usize * COMPUTE_F1_SIMD_FACTOR / u8::BITS as usize }>: Sized,
544    {
545        let partial_ys = partial_ys::<K>(seed);
546
547        let mut t_1 = Vec::with_capacity(1_usize << K);
548        for (x_batch, partial_ys) in partial_ys
549            .as_chunks::<{ K as usize * COMPUTE_F1_SIMD_FACTOR / u8::BITS as usize }>()
550            .0
551            .iter()
552            .copied()
553            .enumerate()
554        {
555            let xs = array::from_fn::<_, COMPUTE_F1_SIMD_FACTOR, _>(|i| {
556                (x_batch * COMPUTE_F1_SIMD_FACTOR + i) as u32
557            });
558            let ys = compute_f1_simd::<K>(xs, &partial_ys);
559            t_1.extend(ys.into_iter().zip(X::array_from_repr(xs)));
560        }
561
562        t_1.par_sort_unstable();
563
564        let (ys, xs) = t_1.into_iter().unzip();
565
566        Self::First { ys, xs }
567    }
568
569    /// All `x`s as [`BitSlice`], for individual `x`s needs to be slices into [`K`] bits slices
570    pub(super) fn xs(&self) -> &[X] {
571        match self {
572            Table::First { xs, .. } => xs,
573            _ => {
574                unreachable!()
575            }
576        }
577    }
578}
579
580mod private {
581    pub(in super::super) trait SupportedOtherTables {}
582}
583
584impl<const K: u8> private::SupportedOtherTables for Table<K, 2> where
585    EvaluatableUsize<{ metadata_size_bytes(K, 2) }>: Sized
586{
587}
588
589impl<const K: u8> private::SupportedOtherTables for Table<K, 3> where
590    EvaluatableUsize<{ metadata_size_bytes(K, 3) }>: Sized
591{
592}
593
594impl<const K: u8> private::SupportedOtherTables for Table<K, 4> where
595    EvaluatableUsize<{ metadata_size_bytes(K, 4) }>: Sized
596{
597}
598
599impl<const K: u8> private::SupportedOtherTables for Table<K, 5> where
600    EvaluatableUsize<{ metadata_size_bytes(K, 5) }>: Sized
601{
602}
603
604impl<const K: u8> private::SupportedOtherTables for Table<K, 6> where
605    EvaluatableUsize<{ metadata_size_bytes(K, 6) }>: Sized
606{
607}
608
609impl<const K: u8> private::SupportedOtherTables for Table<K, 7> where
610    EvaluatableUsize<{ metadata_size_bytes(K, 7) }>: Sized
611{
612}
613
614impl<const K: u8, const TABLE_NUMBER: u8> Table<K, TABLE_NUMBER>
615where
616    Self: private::SupportedOtherTables,
617    EvaluatableUsize<{ metadata_size_bytes(K, TABLE_NUMBER) }>: Sized,
618{
619    /// Creates new [`TABLE_NUMBER`] table. There also exists [`Self::create_parallel()`] that
620    /// trades CPU efficiency and memory usage for lower latency.
621    pub(super) fn create<const PARENT_TABLE_NUMBER: u8>(
622        last_table: &Table<K, PARENT_TABLE_NUMBER>,
623        cache: &mut TablesCache<K>,
624    ) -> Self
625    where
626        EvaluatableUsize<{ metadata_size_bytes(K, PARENT_TABLE_NUMBER) }>: Sized,
627    {
628        let buckets = &mut cache.buckets;
629        let rmap_scratch = &mut cache.rmap_scratch;
630        let left_targets = &cache.left_targets;
631
632        let mut bucket = Bucket {
633            bucket_index: 0,
634            start_position: Position::ZERO,
635            size: Position::ZERO,
636        };
637
638        let last_y = *last_table
639            .ys()
640            .last()
641            .expect("List of y values is never empty; qed");
642        buckets.clear();
643        buckets.reserve(1 + usize::from(last_y) / usize::from(PARAM_BC));
644        last_table
645            .ys()
646            .iter()
647            .zip(Position::ZERO..)
648            .for_each(|(&y, position)| {
649                let bucket_index = u32::from(y) / u32::from(PARAM_BC);
650
651                if bucket_index == bucket.bucket_index {
652                    bucket.size += Position::ONE;
653                    return;
654                }
655
656                buckets.push(bucket);
657
658                bucket = Bucket {
659                    bucket_index,
660                    start_position: position,
661                    size: Position::ONE,
662                };
663            });
664        // Iteration stopped, but we did not store the last bucket yet
665        buckets.push(bucket);
666
667        let num_values = 1 << K;
668        let mut t_n = Vec::with_capacity(num_values);
669        buckets
670            .array_windows::<2>()
671            .for_each(|&[left_bucket, right_bucket]| {
672                match_and_compute_fn::<K, TABLE_NUMBER, PARENT_TABLE_NUMBER>(
673                    last_table,
674                    left_bucket,
675                    right_bucket,
676                    rmap_scratch,
677                    left_targets,
678                    &mut t_n,
679                );
680            });
681
682        t_n.sort_unstable();
683
684        let mut ys = Vec::with_capacity(t_n.len());
685        let mut positions = Vec::with_capacity(t_n.len());
686        let mut metadatas = Vec::with_capacity(t_n.len());
687
688        for (y, [left_position, right_position], metadata) in t_n {
689            ys.push(y);
690            positions.push([left_position, right_position]);
691            // Last table doesn't have metadata
692            if metadata_size_bits(K, TABLE_NUMBER) > 0 {
693                metadatas.push(metadata);
694            }
695        }
696
697        Self::Other {
698            ys,
699            positions,
700            metadatas,
701        }
702    }
703
704    /// Almost the same as [`Self::create()`], but uses parallelism internally for better
705    /// performance (though not efficiency of CPU and memory usage), if you create multiple tables
706    /// in parallel, prefer [`Self::create()`] for better overall performance.
707    #[cfg(any(feature = "parallel", test))]
708    pub(super) fn create_parallel<const PARENT_TABLE_NUMBER: u8>(
709        last_table: &Table<K, PARENT_TABLE_NUMBER>,
710        cache: &mut TablesCache<K>,
711    ) -> Self
712    where
713        EvaluatableUsize<{ metadata_size_bytes(K, PARENT_TABLE_NUMBER) }>: Sized,
714    {
715        let left_targets = &cache.left_targets;
716
717        let mut first_bucket = Bucket {
718            bucket_index: u32::from(last_table.ys()[0]) / u32::from(PARAM_BC),
719            start_position: Position::ZERO,
720            size: Position::ZERO,
721        };
722        for &y in last_table.ys() {
723            let bucket_index = u32::from(y) / u32::from(PARAM_BC);
724
725            if bucket_index == first_bucket.bucket_index {
726                first_bucket.size += Position::ONE;
727            } else {
728                break;
729            }
730        }
731
732        let previous_bucket = Mutex::new(first_bucket);
733
734        let t_n = rayon::broadcast(|_ctx| {
735            let mut entries = Vec::new();
736            let mut rmap_scratch = Vec::new();
737
738            loop {
739                let left_bucket;
740                let right_bucket;
741                {
742                    let mut previous_bucket = previous_bucket.lock();
743
744                    let right_bucket_start_position =
745                        previous_bucket.start_position + previous_bucket.size;
746                    let right_bucket_index = match last_table
747                        .ys()
748                        .get(usize::from(right_bucket_start_position))
749                    {
750                        Some(&y) => u32::from(y) / u32::from(PARAM_BC),
751                        None => {
752                            break;
753                        }
754                    };
755                    let mut right_bucket_size = Position::ZERO;
756
757                    for &y in &last_table.ys()[usize::from(right_bucket_start_position)..] {
758                        let bucket_index = u32::from(y) / u32::from(PARAM_BC);
759
760                        if bucket_index == right_bucket_index {
761                            right_bucket_size += Position::ONE;
762                        } else {
763                            break;
764                        }
765                    }
766
767                    right_bucket = Bucket {
768                        bucket_index: right_bucket_index,
769                        start_position: right_bucket_start_position,
770                        size: right_bucket_size,
771                    };
772
773                    left_bucket = *previous_bucket;
774                    *previous_bucket = right_bucket;
775                }
776
777                match_and_compute_fn::<K, TABLE_NUMBER, PARENT_TABLE_NUMBER>(
778                    last_table,
779                    left_bucket,
780                    right_bucket,
781                    &mut rmap_scratch,
782                    left_targets,
783                    &mut entries,
784                );
785            }
786
787            entries
788        });
789
790        let mut t_n = t_n.into_iter().flatten().collect::<Vec<_>>();
791        t_n.par_sort_unstable();
792
793        let mut ys = Vec::with_capacity(t_n.len());
794        let mut positions = Vec::with_capacity(t_n.len());
795        let mut metadatas = Vec::with_capacity(t_n.len());
796
797        for (y, [left_position, right_position], metadata) in t_n.drain(..) {
798            ys.push(y);
799            positions.push([left_position, right_position]);
800            // Last table doesn't have metadata
801            if metadata_size_bits(K, TABLE_NUMBER) > 0 {
802                metadatas.push(metadata);
803            }
804        }
805
806        // Drop from a background thread, which typically helps with overall concurrency
807        rayon::spawn(move || {
808            drop(t_n);
809        });
810
811        Self::Other {
812            ys,
813            positions,
814            metadatas,
815        }
816    }
817}
818
819impl<const K: u8, const TABLE_NUMBER: u8> Table<K, TABLE_NUMBER>
820where
821    EvaluatableUsize<{ metadata_size_bytes(K, TABLE_NUMBER) }>: Sized,
822{
823    /// All `y`s as [`BitSlice`], for individual `x`s needs to be slices into [`K`] bits slices
824    pub(super) fn ys(&self) -> &[Y] {
825        let (Table::First { ys, .. } | Table::Other { ys, .. }) = self;
826        ys
827    }
828
829    /// Returns `None` on invalid position or first table, `Some(left_position, right_position)` in
830    /// previous table on success
831    pub(super) fn position(&self, position: Position) -> Option<[Position; 2]> {
832        match self {
833            Table::First { .. } => None,
834            Table::Other { positions, .. } => positions.get(usize::from(position)).copied(),
835        }
836    }
837
838    /// Returns `None` on invalid position or for table number 7
839    pub(super) fn metadata(&self, position: Position) -> Option<Metadata<K, TABLE_NUMBER>> {
840        match self {
841            Table::First { xs, .. } => xs.get(usize::from(position)).map(|&x| Metadata::from(x)),
842            Table::Other { metadatas, .. } => metadatas.get(usize::from(position)).copied(),
843        }
844    }
845}