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

ab_proof_of_space/
chiapos.rs
