Skip to main content

ab_proof_of_space/
lib.rs

1//! Proof of space implementation
2#![no_std]
3#![expect(incomplete_features, reason = "generic_const_*")]
4#![warn(rust_2018_idioms, missing_debug_implementations, missing_docs)]
5#![feature(
6    const_block_items,
7    const_convert,
8    const_trait_impl,
9    funnel_shifts,
10    generic_const_args,
11    generic_const_items,
12    impl_restriction,
13    inherent_associated_types,
14    macroless_generic_const_args,
15    min_generic_const_args,
16    mut_restriction,
17    portable_simd,
18    step_trait
19)]
20#![cfg_attr(test, feature(float_erf))]
21#![cfg_attr(feature = "parallel", feature(exact_size_is_empty, sync_unsafe_cell))]
22#![cfg_attr(
23    feature = "alloc",
24    feature(iter_array_chunks, maybe_uninit_fill, ptr_as_uninit)
25)]
26
27pub mod chia;
28pub mod chiapos;
29pub mod shim;
30
31#[cfg(feature = "alloc")]
32extern crate alloc;
33
34#[cfg(feature = "alloc")]
35use ab_core_primitives::pieces::Record;
36use ab_core_primitives::pos::{PosProof, PosSeed};
37use ab_core_primitives::sectors::SBucket;
38use ab_core_primitives::solutions::SolutionPotVerifier;
39#[cfg(feature = "alloc")]
40use alloc::boxed::Box;
41#[cfg(feature = "alloc")]
42use core::fmt;
43#[cfg(feature = "alloc")]
44use core::hint;
45
46/// Proof of space table type
47#[derive(Debug, Clone, Copy)]
48pub enum PosTableType {
49    /// Chia table
50    Chia,
51    /// Shim table
52    Shim,
53}
54
55// TODO: Return a single full proof and the rest as hashes instead to optimize memory usage and
56//  parallelize compute more easily
57/// Proof-of-space proofs
58#[derive(Debug)]
59#[cfg(feature = "alloc")]
60#[repr(C)]
61pub struct PosProofs {
62    /// S-buckets at which proofs were found.
63    ///
64    /// S-buckets are grouped by 8, within each `u8` bits right to left (LSB) indicate the presence
65    /// of a proof for corresponding s-bucket, so that the whole array of bytes can be thought as a
66    /// large set of bits.
67    ///
68    /// There will be at most [`Record::NUM_CHUNKS`] proofs produced/bits set to `1`.
69    pub mut(self) found_proofs: [u8; Record::NUM_S_BUCKETS / u8::BITS as usize],
70    /// [`Record::NUM_CHUNKS`] proofs, corresponding to set bits of `found_proofs`.
71    pub mut(self) proofs: [PosProof; const { Record::NUM_CHUNKS }],
72}
73
74// TODO: A method that returns hashed proofs (with SIMD) for all s-buckets for plotting
75#[cfg(feature = "alloc")]
76impl PosProofs {
77    /// Get proof for specified s-bucket (if exists).
78    ///
79    /// Note that this is not the most efficient API possible, so prefer using the `proofs` field
80    /// directly if the use case allows.
81    #[inline]
82    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
83    pub fn for_s_bucket(&self, s_bucket: SBucket) -> Option<PosProof> {
84        let proof_index = Self::proof_index_for_s_bucket(&self.found_proofs, s_bucket)?;
85
86        // SAFETY: Protected invariant of the data structure
87        unsafe {
88            hint::assert_unchecked(proof_index < Record::NUM_CHUNKS);
89        }
90
91        Some(self.proofs[proof_index])
92    }
93
94    #[inline(always)]
95    #[cfg_attr(feature = "no-panic", no_panic::no_panic)]
96    fn proof_index_for_s_bucket(
97        found_proofs: &[u8; Record::NUM_S_BUCKETS / u8::BITS as usize],
98        s_bucket: SBucket,
99    ) -> Option<usize> {
100        let bits_offset = usize::from(s_bucket);
101        let found_proofs_byte_offset = bits_offset / u8::BITS as usize;
102        let found_proofs_bit_offset = bits_offset as u32 % u8::BITS;
103        let (found_proofs_before, found_proofs_after) =
104            found_proofs.split_at(found_proofs_byte_offset);
105        if (found_proofs_after[0] & (1 << found_proofs_bit_offset)) == 0 {
106            return None;
107        }
108        let proof_index = found_proofs_before
109            .iter()
110            .map(|&bits| bits.count_ones())
111            .sum::<u32>()
112            + found_proofs_after[0]
113                .unbounded_shl(u8::BITS - found_proofs_bit_offset)
114                .count_ones();
115
116        Some(proof_index as usize)
117    }
118}
119
120// TODO: Think about redesigning this API now that proofs are the output rather than tables
121/// Stateful table generator with better performance.
122///
123/// Prefer cloning it over creating multiple separate generators.
124#[cfg(feature = "alloc")]
125pub trait TableGenerator<T: Table>:
126    fmt::Debug + Default + Clone + Send + Sync + Sized + 'static
127{
128    /// Create proofs with 32 bytes seed.
129    ///
130    /// There is also `Self::create_proofs_parallel()` that can achieve higher performance and
131    /// lower latency at the cost of lower CPU efficiency and higher memory usage.
132    fn create_proofs(&self, seed: &PosSeed) -> Box<PosProofs>;
133
134    /// Almost the same as [`Self::create_proofs()`], but uses parallelism internally for better
135    /// performance and lower latency at the cost of lower CPU efficiency and higher memory usage
136    #[cfg(feature = "parallel")]
137    fn create_proofs_parallel(&self, seed: &PosSeed) -> Box<PosProofs> {
138        self.create_proofs(seed)
139    }
140}
141
142/// Proof of space kind
143pub trait Table: SolutionPotVerifier + Sized + Send + Sync + 'static {
144    /// Proof of space table type
145    const TABLE_TYPE: PosTableType;
146    /// Instance that can be used to generate tables with better performance
147    #[cfg(feature = "alloc")]
148    type Generator: TableGenerator<Self>;
149
150    /// Check whether proof created earlier is valid
151    fn is_proof_valid(seed: &PosSeed, s_bucket: SBucket, proof: &PosProof) -> bool;
152
153    /// Returns a stateful table generator with better performance
154    #[cfg(feature = "alloc")]
155    fn generator() -> Self::Generator {
156        Self::Generator::default()
157    }
158}