chia_datalayer/merkle/
format.rs

1use crate::merkle::error::Error;
2use crate::{Side, block_range, internal_hash};
3use chia_protocol::Bytes32;
4#[cfg(feature = "py-bindings")]
5use chia_py_streamable_macro::{PyJsonDict, PyStreamable};
6use chia_streamable_macro::Streamable;
7use chia_traits::Streamable;
8#[cfg(feature = "py-bindings")]
9use pyo3::{Bound, FromPyObject, IntoPyObject, PyAny, PyErr, Python, pyclass, pymethods};
10use std::ops::Range;
11
12pub type TreeIndexType = u32;
13
14#[cfg_attr(
15    feature = "py-bindings",
16    pyclass(from_py_object),
17    derive(PyJsonDict, PyStreamable)
18)]
19#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Streamable)]
20// ISSUE: this cfg()/cfg(not()) is terrible, but there's an issue with pyo3
21//        being found with a cfg_attr
22//        https://github.com/PyO3/pyo3/issues/5125
23#[cfg(feature = "py-bindings")]
24pub struct TreeIndex(#[pyo3(get, name = "raw")] pub TreeIndexType);
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Streamable)]
27#[cfg(not(feature = "py-bindings"))]
28pub struct TreeIndex(pub TreeIndexType);
29
30#[cfg(feature = "py-bindings")]
31#[pymethods]
32impl TreeIndex {
33    #[new]
34    pub fn py_new(raw: TreeIndexType) -> Self {
35        Self(raw)
36    }
37}
38
39impl std::fmt::Display for TreeIndex {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        self.0.fmt(f)
42    }
43}
44
45#[cfg_attr(
46    feature = "py-bindings",
47    derive(FromPyObject, IntoPyObject, PyJsonDict),
48    pyo3(transparent)
49)]
50#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Streamable)]
51pub struct Parent(pub Option<TreeIndex>);
52
53#[cfg_attr(
54    feature = "py-bindings",
55    derive(FromPyObject, IntoPyObject, PyJsonDict),
56    pyo3(transparent)
57)]
58#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
59#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Streamable)]
60pub struct Hash(pub Bytes32);
61
62/// Key and value ids are provided from outside of this code and are implemented as
63/// the row id from sqlite which is a signed 8 byte integer.  The actual key and
64/// value data bytes will not be handled within this code, only outside.
65#[cfg_attr(
66    feature = "py-bindings",
67    pyclass(from_py_object),
68    derive(PyJsonDict, PyStreamable)
69)]
70#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
71#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Streamable)]
72// ISSUE: this cfg()/cfg(not()) is terrible, but there's an issue with pyo3
73//        being found with a cfg_attr
74//        https://github.com/PyO3/pyo3/issues/5125
75#[cfg(feature = "py-bindings")]
76pub struct KeyId(#[pyo3(get, name = "raw")] pub i64);
77
78#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
79#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Streamable)]
80#[cfg(not(feature = "py-bindings"))]
81pub struct KeyId(pub i64);
82
83#[cfg(feature = "py-bindings")]
84#[pymethods]
85impl KeyId {
86    #[new]
87    pub fn py_new(raw: i64) -> Self {
88        Self(raw)
89    }
90}
91
92#[cfg_attr(
93    feature = "py-bindings",
94    pyclass(from_py_object),
95    derive(PyJsonDict, PyStreamable)
96)]
97#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
98#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Streamable)]
99// ISSUE: this cfg()/cfg(not()) is terrible, but there's an issue with pyo3
100//        being found with a cfg_attr
101//        https://github.com/PyO3/pyo3/issues/5125
102#[cfg(feature = "py-bindings")]
103pub struct ValueId(#[pyo3(get, name = "raw")] pub i64);
104
105#[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))]
106#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Streamable)]
107#[cfg(not(feature = "py-bindings"))]
108pub struct ValueId(pub i64);
109
110#[cfg(feature = "py-bindings")]
111#[pymethods]
112impl ValueId {
113    #[new]
114    pub fn py_new(raw: i64) -> Self {
115        Self(raw)
116    }
117}
118
119impl std::fmt::Display for ValueId {
120    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
121        self.0.fmt(f)
122    }
123}
124
125impl std::fmt::Display for KeyId {
126    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
127        self.0.fmt(f)
128    }
129}
130
131// define the serialized block format
132const METADATA_RANGE: Range<usize> = 0..METADATA_SIZE;
133pub const METADATA_SIZE: usize = 2;
134// TODO: figure out the real max better than trial and error?
135pub const DATA_SIZE: usize = 53;
136pub const BLOCK_SIZE: usize = METADATA_SIZE + DATA_SIZE;
137
138pub type BlockBytes = [u8; BLOCK_SIZE];
139type MetadataBytes = [u8; METADATA_SIZE];
140type DataBytes = [u8; DATA_SIZE];
141
142const DATA_RANGE: Range<usize> = METADATA_SIZE..METADATA_SIZE + DATA_SIZE;
143
144pub(crate) fn streamable_from_bytes_ignore_extra_bytes<T>(
145    bytes: &[u8],
146) -> Result<T, chia_traits::chia_error::Error>
147where
148    T: Streamable,
149{
150    let mut cursor = std::io::Cursor::new(bytes);
151    T::parse::<false>(&mut cursor)
152}
153
154#[repr(u8)]
155#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, Streamable)]
156pub enum NodeType {
157    Internal = 0,
158    Leaf = 1,
159}
160
161#[derive(Copy, Clone, Hash, Debug, PartialEq, Eq, Streamable)]
162pub struct NodeMetadata {
163    // OPT: could save 1-2% of tree space by packing (and maybe don't do that)
164    pub node_type: NodeType,
165    pub dirty: bool,
166}
167
168#[cfg_attr(
169    feature = "py-bindings",
170    pyclass(get_all, from_py_object),
171    derive(PyJsonDict, PyStreamable)
172)]
173#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, Streamable)]
174pub struct InternalNode {
175    pub hash: Hash,
176    pub parent: Parent,
177    pub left: TreeIndex,
178    pub right: TreeIndex,
179}
180
181impl InternalNode {
182    pub fn sibling_index(&self, index: TreeIndex) -> Result<TreeIndex, Error> {
183        if index == self.right {
184            Ok(self.left)
185        } else if index == self.left {
186            Ok(self.right)
187        } else {
188            Err(Error::IndexIsNotAChild(index))
189        }
190    }
191
192    pub fn get_sibling_side(&self, index: TreeIndex) -> Result<Side, Error> {
193        if self.left == index {
194            Ok(Side::Right)
195        } else if self.right == index {
196            Ok(Side::Left)
197        } else {
198            Err(Error::IndexIsNotAChild(index))
199        }
200    }
201}
202
203#[cfg_attr(
204    feature = "py-bindings",
205    pyclass(get_all, from_py_object),
206    derive(PyJsonDict, PyStreamable)
207)]
208#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, Streamable)]
209pub struct LeafNode {
210    pub hash: Hash,
211    pub parent: Parent,
212    pub key: KeyId,
213    pub value: ValueId,
214}
215
216// TODO: consider forcing ::new() with validity checks
217#[derive(Clone, Copy, Debug, PartialEq, Eq)]
218pub enum Node {
219    Internal(InternalNode),
220    Leaf(LeafNode),
221}
222
223impl Node {
224    pub fn parent(&self) -> Parent {
225        match self {
226            Node::Internal(node) => node.parent,
227            Node::Leaf(node) => node.parent,
228        }
229    }
230
231    pub fn set_parent(&mut self, parent: Parent) {
232        match self {
233            Node::Internal(node) => node.parent = parent,
234            Node::Leaf(node) => node.parent = parent,
235        }
236    }
237
238    pub fn hash(&self) -> Hash {
239        match self {
240            Node::Internal(node) => node.hash,
241            Node::Leaf(node) => node.hash,
242        }
243    }
244
245    pub fn set_hash(&mut self, hash: Hash) {
246        match self {
247            Node::Internal(node) => node.hash = hash,
248            Node::Leaf(node) => node.hash = hash,
249        }
250    }
251
252    pub fn from_bytes(
253        metadata: &NodeMetadata,
254        blob: &DataBytes,
255    ) -> Result<Self, chia_traits::chia_error::Error> {
256        Ok(match metadata.node_type {
257            NodeType::Internal => Node::Internal(streamable_from_bytes_ignore_extra_bytes(blob)?),
258            NodeType::Leaf => Node::Leaf(streamable_from_bytes_ignore_extra_bytes(blob)?),
259        })
260    }
261
262    pub fn to_bytes(&self) -> Result<DataBytes, Error> {
263        let mut base = match self {
264            Node::Internal(node) => node.to_bytes(),
265            Node::Leaf(node) => node.to_bytes(),
266        }
267        .map_err(Error::Streaming)?;
268        assert!(base.len() <= DATA_SIZE);
269        base.resize(DATA_SIZE, 0);
270        Ok(base
271            .as_slice()
272            .try_into()
273            .expect("padding was added above, might be too large"))
274    }
275
276    pub fn expect_leaf(&self, message: &str) -> LeafNode {
277        let Node::Leaf(leaf) = self else {
278            let message = message.replace("<<self>>", &format!("{self:?}"));
279            panic!("{}", message)
280        };
281
282        *leaf
283    }
284
285    pub fn expect_internal(&self, message: &str) -> InternalNode {
286        let Node::Internal(internal) = self else {
287            let message = message.replace("<<self>>", &format!("{self:?}"));
288            panic!("{}", message)
289        };
290
291        *internal
292    }
293
294    pub fn try_into_leaf(self) -> Result<LeafNode, Error> {
295        match self {
296            Node::Leaf(leaf) => Ok(leaf),
297            Node::Internal(internal) => Err(Error::NodeNotALeaf(internal)),
298        }
299    }
300}
301
302#[cfg(feature = "py-bindings")]
303impl<'py> IntoPyObject<'py> for Node {
304    type Target = PyAny;
305    type Output = Bound<'py, Self::Target>;
306    type Error = PyErr;
307
308    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
309        match self {
310            Node::Internal(node) => Ok(node.into_pyobject(py)?.into_any()),
311            Node::Leaf(node) => Ok(node.into_pyobject(py)?.into_any()),
312        }
313    }
314}
315
316// TODO: consider forcing ::new() with validity checks
317#[derive(Clone, Copy, Debug, PartialEq)]
318pub struct Block {
319    // NOTE: metadata node type and node's type not verified for agreement
320    pub metadata: NodeMetadata,
321    pub node: Node,
322}
323
324impl Block {
325    pub fn to_bytes(&self) -> Result<BlockBytes, Error> {
326        let mut blob: BlockBytes = [0; BLOCK_SIZE];
327        blob[METADATA_RANGE].copy_from_slice(&self.metadata.to_bytes().map_err(Error::Streaming)?);
328        blob[DATA_RANGE].copy_from_slice(&self.node.to_bytes()?);
329
330        Ok(blob)
331    }
332
333    pub fn from_bytes(blob: BlockBytes) -> Result<Self, Error> {
334        let metadata_blob: MetadataBytes = blob[METADATA_RANGE].try_into().unwrap();
335        let data_blob: DataBytes = blob[DATA_RANGE].try_into().unwrap();
336        let metadata =
337            NodeMetadata::from_bytes(&metadata_blob).map_err(Error::FailedLoadingMetadata)?;
338        let node = Node::from_bytes(&metadata, &data_blob).map_err(Error::FailedLoadingNode)?;
339
340        Ok(Block { metadata, node })
341    }
342
343    pub fn update_hash(&mut self, left: &Hash, right: &Hash) {
344        self.node.set_hash(internal_hash(left, right));
345        self.metadata.dirty = false;
346    }
347}
348
349pub fn try_get_block(blob: &[u8], index: TreeIndex) -> Result<Block, Error> {
350    let range = block_range(index);
351    let block_bytes: BlockBytes = blob
352        .get(range)
353        .ok_or(Error::BlockIndexOutOfBounds(index))?
354        .try_into()
355        .expect("used block_range() so should be correct length");
356
357    Block::from_bytes(block_bytes)
358}
chia_datalayer/merkle/format.rs

chia_datalayer/merkle/
format.rs
