chik_protocol/
proof_of_space.rs

1use crate::bytes::{Bytes, Bytes32};
2use chik_bls::G1Element;
3use chik_streamable_macro::streamable;
4
5#[streamable(no_json)]
6pub struct ProofOfSpace {
7    challenge: Bytes32,
8    pool_public_key: Option<G1Element>,
9    pool_contract_puzzle_hash: Option<Bytes32>,
10    plot_public_key: G1Element,
11    // this field was renamed when adding support for v2 plots since the top
12    // bit now means whether it's v1 or v2. To stay backwards compabible with
13    // JSON serialization, we still serialize this as its original name
14    #[cfg_attr(feature = "serde", serde(rename = "size", alias = "version_and_size"))]
15    version_and_size: u8,
16    proof: Bytes,
17}
18
19#[derive(Debug, PartialEq)]
20pub enum PlotSize {
21    V1(u8),
22    V2(u8),
23}
24
25impl ProofOfSpace {
26    pub fn size(&self) -> PlotSize {
27        if (self.version_and_size & 0x80) == 0 {
28            PlotSize::V1(self.version_and_size)
29        } else {
30            PlotSize::V2(self.version_and_size & 0x7f)
31        }
32    }
33}
34
35#[cfg(feature = "py-bindings")]
36use chik_traits::{FromJsonDict, ToJsonDict};
37#[cfg(feature = "py-bindings")]
38use pyo3::prelude::*;
39
40#[cfg(feature = "py-bindings")]
41#[pyclass(name = "PlotSize")]
42pub struct PyPlotSize {
43    #[pyo3(get)]
44    pub size_v1: Option<u8>,
45    #[pyo3(get)]
46    pub size_v2: Option<u8>,
47}
48
49#[cfg(feature = "py-bindings")]
50#[pymethods]
51impl PyPlotSize {
52    #[staticmethod]
53    fn make_v1(s: u8) -> Self {
54        Self {
55            size_v1: Some(s),
56            size_v2: None,
57        }
58    }
59
60    #[staticmethod]
61    fn make_v2(s: u8) -> Self {
62        Self {
63            size_v1: None,
64            size_v2: Some(s),
65        }
66    }
67}
68
69#[cfg(feature = "py-bindings")]
70#[pymethods]
71impl ProofOfSpace {
72    #[pyo3(name = "size")]
73    fn py_size(&self) -> PyPlotSize {
74        match self.size() {
75            PlotSize::V1(s) => PyPlotSize {
76                size_v1: Some(s),
77                size_v2: None,
78            },
79            PlotSize::V2(s) => PyPlotSize {
80                size_v1: None,
81                size_v2: Some(s),
82            },
83        }
84    }
85}
86
87#[cfg(feature = "py-bindings")]
88impl ToJsonDict for ProofOfSpace {
89    fn to_json_dict(&self, py: pyo3::Python<'_>) -> pyo3::PyResult<pyo3::PyObject> {
90        use pyo3::prelude::PyDictMethods;
91        let ret = pyo3::types::PyDict::new(py);
92
93        ret.set_item("challenge", self.challenge.to_json_dict(py)?)?;
94        ret.set_item("pool_public_key", self.pool_public_key.to_json_dict(py)?)?;
95        ret.set_item(
96            "pool_contract_puzzle_hash",
97            self.pool_contract_puzzle_hash.to_json_dict(py)?,
98        )?;
99        ret.set_item("plot_public_key", self.plot_public_key.to_json_dict(py)?)?;
100
101        // "size" was the original name of this field. We keep it to remain backwards compatible
102        ret.set_item("size", self.version_and_size.to_json_dict(py)?)?;
103        ret.set_item("proof", self.proof.to_json_dict(py)?)?;
104
105        Ok(ret.into())
106    }
107}
108
109#[cfg(feature = "py-bindings")]
110impl FromJsonDict for ProofOfSpace {
111    fn from_json_dict(o: &pyo3::Bound<'_, pyo3::PyAny>) -> pyo3::PyResult<Self> {
112        use pyo3::prelude::PyAnyMethods;
113        Ok(Self {
114            challenge: <Bytes32 as FromJsonDict>::from_json_dict(&o.get_item("challenge")?)?,
115            pool_public_key: <Option<G1Element> as FromJsonDict>::from_json_dict(
116                &o.get_item("pool_public_key")?,
117            )?,
118            pool_contract_puzzle_hash: <Option<Bytes32> as FromJsonDict>::from_json_dict(
119                &o.get_item("pool_contract_puzzle_hash")?,
120            )?,
121            plot_public_key: <G1Element as FromJsonDict>::from_json_dict(
122                &o.get_item("plot_public_key")?,
123            )?,
124            version_and_size: <u8 as FromJsonDict>::from_json_dict(&o.get_item("size")?)?,
125            proof: <Bytes as FromJsonDict>::from_json_dict(&o.get_item("proof")?)?,
126        })
127    }
128}
129
130#[cfg(test)]
131#[allow(clippy::needless_pass_by_value)]
132mod tests {
133    use super::*;
134    use rstest::rstest;
135
136    #[rstest]
137    #[case(0x00, PlotSize::V1(0))]
138    #[case(0x01, PlotSize::V1(1))]
139    #[case(0x08, PlotSize::V1(8))]
140    #[case(0x7f, PlotSize::V1(0x7f))]
141    #[case(0x80, PlotSize::V2(0))]
142    #[case(0x81, PlotSize::V2(1))]
143    #[case(0x80 + 28, PlotSize::V2(28))]
144    #[case(0x80 + 30, PlotSize::V2(30))]
145    #[case(0x80 + 32, PlotSize::V2(32))]
146    #[case(0xff, PlotSize::V2(0x7f))]
147    fn proof_of_space_size(#[case] size_field: u8, #[case] expect: PlotSize) {
148        let pos = ProofOfSpace::new(
149            Bytes32::from(b"abababababababababababababababab"),
150            None,
151            None,
152            G1Element::default(),
153            size_field,
154            Bytes::from(vec![]),
155        );
156
157        assert_eq!(pos.size(), expect);
158    }
159}