ab_proof_of_space/
chiapos.rs

1//! Chia proof of space reimplementation in Rust
2
3mod constants;
4mod table;
5#[cfg(all(feature = "alloc", test))]
6mod tests;
7
8#[cfg(feature = "alloc")]
9use crate::PosProofs;
10use crate::chiapos::constants::NUM_TABLES;
11#[cfg(feature = "alloc")]
12use crate::chiapos::table::types::Position;
13use crate::chiapos::table::types::{Metadata, X, Y};
14#[cfg(feature = "alloc")]
15use crate::chiapos::table::{PrunedTable, Table};
16use crate::chiapos::table::{compute_f1, compute_fn, has_match};
17#[cfg(feature = "alloc")]
18use alloc::boxed::Box;
19#[cfg(feature = "alloc")]
20use array_reshape::Flatten;
21use core::mem::MaybeUninit;
22#[cfg(feature = "alloc")]
23use core::mem::offset_of;
24use core::{array, hint};
25#[cfg(feature = "parallel")]
26use rayon::prelude::*;
27#[cfg(any(feature = "full-chiapos", test))]
28use sha2::{Digest, Sha256};
29#[cfg(feature = "alloc")]
30use subspace_core_primitives::pieces::Record;
31#[cfg(feature = "alloc")]
32use subspace_core_primitives::pos::PosProof;
33#[cfg(feature = "alloc")]
34use subspace_core_primitives::sectors::SBucket;
35
36/// Supported K values for Chia proof of space tables
37pub impl(self) trait SupportedKValue {}
38
39/// Size of the proof in bytes for a given `K` value
40pub const PROOF_SIZE<const K: u8>: usize =
41    2usize.pow(u32::from(NUM_TABLES - 1)) * usize::from(K) / u8::BITS as usize;
42
43/// Proof-of-space proofs
44#[derive(Debug)]
45#[cfg(feature = "alloc")]
46#[repr(C)]
47pub struct Proofs<const K: u8> {
48    /// S-buckets at which proofs were found.
49    ///
50    /// S-buckets are grouped by 8, within each `u8` bits right to left (LSB) indicate the presence
51    /// of a proof for corresponding s-bucket, so that the whole array of bytes can be thought as a
52    /// large set of bits.
53    ///
54    /// There will be at most [`Record::NUM_CHUNKS`] proofs produced/bits set to `1`.
55    pub mut(self) found_proofs: [u8; Record::NUM_S_BUCKETS / u8::BITS as usize],
56    /// [`Record::NUM_CHUNKS`] proofs, corresponding to set bits of `found_proofs`.
57    pub mut(self) proofs: [[u8; PROOF_SIZE::<K>]; const { Record::NUM_CHUNKS }],
58}
59
60#[cfg(feature = "alloc")]
61impl From<Box<Proofs<const { PosProof::K }>>> for Box<PosProofs> {
62    // TODO: `no_panic::no_panic` fails to parse const generic arguments
63    // #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
64    fn from(proofs: Box<Proofs<const { PosProof::K }>>) -> Self {
65        // Statically ensure types are the same
66        const {
67            assert!(size_of::<Proofs<const { PosProof::K }>>() == size_of::<PosProofs>());
68            assert!(align_of::<Proofs<const { PosProof::K }>>() == align_of::<PosProofs>());
69            assert!(
70                offset_of!(Proofs<const { PosProof::K }>, found_proofs)
71                    == offset_of!(PosProofs, found_proofs)
72            );
73            assert!(
74                offset_of!(Proofs<const { PosProof::K }>, proofs) == offset_of!(PosProofs, proofs)
75            );
76        }
77        // SAFETY: Both structs have an identical layout with `#[repr(C)]` internals
78        unsafe { Box::from_raw(Box::into_raw(proofs).cast()) }
79    }
80}
81
82#[cfg(feature = "alloc")]
83impl<const K: u8> Proofs<K> {
84    /// Get proof for specified s-bucket (if exists).
85    ///
86    /// Note that this is not the most efficient API possible, so prefer using the `proofs` field
87    /// directly if the use case allows.
88    #[inline]
89    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
90    pub fn for_s_bucket(&self, s_bucket: SBucket) -> Option<[u8; PROOF_SIZE::<K>]> {
91        let proof_index = PosProofs::proof_index_for_s_bucket(&self.found_proofs, s_bucket)?;
92
93        // SAFETY: Protected invariant of the data structure
94        unsafe {
95            hint::assert_unchecked(proof_index < Record::NUM_CHUNKS);
96        }
97
98        Some(self.proofs[proof_index])
99    }
100
101    /// Initialize `found_proofs` with zeroes and return both fields separately so that proofs can
102    /// be written into still uninitialized `proofs`
103    #[inline(always)]
104    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
105    fn split_uninit(
106        proofs: &mut MaybeUninit<Self>,
107    ) -> (
108        &mut [u8; Record::NUM_S_BUCKETS / u8::BITS as usize],
109        &mut [MaybeUninit<[u8; PROOF_SIZE::<K>]>; const { Record::NUM_CHUNKS }],
110    ) {
111        let proofs_ptr = proofs.as_mut_ptr();
112        // SAFETY: This is the correct way to access uninit reference to the inner field
113        let found_proofs = unsafe {
114            (&raw mut (*proofs_ptr).found_proofs)
115                .as_uninit_mut()
116                .expect("Not null; qed")
117        };
118        let found_proofs = found_proofs.write([0; _]);
119        // SAFETY: This is the correct way to access uninit reference to the inner field
120        let proofs = unsafe {
121            (&raw mut (*proofs_ptr).proofs)
122                .cast::<[MaybeUninit<_>; const { Record::NUM_CHUNKS }]>()
123                .as_mut_unchecked()
124        };
125
126        (found_proofs, proofs)
127    }
128}
129
130type Seed = [u8; 32];
131#[cfg(any(feature = "full-chiapos", test))]
132type Challenge = [u8; 32];
133#[cfg(any(feature = "full-chiapos", test))]
134type Quality = [u8; 32];
135
136/// Pick position in `table_number` based on challenge bits
137#[cfg(all(feature = "alloc", any(feature = "full-chiapos", test)))]
138const fn pick_position(
139    [left_position, right_position]: [Position; 2],
140    last_5_challenge_bits: u8,
141    table_number: u8,
142) -> Position {
143    if ((last_5_challenge_bits >> (table_number - 2)) & 1) == 0 {
144        left_position
145    } else {
146        right_position
147    }
148}
149
150/// Number of positions after expanding each of `N` positions into a pair
151#[cfg(feature = "alloc")]
152const EXPANDED_POSITIONS<const N: usize>: usize = N * 2;
153
154/// Expand each position into the pair of positions in the parent table it was derived from
155#[cfg(feature = "alloc")]
156#[inline(always)]
157#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
158fn expand_positions<const N: usize>(
159    positions: [Position; N],
160    expand: impl Fn(Position) -> [Position; 2],
161) -> [Position; EXPANDED_POSITIONS::<N>] {
162    positions.map(expand).flatten()
163}
164
165/// Collection of Chia tables
166#[derive(Debug)]
167pub struct Tables<const K: u8>
168where
169    Self: SupportedKValue,
170{
171    #[cfg(feature = "alloc")]
172    table_2: PrunedTable<K, 2>,
173    #[cfg(feature = "alloc")]
174    table_3: PrunedTable<K, 3>,
175    #[cfg(feature = "alloc")]
176    table_4: PrunedTable<K, 4>,
177    #[cfg(feature = "alloc")]
178    table_5: PrunedTable<K, 5>,
179    #[cfg(feature = "alloc")]
180    table_6: PrunedTable<K, 6>,
181    #[cfg(feature = "alloc")]
182    table_7: Table<K, 7>,
183}
184
185impl<const K: u8> Tables<K>
186where
187    Self: SupportedKValue,
188{
189    /// Create Chia proof of space tables.
190    ///
191    /// There is also `Self::create_parallel()` that can achieve higher performance and lower
192    /// latency at the cost of lower CPU efficiency and higher memory usage.
193    #[cfg(all(feature = "alloc", any(feature = "full-chiapos", test)))]
194    pub fn create(seed: Seed) -> Self {
195        let table_1 = Table::<K, 1>::create(seed);
196        let (table_2, _) = Table::<K, 2>::create(table_1);
197        let (table_3, table_2) = Table::<K, 3>::create(table_2);
198        let (table_4, table_3) = Table::<K, 4>::create(table_3);
199        let (table_5, table_4) = Table::<K, 5>::create(table_4);
200        let (table_6, table_5) = Table::<K, 6>::create(table_5);
201        let (table_7, table_6) = Table::<K, 7>::create(table_6);
202
203        Self {
204            table_2,
205            table_3,
206            table_4,
207            table_5,
208            table_6,
209            table_7,
210        }
211    }
212
213    /// Create proofs.
214    ///
215    /// This is an optimized combination of `Self::create()` and `Self::find_proof()`.
216    ///
217    /// There is also `Self::create_proofs_parallel()` that can achieve higher performance and lower
218    /// latency at the cost of lower CPU efficiency and higher memory usage.
219    #[cfg(feature = "alloc")]
220    pub fn create_proofs(seed: Seed) -> Box<Proofs<K>> {
221        let table_1 = Table::<K, 1>::create(seed);
222        let (table_2, _) = Table::<K, 2>::create(table_1);
223        let (table_3, table_2) = Table::<K, 3>::create(table_2);
224        let (table_4, table_3) = Table::<K, 4>::create(table_3);
225        let (table_5, table_4) = Table::<K, 5>::create(table_4);
226        let (table_6, table_5) = Table::<K, 6>::create(table_5);
227        let (table_6_proof_targets, table_6) = Table::<K, 7>::create_proof_targets(table_6);
228
229        let mut proofs = Box::<Proofs<K>>::new_uninit();
230        Self::find_proofs_internal(
231            &table_2,
232            &table_3,
233            &table_4,
234            &table_5,
235            &table_6,
236            &table_6_proof_targets,
237            &mut proofs,
238        );
239
240        // SAFETY: Fully and correctly initialized above
241        unsafe { proofs.assume_init() }
242    }
243
244    /// Find a proof for each s-bucket that has a target in the last table
245    // TODO: Rewrite this more efficiently
246    #[cfg(feature = "alloc")]
247    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
248    fn find_proofs_internal(
249        table_2: &PrunedTable<K, 2>,
250        table_3: &PrunedTable<K, 3>,
251        table_4: &PrunedTable<K, 4>,
252        table_5: &PrunedTable<K, 5>,
253        table_6: &PrunedTable<K, 6>,
254        table_6_proof_targets: &[[Position; 2]; const { Record::NUM_S_BUCKETS }],
255        proofs: &mut MaybeUninit<Proofs<K>>,
256    ) {
257        let (found_proofs, proofs) = Proofs::<K>::split_uninit(proofs);
258
259        let mut num_found_proofs = 0_usize;
260        'outer: for (table_6_proof_targets, found_proofs) in table_6_proof_targets
261            .as_chunks::<{ u8::BITS as usize }>()
262            .0
263            .iter()
264            .zip(found_proofs)
265        {
266            for (proof_offset, table_6_proof_targets) in table_6_proof_targets.iter().enumerate() {
267                if table_6_proof_targets != &[Position::ZERO; 2] {
268                    let proof = Self::find_proof_raw_internal(
269                        table_2,
270                        table_3,
271                        table_4,
272                        table_5,
273                        table_6,
274                        *table_6_proof_targets,
275                    );
276
277                    *found_proofs |= 1 << proof_offset;
278
279                    // TODO: Remove once https://github.com/rust-lang/rust/issues/162834 is resolved
280                    // SAFETY: The loop is stopped as soon as `Record::NUM_CHUNKS` proofs are found
281                    unsafe {
282                        hint::assert_unchecked(num_found_proofs < Record::NUM_CHUNKS);
283                    }
284                    proofs[num_found_proofs].write(proof);
285                    num_found_proofs += 1;
286
287                    if num_found_proofs == Record::NUM_CHUNKS {
288                        break 'outer;
289                    }
290                }
291            }
292        }
293
294        // It is statically known to be the case, and there is a test that checks the lower bound
295        debug_assert_eq!(num_found_proofs, Record::NUM_CHUNKS);
296    }
297
298    /// Almost the same as [`Self::create()`], but uses parallelism internally for better
299    /// performance and lower latency at the cost of lower CPU efficiency and higher memory usage
300    #[cfg(all(feature = "parallel", any(feature = "full-chiapos", test)))]
301    pub fn create_parallel(seed: Seed) -> Self {
302        let table_1 = Table::<K, 1>::create_parallel(seed);
303        let (table_2, _) = Table::<K, 2>::create_parallel(table_1);
304        let (table_3, table_2) = Table::<K, 3>::create_parallel(table_2);
305        let (table_4, table_3) = Table::<K, 4>::create_parallel(table_3);
306        let (table_5, table_4) = Table::<K, 5>::create_parallel(table_4);
307        let (table_6, table_5) = Table::<K, 6>::create_parallel(table_5);
308        let (table_7, table_6) = Table::<K, 7>::create_parallel(table_6);
309
310        Self {
311            table_2,
312            table_3,
313            table_4,
314            table_5,
315            table_6,
316            table_7,
317        }
318    }
319
320    /// Almost the same as [`Self::create_proofs()`], but uses parallelism internally for better
321    /// performance and lower latency at the cost of lower CPU efficiency and higher memory usage
322    #[cfg(feature = "parallel")]
323    pub fn create_proofs_parallel(seed: Seed) -> Box<Proofs<K>> {
324        let table_1 = Table::<K, 1>::create_parallel(seed);
325        let (table_2, _) = Table::<K, 2>::create_parallel(table_1);
326        let (table_3, table_2) = Table::<K, 3>::create_parallel(table_2);
327        let (table_4, table_3) = Table::<K, 4>::create_parallel(table_3);
328        let (table_5, table_4) = Table::<K, 5>::create_parallel(table_4);
329        let (table_6, table_5) = Table::<K, 6>::create_parallel(table_5);
330        let (table_6_proof_targets, table_6) =
331            Table::<K, 7>::create_proof_targets_parallel(table_6);
332
333        let mut proofs = Box::<Proofs<K>>::new_uninit();
334        // SAFETY: Contents is `MaybeUninit`
335        let mut targets = unsafe {
336            Box::<[MaybeUninit<[Position; 2]>; const { Record::NUM_CHUNKS }]>::new_uninit()
337                .assume_init()
338        };
339        {
340            let (found_proofs, proofs) = Proofs::<K>::split_uninit(&mut proofs);
341
342            // Deciding which s-buckets have a proof is cheap, so it is done sequentially, leaving
343            // only the expensive part below to do in parallel
344            let targets =
345                Self::collect_proof_targets(&table_6_proof_targets, found_proofs, &mut targets);
346
347            // Work items here are large enough that `rayon::broadcast()` with manual batching (as
348            // used elsewhere in this crate) measures the same, so the safe version is used
349            proofs[..targets.len()]
350                .par_iter_mut()
351                .zip(targets)
352                .for_each(|(proof, &table_6_proof_targets)| {
353                    proof.write(Self::find_proof_raw_internal(
354                        &table_2,
355                        &table_3,
356                        &table_4,
357                        &table_5,
358                        &table_6,
359                        table_6_proof_targets,
360                    ));
361                });
362        }
363
364        // SAFETY: Fully and correctly initialized
365        unsafe { proofs.assume_init() }
366    }
367
368    /// Collect targets in the last table for s-buckets that have a proof, which is the cheap
369    /// sequential part of [`Self::create_proofs_parallel()`]
370    #[cfg(feature = "parallel")]
371    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
372    fn collect_proof_targets<'a>(
373        table_6_proof_targets: &[[Position; 2]; const { Record::NUM_S_BUCKETS }],
374        found_proofs: &mut [u8; Record::NUM_S_BUCKETS / u8::BITS as usize],
375        targets: &'a mut [MaybeUninit<[Position; 2]>; const { Record::NUM_CHUNKS }],
376    ) -> &'a [[Position; 2]] {
377        let mut num_found_proofs = 0_usize;
378
379        'outer: for (table_6_proof_targets, found_proofs) in table_6_proof_targets
380            .as_chunks::<{ u8::BITS as usize }>()
381            .0
382            .iter()
383            .zip(found_proofs)
384        {
385            for (proof_offset, table_6_proof_targets) in table_6_proof_targets.iter().enumerate() {
386                if table_6_proof_targets != &[Position::ZERO; 2] {
387                    *found_proofs |= 1 << proof_offset;
388
389                    // TODO: Remove once https://github.com/rust-lang/rust/issues/162834 is resolved
390                    // SAFETY: The loop is stopped as soon as `Record::NUM_CHUNKS` targets are
391                    // collected
392                    unsafe {
393                        hint::assert_unchecked(num_found_proofs < Record::NUM_CHUNKS);
394                    }
395                    targets[num_found_proofs].write(*table_6_proof_targets);
396                    num_found_proofs += 1;
397
398                    if num_found_proofs == Record::NUM_CHUNKS {
399                        break 'outer;
400                    }
401                }
402            }
403        }
404
405        // TODO: Remove once https://github.com/rust-lang/rust/issues/162834 is resolved
406        // SAFETY: The loop above is stopped as soon as `Record::NUM_CHUNKS` targets are collected
407        unsafe {
408            hint::assert_unchecked(num_found_proofs <= Record::NUM_CHUNKS);
409        }
410
411        // SAFETY: Initialized this many elements above
412        unsafe { targets[..num_found_proofs].assume_init_ref() }
413    }
414
415    /// Find proof of space quality for a given challenge
416    #[cfg(all(feature = "alloc", any(feature = "full-chiapos", test)))]
417    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
418    pub fn find_quality<'a>(
419        &'a self,
420        challenge: &'a Challenge,
421    ) -> impl Iterator<Item = Quality> + 'a {
422        let last_5_challenge_bits = challenge[challenge.len() - 1] & 0b0001_1111;
423
424        let first_k_challenge_bits = u32::from_be_bytes(
425            challenge[..size_of::<u32>()]
426                .try_into()
427                .expect("Challenge is known to statically have enough bytes; qed"),
428        ) >> (u32::BITS as usize - usize::from(K));
429
430        // SAFETY: Bucket range is by definition in bounds
431        let bucket = unsafe {
432            self.table_7
433                .buckets()
434                .get_unchecked(Y::bucket_range_from_first_k_bits(first_k_challenge_bits))
435        };
436        // Iterate just over elements that are matching `first_k_challenge_bits` prefix
437        bucket
438            .iter()
439            .flat_map(move |positions| {
440                positions
441                    .iter()
442                    .take_while(|&&(position, _y)| position != Position::SENTINEL)
443                    .filter(move |&&(_position, y)| y.first_k_bits() == first_k_challenge_bits)
444            })
445            .map(move |&(position, _y)| {
446                // SAFETY: Internally generated positions that come from the parent table
447                let positions = unsafe { self.table_7.position(position) };
448                // SAFETY: Internally generated positions that come from the parent table
449                let positions = unsafe {
450                    self.table_6
451                        .position(pick_position(positions, last_5_challenge_bits, 6))
452                };
453                // SAFETY: Internally generated positions that come from the parent table
454                let positions = unsafe {
455                    self.table_5
456                        .position(pick_position(positions, last_5_challenge_bits, 5))
457                };
458                // SAFETY: Internally generated positions that come from the parent table
459                let positions = unsafe {
460                    self.table_4
461                        .position(pick_position(positions, last_5_challenge_bits, 4))
462                };
463                // SAFETY: Internally generated positions that come from the parent table
464                let positions = unsafe {
465                    self.table_3
466                        .position(pick_position(positions, last_5_challenge_bits, 3))
467                };
468                // SAFETY: Internally generated positions that come from the parent table
469                let [left_position, right_position] = unsafe {
470                    self.table_2
471                        .position(pick_position(positions, last_5_challenge_bits, 2))
472                };
473
474                // X matches position
475                let left_x = X::from(u32::from(left_position));
476                let right_x = X::from(u32::from(right_position));
477
478                let mut hasher = Sha256::new();
479                hasher.update(challenge);
480                let left_right_xs = (u64::from(left_x) << (u64::BITS as usize - usize::from(K)))
481                    | (u64::from(right_x) << (u64::BITS as usize - usize::from(K * 2)));
482                hasher.update(
483                    &left_right_xs.to_be_bytes()
484                        [..(usize::from(K) * 2).div_ceil(u8::BITS as usize)],
485                );
486                hasher.finalize().into()
487            })
488    }
489
490    /// Similar to `Self::find_proof()`, but takes the first `k` challenge bits in the least
491    /// significant bits of `u32` as a challenge instead
492    #[cfg(feature = "alloc")]
493    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
494    pub fn find_proof_raw(
495        &self,
496        first_k_challenge_bits: u32,
497    ) -> impl Iterator<Item = [u8; PROOF_SIZE::<K>]> + '_ {
498        // Iterate just over elements that are matching `first_k_challenge_bits` prefix
499        self.table_7
500            .buckets()
501            .get(Y::bucket_range_from_first_k_bits(first_k_challenge_bits))
502            .unwrap_or(&[])
503            .iter()
504            .flat_map(move |positions| {
505                positions
506                    .iter()
507                    .take_while(|&&(position, _y)| position != Position::SENTINEL)
508                    .filter(move |&&(_position, y)| y.first_k_bits() == first_k_challenge_bits)
509            })
510            .map(move |&(position, _y)| {
511                // SAFETY: Internally generated positions that come from the parent table
512                let table_6_proof_targets = unsafe { self.table_7.position(position) };
513
514                Self::find_proof_raw_internal(
515                    &self.table_2,
516                    &self.table_3,
517                    &self.table_4,
518                    &self.table_5,
519                    &self.table_6,
520                    table_6_proof_targets,
521                )
522            })
523    }
524
525    #[cfg(feature = "alloc")]
526    #[inline(always)]
527    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
528    fn find_proof_raw_internal(
529        table_2: &PrunedTable<K, 2>,
530        table_3: &PrunedTable<K, 3>,
531        table_4: &PrunedTable<K, 4>,
532        table_5: &PrunedTable<K, 5>,
533        table_6: &PrunedTable<K, 6>,
534        table_6_proof_targets: [Position; 2],
535    ) -> [u8; PROOF_SIZE::<K>] {
536        let mut proof = [0u8; _];
537
538        // Positions are expanded one table at a time rather than by walking the tree of positions
539        // depth-first. All lookups within a single table are independent of each other, which
540        // allows many of these cache misses to be in flight at the same time instead of waiting
541        // for each other. This mirrors [`Self::verify_only_raw()`], which walks the same tree in
542        // the opposite direction a table at a time for the same reason.
543        let positions = expand_positions(table_6_proof_targets, |position| {
544            // SAFETY: Internally generated positions that come from the parent table
545            unsafe { table_6.position(position) }
546        });
547        let positions = expand_positions(positions, |position| {
548            // SAFETY: Internally generated positions that come from the parent table
549            unsafe { table_5.position(position) }
550        });
551        let positions = expand_positions(positions, |position| {
552            // SAFETY: Internally generated positions that come from the parent table
553            unsafe { table_4.position(position) }
554        });
555        let positions = expand_positions(positions, |position| {
556            // SAFETY: Internally generated positions that come from the parent table
557            unsafe { table_3.position(position) }
558        });
559        let positions = expand_positions(positions, |position| {
560            // SAFETY: Internally generated positions that come from the parent table
561            unsafe { table_2.position(position) }
562        });
563
564        positions
565            .iter()
566            .map(|&position| {
567                // X matches position
568                X::from(u32::from(position))
569            })
570            .enumerate()
571            .for_each(|(offset, x)| {
572                let x_offset_in_bits = usize::from(K) * offset;
573                // Collect bytes where bits of `x` will be written
574                let proof_bytes = &mut proof[x_offset_in_bits / u8::BITS as usize..]
575                    [..(x_offset_in_bits % u8::BITS as usize + usize::from(K))
576                        .div_ceil(u8::BITS as usize)];
577
578                // Bits of `x` already shifted to the correct location as they will appear
579                // in `proof`
580                let x_shifted = u32::from(x)
581                    << (u32::BITS as usize
582                        - (usize::from(K) + x_offset_in_bits % u8::BITS as usize));
583
584                // TODO: Store proofs in words, like GPU version does
585                // Copy `x` bits into proof
586                x_shifted
587                    .to_be_bytes()
588                    .iter()
589                    .zip(proof_bytes)
590                    .for_each(|(from, to)| {
591                        *to |= from;
592                    });
593            });
594
595        proof
596    }
597
598    /// Find proof of space for a given challenge
599    #[cfg(all(feature = "alloc", any(feature = "full-chiapos", test)))]
600    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
601    pub fn find_proof(
602        &self,
603        first_challenge_bytes: [u8; 4],
604    ) -> impl Iterator<Item = [u8; PROOF_SIZE::<K>]> + '_ {
605        let first_k_challenge_bits =
606            u32::from_be_bytes(first_challenge_bytes) >> (u32::BITS as usize - usize::from(K));
607
608        self.find_proof_raw(first_k_challenge_bits)
609    }
610
611    /// Similar to `Self::verify()`, but takes the first `k` challenge bits in the least significant
612    /// bits of `u32` as a challenge instead and doesn't compute quality
613    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
614    pub fn verify_only_raw(
615        seed: &Seed,
616        first_k_challenge_bits: u32,
617        proof_of_space: &[u8; PROOF_SIZE::<K>],
618    ) -> bool {
619        let ys_and_metadata = array::from_fn::<_, 64, _>(|offset| {
620            let mut pre_x_bytes = 0u64.to_be_bytes();
621            let offset_in_bits = usize::from(K) * offset;
622            let bytes_to_copy =
623                (offset_in_bits % u8::BITS as usize + usize::from(K)).div_ceil(u8::BITS as usize);
624            // Copy full bytes that contain bits of `x`
625            pre_x_bytes[..bytes_to_copy].copy_from_slice(
626                &proof_of_space[offset_in_bits / u8::BITS as usize..][..bytes_to_copy],
627            );
628            // Extract `pre_x` whose last `K` bits start with `x`
629            let pre_x = u64::from_be_bytes(pre_x_bytes)
630                >> (u64::BITS as usize - (usize::from(K) + offset_in_bits % u8::BITS as usize));
631            // Convert to the desired type and clear extra bits
632            let x = X::from(pre_x as u32 & (u32::MAX >> (u32::BITS as usize - usize::from(K))));
633
634            let y = compute_f1::<K>(x, seed);
635
636            (y, Metadata::from(x))
637        });
638
639        let mut next_ys_and_metadata = [MaybeUninit::uninit(); _];
640        let ys_and_metadata =
641            Self::collect_ys_and_metadata::<2, 1, 64>(&ys_and_metadata, &mut next_ys_and_metadata);
642        let mut next_ys_and_metadata = [MaybeUninit::uninit(); _];
643        let ys_and_metadata =
644            Self::collect_ys_and_metadata::<3, 2, 32>(ys_and_metadata, &mut next_ys_and_metadata);
645        let mut next_ys_and_metadata = [MaybeUninit::uninit(); _];
646        let ys_and_metadata =
647            Self::collect_ys_and_metadata::<4, 3, 16>(ys_and_metadata, &mut next_ys_and_metadata);
648        let mut next_ys_and_metadata = [MaybeUninit::uninit(); _];
649        let ys_and_metadata =
650            Self::collect_ys_and_metadata::<5, 4, 8>(ys_and_metadata, &mut next_ys_and_metadata);
651        let mut next_ys_and_metadata = [MaybeUninit::uninit(); _];
652        let ys_and_metadata =
653            Self::collect_ys_and_metadata::<6, 5, 4>(ys_and_metadata, &mut next_ys_and_metadata);
654        let mut next_ys_and_metadata = [MaybeUninit::uninit(); _];
655        let ys_and_metadata =
656            Self::collect_ys_and_metadata::<7, 6, 2>(ys_and_metadata, &mut next_ys_and_metadata);
657
658        let Some((y, _metadata)) = ys_and_metadata.first() else {
659            return false;
660        };
661
662        // Check if the first K bits of `y` match
663        y.first_k_bits() == first_k_challenge_bits
664    }
665
666    /// Verify proof of space for a given seed and challenge
667    // TODO: `no_panic::no_panic` can't prove lack of panics in `sha2`
668    // #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
669    #[cfg(any(feature = "full-chiapos", test))]
670    pub fn verify(
671        seed: &Seed,
672        challenge: &Challenge,
673        proof_of_space: &[u8; PROOF_SIZE::<K>],
674    ) -> Option<Quality> {
675        let first_k_challenge_bits =
676            u32::from_be_bytes([challenge[0], challenge[1], challenge[2], challenge[3]])
677                >> (u32::BITS as usize - usize::from(K));
678
679        if !Self::verify_only_raw(seed, first_k_challenge_bits, proof_of_space) {
680            return None;
681        }
682
683        let last_5_challenge_bits = challenge[challenge.len() - 1] & 0b0001_1111;
684
685        let mut quality_index = 0_usize.to_be_bytes();
686        quality_index[0] = last_5_challenge_bits;
687        let quality_index = usize::from_be_bytes(quality_index);
688
689        // NOTE: this works correctly but may overflow if `quality_index` is changed to
690        // not be zero-initialized anymore
691        let left_right_xs_bit_offset = quality_index * usize::from(K * 2);
692        // Collect `left_x` and `right_x` bits, potentially with extra bits at the beginning
693        // and the end
694        let left_right_xs_bytes = &proof_of_space[left_right_xs_bit_offset / u8::BITS as usize..]
695            [..(left_right_xs_bit_offset % u8::BITS as usize + usize::from(K * 2))
696                .div_ceil(u8::BITS as usize)];
697
698        let mut left_right_xs = 0u64.to_be_bytes();
699        left_right_xs[..left_right_xs_bytes.len()].copy_from_slice(left_right_xs_bytes);
700        // Move `left_x` and `right_x` bits to most significant bits
701        let left_right_xs =
702            u64::from_be_bytes(left_right_xs) << (left_right_xs_bit_offset % u8::BITS as usize);
703        // Clear extra bits
704        let left_right_xs_mask = u64::MAX << (u64::BITS as usize - usize::from(K * 2));
705        let left_right_xs = left_right_xs & left_right_xs_mask;
706
707        let mut hasher = Sha256::new();
708        hasher.update(challenge);
709        hasher
710            .update(&left_right_xs.to_be_bytes()[..usize::from(K * 2).div_ceil(u8::BITS as usize)]);
711        Some(hasher.finalize().into())
712    }
713
714    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
715    fn collect_ys_and_metadata<
716        'a,
717        const TABLE_NUMBER: u8,
718        const PARENT_TABLE_NUMBER: u8,
719        const N: usize,
720    >(
721        ys_and_metadata: &[(Y, Metadata<K, PARENT_TABLE_NUMBER>)],
722        next_ys_and_metadata: &'a mut [MaybeUninit<(Y, Metadata<K, TABLE_NUMBER>)>; N],
723    ) -> &'a [(Y, Metadata<K, TABLE_NUMBER>)] {
724        let mut next_offset = 0_usize;
725        for &[(left_y, left_metadata), (right_y, right_metadata)] in
726            ys_and_metadata.as_chunks::<2>().0
727        {
728            if !has_match(left_y, right_y) {
729                continue;
730            }
731
732            // SAFETY: Inputs are limited to `N * 2` elements, and at most one match is produced for
733            // every two of them
734            unsafe {
735                hint::assert_unchecked(next_offset < N);
736            }
737            next_ys_and_metadata[next_offset].write(compute_fn::<
738                K,
739                TABLE_NUMBER,
740                PARENT_TABLE_NUMBER,
741            >(
742                left_y, left_metadata, right_metadata
743            ));
744            next_offset += 1;
745        }
746
747        // TODO: Remove once https://github.com/rust-lang/rust/issues/162834 is resolved
748        // SAFETY: Inputs are limited to `N * 2` elements, and at most one match is produced for
749        // every two of them
750        unsafe {
751            hint::assert_unchecked(next_offset <= N);
752        }
753
754        // SAFETY: Initialized `next_offset` elements
755        unsafe { next_ys_and_metadata[..next_offset].assume_init_ref() }
756    }
757}
758
759macro_rules! impl_supported {
760    ($($k: expr$(,)? )*) => {
761        $(
762impl SupportedKValue for Tables<$k> {}
763        )*
764    }
765}
766
767// Only these k values are supported by the current implementation
768#[cfg(feature = "full-chiapos")]
769impl_supported!(15, 16, 18, 19, 21, 22, 23, 24, 25);
770#[cfg(any(feature = "full-chiapos", test))]
771impl_supported!(17);
772impl_supported!(20);
ab_proof_of_space/chiapos.rs

ab_proof_of_space/
chiapos.rs
