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

ab_proof_of_space/
chiapos.rs
