From f4c3761a9c8b987ddb5697e184c0de70f4bc4a14 Mon Sep 17 00:00:00 2001 From: Rigidity Date: Wed, 26 Jun 2024 09:38:47 -0400 Subject: [PATCH 1/2] Derive ToClvm/FromClvm for Coin --- crates/chia-protocol/src/coin.rs | 28 ++++------------------------ 1 file changed, 4 insertions(+), 24 deletions(-) diff --git a/crates/chia-protocol/src/coin.rs b/crates/chia-protocol/src/coin.rs index 42e6a4074..afa0d2aaf 100644 --- a/crates/chia-protocol/src/coin.rs +++ b/crates/chia-protocol/src/coin.rs @@ -1,16 +1,14 @@ -use crate::{Bytes32, BytesImpl}; +use crate::Bytes32; use chia_streamable_macro::streamable; -use clvm_traits::{ - clvm_list, destructure_list, match_list, ClvmDecoder, ClvmEncoder, FromClvm, FromClvmError, - ToClvm, ToClvmError, -}; +use clvm_traits::{FromClvm, ToClvm}; use sha2::{Digest, Sha256}; #[cfg(feature = "py-bindings")] use pyo3::prelude::*; #[streamable] -#[derive(Copy)] +#[derive(Copy, ToClvm, FromClvm)] +#[clvm(list)] pub struct Coin { parent_coin_info: Bytes32, puzzle_hash: Bytes32, @@ -55,24 +53,6 @@ impl Coin { } } -impl ToClvm for Coin { - fn to_clvm(&self, encoder: &mut impl ClvmEncoder) -> Result { - clvm_list!(self.parent_coin_info, self.puzzle_hash, self.amount).to_clvm(encoder) - } -} - -impl FromClvm for Coin { - fn from_clvm(decoder: &impl ClvmDecoder, node: N) -> Result { - let destructure_list!(parent_coin_info, puzzle_hash, amount) = - , BytesImpl<32>, u64)>::from_clvm(decoder, node)?; - Ok(Coin { - parent_coin_info, - puzzle_hash, - amount, - }) - } -} - #[cfg(test)] mod tests { use super::*; From 1b47b3b6aa1b472e33b7110fe7ded35a40c40033 Mon Sep 17 00:00:00 2001 From: Rigidity Date: Wed, 26 Jun 2024 10:40:03 -0400 Subject: [PATCH 2/2] Python-only version of WeightProof --- crates/chia_py_streamable_macro/src/lib.rs | 6 + wheel/src/api.rs | 5 +- wheel/src/lib.rs | 1 + wheel/src/weight_proof.rs | 241 +++++++++++++++++++++ 4 files changed, 251 insertions(+), 2 deletions(-) create mode 100644 wheel/src/weight_proof.rs diff --git a/crates/chia_py_streamable_macro/src/lib.rs b/crates/chia_py_streamable_macro/src/lib.rs index 7c79a4839..77c2df9ec 100644 --- a/crates/chia_py_streamable_macro/src/lib.rs +++ b/crates/chia_py_streamable_macro/src/lib.rs @@ -101,6 +101,12 @@ pub fn py_streamable_macro(input: proc_macro::TokenStream) -> proc_macro::TokenS Ok(pyo3::IntoPy::into_py(self.clone(), py).into_bound(py)) } } + + impl pyo3::ToPyObject for #ident { + fn to_object(&self, py: pyo3::Python) -> pyo3::PyObject { + pyo3::Bound::new(py, self.clone()).unwrap().to_object(py) + } + } }; let mut fnames = Vec::::new(); diff --git a/wheel/src/api.rs b/wheel/src/api.rs index c83da5ef4..4d09bb201 100644 --- a/wheel/src/api.rs +++ b/wheel/src/api.rs @@ -1,4 +1,5 @@ use crate::run_generator::{run_block_generator, run_block_generator2}; +use crate::weight_proof::PyWeightProof; use chia_consensus::allocator::make_allocator; use chia_consensus::consensus_constants::ConsensusConstants; use chia_consensus::gen::conditions::MempoolVisitor; @@ -37,7 +38,7 @@ use chia_protocol::{ RespondUnfinishedBlock, RewardChainBlock, RewardChainBlockUnfinished, RewardChainSubSlot, SendTransaction, SpendBundle, SubEpochChallengeSegment, SubEpochData, SubEpochSegments, SubEpochSummary, SubSlotData, SubSlotProofs, TimestampedPeerInfo, TransactionAck, - TransactionsInfo, UnfinishedBlock, UnfinishedHeaderBlock, VDFInfo, VDFProof, WeightProof, + TransactionsInfo, UnfinishedBlock, UnfinishedHeaderBlock, VDFInfo, VDFProof, }; use clvm_utils::tree_hash_from_bytes; use clvmr::{ENABLE_BLS_OPS_OUTSIDE_GUARD, ENABLE_FIXED_DIV, LIMIT_HEAP, NO_UNKNOWN_OPS}; @@ -437,7 +438,7 @@ pub fn chia_rs(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; - m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/wheel/src/lib.rs b/wheel/src/lib.rs index 3b42df578..671ecef6d 100644 --- a/wheel/src/lib.rs +++ b/wheel/src/lib.rs @@ -4,3 +4,4 @@ mod adapt_response; mod api; mod run_generator; mod run_program; +mod weight_proof; diff --git a/wheel/src/weight_proof.rs b/wheel/src/weight_proof.rs new file mode 100644 index 000000000..4925da8db --- /dev/null +++ b/wheel/src/weight_proof.rs @@ -0,0 +1,241 @@ +use std::hash::{DefaultHasher, Hash, Hasher}; +use std::io::Cursor; + +use chia_protocol::{HeaderBlock, SubEpochChallengeSegment, SubEpochData, WeightProof}; +use chia_traits::{FromJsonDict, Streamable, ToJsonDict}; +use pyo3::buffer::PyBuffer; +use pyo3::exceptions::PyKeyError; +use pyo3::prelude::*; +use pyo3::types::{PyBytes, PyDict, PyList}; +use sha2::{Digest, Sha256}; + +#[pyclass(frozen, get_all, name = "WeightProof")] +#[derive(Debug, Clone)] +pub struct PyWeightProof { + pub sub_epochs: Py, + pub sub_epoch_segments: Py, + pub recent_chain_data: Py, +} + +impl Streamable for PyWeightProof { + fn parse(input: &mut std::io::Cursor<&[u8]>) -> chia_traits::Result + where + Self: Sized, + { + let wp = WeightProof::parse::(input)?; + Ok(Python::with_gil(|py| { + Self::py_new( + py, + wp.sub_epochs, + wp.sub_epoch_segments, + wp.recent_chain_data, + ) + })) + } + + fn update_digest(&self, digest: &mut sha2::Sha256) { + Python::with_gil(|py| { + let wp = self.to_rust(py).unwrap(); + wp.update_digest(digest); + }); + } + + fn stream(&self, out: &mut Vec) -> chia_traits::Result<()> { + Python::with_gil(|py| { + let wp = self.to_rust(py).unwrap(); + wp.stream(out) + }) + } +} + +impl PyWeightProof { + pub fn to_rust(&self, py: Python<'_>) -> PyResult { + let list = self.sub_epochs.bind(py); + let mut sub_epochs = Vec::with_capacity(list.len()); + for item in list.iter() { + sub_epochs.push(item.extract::()?); + } + + let list = self.sub_epoch_segments.bind(py); + let mut sub_epoch_segments = Vec::with_capacity(list.len()); + for item in list.iter() { + sub_epoch_segments.push(item.extract::()?); + } + + let list = self.recent_chain_data.bind(py); + let mut recent_chain_data = Vec::with_capacity(list.len()); + for item in list.iter() { + recent_chain_data.push(item.extract::()?); + } + + Ok(WeightProof { + sub_epochs, + sub_epoch_segments, + recent_chain_data, + }) + } +} + +#[pymethods] +impl PyWeightProof { + #[new] + #[pyo3(signature = (sub_epochs, sub_epoch_segments, recent_chain_data))] + pub fn py_new( + py: Python<'_>, + sub_epochs: Vec, + sub_epoch_segments: Vec, + recent_chain_data: Vec, + ) -> Self { + let sub_epochs = PyList::new_bound(py, sub_epochs).into(); + let sub_epoch_segments = PyList::new_bound(py, sub_epoch_segments).into(); + let recent_chain_data = PyList::new_bound(py, recent_chain_data).into(); + + Self { + sub_epochs, + sub_epoch_segments, + recent_chain_data, + } + } + + #[staticmethod] + #[pyo3(signature=(json_dict))] + pub fn from_json_dict(json_dict: &Bound<'_, PyAny>) -> PyResult { + ::from_json_dict(json_dict) + } + + pub fn to_json_dict(&self, py: Python<'_>) -> PyResult { + ToJsonDict::to_json_dict(self, py) + } + + #[pyo3(signature = (**kwargs))] + fn replace(&self, kwargs: Option<&Bound<'_, PyDict>>) -> PyResult { + let mut ret = self.clone(); + if let Some(kwargs) = kwargs { + let iter = kwargs.iter(); + for (field, value) in iter { + let field = field.extract::()?; + match field.as_str() { + "sub_epochs" => ret.sub_epochs = value.extract()?, + "sub_epoch_segments" => ret.sub_epoch_segments = value.extract()?, + "recent_chain_data" => ret.recent_chain_data = value.extract()?, + _ => { + return Err(PyKeyError::new_err(format!("unknown field {field}"))); + } + } + } + } + Ok(ret) + } + + fn __repr__(&self) -> String { + format!("{self:?}") + } + + fn __hash__(&self, py: Python<'_>) -> PyResult { + let mut hasher = DefaultHasher::new(); + Hash::hash(&self.to_rust(py)?, &mut hasher); + Ok(Hasher::finish(&hasher)) + } + + #[staticmethod] + #[pyo3(name = "from_bytes")] + pub fn py_from_bytes(blob: PyBuffer) -> PyResult { + assert!( + blob.is_c_contiguous(), + "from_bytes() must be called with a contiguous buffer" + ); + let slice = + unsafe { std::slice::from_raw_parts(blob.buf_ptr() as *const u8, blob.len_bytes()) }; + Self::from_bytes(slice).map_err(PyErr::from) + } + + #[staticmethod] + #[pyo3(name = "from_bytes_unchecked")] + pub fn py_from_bytes_unchecked(blob: PyBuffer) -> PyResult { + assert!( + blob.is_c_contiguous(), + "from_bytes_unchecked() must be called with a contiguous buffer" + ); + let slice = + unsafe { std::slice::from_raw_parts(blob.buf_ptr() as *const u8, blob.len_bytes()) }; + Self::from_bytes_unchecked(slice).map_err(PyErr::from) + } + + // returns the type as well as the number of bytes read from the buffer + #[staticmethod] + #[pyo3(signature= (blob, trusted=false))] + pub fn parse_rust(blob: PyBuffer, trusted: bool) -> PyResult<(Self, u32)> { + assert!( + blob.is_c_contiguous(), + "parse_rust() must be called with a contiguous buffer" + ); + let slice = + unsafe { std::slice::from_raw_parts(blob.buf_ptr() as *const u8, blob.len_bytes()) }; + let mut input = Cursor::<&[u8]>::new(slice); + if trusted { + Self::parse::(&mut input) + .map_err(PyErr::from) + .map(|v| (v, input.position() as u32)) + } else { + Self::parse::(&mut input) + .map_err(PyErr::from) + .map(|v| (v, input.position() as u32)) + } + } + + pub fn get_hash<'p>(&self, py: Python<'p>) -> Bound<'p, PyBytes> { + let mut ctx = Sha256::new(); + Streamable::update_digest(self, &mut ctx); + PyBytes::new_bound(py, ctx.finalize().as_slice()) + } + + #[pyo3(name = "to_bytes")] + pub fn py_to_bytes<'p>(&self, py: Python<'p>) -> PyResult> { + let mut writer = Vec::::new(); + Streamable::stream(self, &mut writer).map_err(PyErr::from)?; + Ok(PyBytes::new_bound(py, &writer)) + } + + pub fn stream_to_bytes<'p>(&self, py: Python<'p>) -> PyResult> { + self.py_to_bytes(py) + } + + pub fn __bytes__<'p>(&self, py: Python<'p>) -> PyResult> { + self.py_to_bytes(py) + } + + pub fn __deepcopy__(&self, _memo: Bound<'_, PyAny>) -> Self { + self.clone() + } + + pub fn __copy__(&self) -> Self { + self.clone() + } +} + +impl ToJsonDict for PyWeightProof { + fn to_json_dict(&self, py: Python<'_>) -> PyResult { + let wp = self.to_rust(py)?; + let ret = PyDict::new_bound(py); + ret.set_item("sub_epochs", wp.sub_epochs.to_json_dict(py)?)?; + ret.set_item( + "sub_epoch_segments", + wp.sub_epoch_segments.to_json_dict(py)?, + )?; + ret.set_item("recent_chain_data", wp.recent_chain_data.to_json_dict(py)?)?; + Ok(ret.into()) + } +} + +impl FromJsonDict for PyWeightProof { + fn from_json_dict(o: &Bound<'_, PyAny>) -> pyo3::PyResult { + Python::with_gil(|py| { + Ok(Self::py_new( + py, + FromJsonDict::from_json_dict(&o.get_item("sub_epochs")?)?, + FromJsonDict::from_json_dict(&o.get_item("sub_epoch_segments")?)?, + FromJsonDict::from_json_dict(&o.get_item("recent_chain_data")?)?, + )) + }) + } +}