Skip to main content

ab_proof_of_space/
shim.rs

1//! Shim proof of space implementation that works much faster than Chia and can be used for testing
2//! purposes to reduce memory and CPU usage
3
4#[cfg(all(feature = "alloc", test, not(miri)))]
5mod tests;
6
7#[cfg(feature = "alloc")]
8use crate::PosProofs;
9#[cfg(feature = "alloc")]
10use crate::TableGenerator;
11use crate::{PosTableType, Table};
12#[cfg(feature = "alloc")]
13use ab_core_primitives::pieces::Record;
14use ab_core_primitives::pos::{PosProof, PosSeed};
15use ab_core_primitives::sectors::SBucket;
16#[cfg(feature = "alloc")]
17use alloc::boxed::Box;
18#[cfg(feature = "alloc")]
19use core::hint;
20use core::iter;
21
22/// Proof of space table generator.
23///
24/// Shim implementation.
25#[derive(Debug, Default, Clone)]
26#[cfg(feature = "alloc")]
27pub struct ShimTableGenerator;
28
29#[cfg(feature = "alloc")]
30impl TableGenerator<ShimTable> for ShimTableGenerator {
31    fn create_proofs(&self, seed: &PosSeed) -> Box<PosProofs> {
32        // SAFETY: Data structure filled with zeroes is a valid invariant
33        let mut proofs = unsafe { Box::<PosProofs>::new_zeroed().assume_init() };
34
35        create_proofs_internal(seed, &mut proofs);
36
37        proofs
38    }
39}
40
41/// Find proofs for as many s-buckets as fit into `proofs`, which must be zero-initialized
42#[cfg(feature = "alloc")]
43#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
44fn create_proofs_internal(seed: &PosSeed, proofs: &mut PosProofs) {
45    let mut num_found_proofs = 0_usize;
46
47    'outer: for (s_buckets, found_proofs) in (0..Record::NUM_S_BUCKETS as u32)
48        .array_chunks::<{ u8::BITS as usize }>()
49        .zip(&mut proofs.found_proofs)
50    {
51        for (proof_offset, s_bucket) in s_buckets.into_iter().enumerate() {
52            if let Some(proof) = find_proof(seed, s_bucket) {
53                *found_proofs |= 1 << proof_offset;
54
55                // TODO: Remove once https://github.com/rust-lang/rust/issues/162834 is resolved
56                // SAFETY: The loop is stopped as soon as `Record::NUM_CHUNKS` proofs are found
57                unsafe {
58                    hint::assert_unchecked(num_found_proofs < Record::NUM_CHUNKS);
59                }
60                proofs.proofs[num_found_proofs] = proof;
61                num_found_proofs += 1;
62
63                if num_found_proofs == Record::NUM_CHUNKS {
64                    break 'outer;
65                }
66            }
67        }
68    }
69}
70
71/// Proof of space table.
72///
73/// Shim implementation.
74#[derive(Debug)]
75pub struct ShimTable;
76
77impl ab_core_primitives::solutions::SolutionPotVerifier for ShimTable {
78    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
79    fn is_proof_valid(seed: &PosSeed, s_bucket: SBucket, proof: &PosProof) -> bool {
80        let Some(correct_proof) = find_proof(seed, u32::from(s_bucket)) else {
81            return false;
82        };
83
84        &correct_proof == proof
85    }
86}
87
88impl Table for ShimTable {
89    const TABLE_TYPE: PosTableType = PosTableType::Shim;
90    #[cfg(feature = "alloc")]
91    type Generator = ShimTableGenerator;
92
93    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
94    fn is_proof_valid(seed: &PosSeed, s_bucket: SBucket, proof: &PosProof) -> bool {
95        <Self as ab_core_primitives::solutions::SolutionPotVerifier>::is_proof_valid(
96            seed, s_bucket, proof,
97        )
98    }
99}
100
101#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
102fn find_proof(seed: &PosSeed, challenge_index: u32) -> Option<PosProof> {
103    let quality = ab_blake3::single_block_hash(&challenge_index.to_le_bytes())
104        .expect("Less than a single block worth of bytes; qed");
105    if quality[0].is_multiple_of(3) {
106        None
107    } else {
108        let mut proof = PosProof::default();
109        proof
110            .iter_mut()
111            .zip(seed.iter().chain(iter::repeat(quality.iter()).flatten()))
112            .for_each(|(output, input)| {
113                *output = *input;
114            });
115
116        Some(proof)
117    }
118}