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;
18use core::iter;
19
20/// Proof of space table generator.
21///
22/// Shim implementation.
23#[derive(Debug, Default, Clone)]
24#[cfg(feature = "alloc")]
25pub struct ShimTableGenerator;
26
27#[cfg(feature = "alloc")]
28impl TableGenerator<ShimTable> for ShimTableGenerator {
29    fn create_proofs(&self, seed: &PosSeed) -> Box<PosProofs> {
30        // SAFETY: Data structure filled with zeroes is a valid invariant
31        let mut proofs = unsafe { Box::<PosProofs>::new_zeroed().assume_init() };
32
33        create_proofs_internal(seed, &mut proofs);
34
35        proofs
36    }
37}
38
39/// Find proofs for exactly one s-bucket of each pair of neighboring s-buckets.
40///
41/// That gives exactly [`Record::NUM_CHUNKS`] proofs, as many as fit into `proofs`, which must be
42/// zero-initialized.
43#[cfg(feature = "alloc")]
44#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
45fn create_proofs_internal(seed: &PosSeed, proofs: &mut PosProofs) {
46    // Each byte of `found_proofs` covers 8 s-buckets, which are 4 pairs
47    const PAIRS_PER_BYTE: usize = u8::BITS as usize / 2;
48    const {
49        assert!(Record::NUM_S_BUCKETS == Record::NUM_CHUNKS * 2);
50    }
51
52    let PosProofs {
53        found_proofs,
54        proofs,
55    } = proofs;
56    let (proofs, _) = proofs.as_chunks_mut::<PAIRS_PER_BYTE>();
57
58    for ((found_proofs, proofs), first_pair_index) in found_proofs
59        .iter_mut()
60        .zip(proofs)
61        .zip((0..).step_by(PAIRS_PER_BYTE))
62    {
63        for ((output, pair_index), bit_offset) in proofs
64            .iter_mut()
65            .zip(first_pair_index..)
66            .zip((0..).step_by(2))
67        {
68            let (second_has_proof, proof) = pair_proof(seed, pair_index);
69            *found_proofs |= 1 << (bit_offset + u32::from(second_has_proof));
70            *output = proof;
71        }
72    }
73}
74
75/// Proof of space table.
76///
77/// Shim implementation.
78#[derive(Debug)]
79pub struct ShimTable;
80
81impl ab_core_primitives::solutions::SolutionPotVerifier for ShimTable {
82    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
83    fn is_proof_valid(seed: &PosSeed, s_bucket: SBucket, proof: &PosProof) -> bool {
84        let Some(correct_proof) = find_proof(seed, u32::from(s_bucket)) else {
85            return false;
86        };
87
88        &correct_proof == proof
89    }
90}
91
92impl Table for ShimTable {
93    const TABLE_TYPE: PosTableType = PosTableType::Shim;
94    #[cfg(feature = "alloc")]
95    type Generator = ShimTableGenerator;
96
97    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
98    fn is_proof_valid(seed: &PosSeed, s_bucket: SBucket, proof: &PosProof) -> bool {
99        <Self as ab_core_primitives::solutions::SolutionPotVerifier>::is_proof_valid(
100            seed, s_bucket, proof,
101        )
102    }
103}
104
105#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
106fn find_proof(seed: &PosSeed, challenge_index: u32) -> Option<PosProof> {
107    let (second_has_proof, proof) = pair_proof(seed, challenge_index / 2);
108    let is_second = !challenge_index.is_multiple_of(2);
109
110    (second_has_proof == is_second).then_some(proof)
111}
112
113/// Proof of the pair of neighboring s-buckets `2 * pair_index` and `2 * pair_index + 1`.
114///
115/// Only one of them has a proof, the second one if `true` is returned.
116#[cfg_attr(feature = "no-panic", no_panic::no_panic)]
117fn pair_proof(seed: &PosSeed, pair_index: u32) -> (bool, PosProof) {
118    let mut bytes_to_hash = [0; PosSeed::SIZE + size_of::<u32>()];
119    bytes_to_hash[..PosSeed::SIZE].copy_from_slice(seed.as_slice());
120    bytes_to_hash[PosSeed::SIZE..].copy_from_slice(&pair_index.to_le_bytes());
121    let quality = ab_blake3::single_block_hash(&bytes_to_hash)
122        .expect("Less than a single block worth of bytes; qed");
123
124    let mut proof = PosProof::default();
125    proof
126        .iter_mut()
127        .zip(seed.iter().chain(iter::repeat(quality.iter()).flatten()))
128        .for_each(|(output, input)| {
129            *output = *input;
130        });
131
132    (quality[0] & 1 == 1, proof)
133}