zcash_client_backend/serialization/
shardtree.rs

1//! Serialization formats for data stored as SQLite BLOBs
2
3use byteorder::{ReadBytesExt, WriteBytesExt};
4use core::ops::Deref;
5use shardtree::{Node, PrunableTree, RetentionFlags, Tree};
6use std::io::{self, Read, Write};
7use std::sync::Arc;
8use zcash_encoding::Optional;
9use zcash_primitives::merkle_tree::HashSer;
10
11const SER_V1: u8 = 1;
12
13const NIL_TAG: u8 = 0;
14const LEAF_TAG: u8 = 1;
15const PARENT_TAG: u8 = 2;
16
17/// Writes a [`PrunableTree`] to the provided [`Write`] instance.
18///
19/// This is the primary method used for ShardTree shard persistence. It writes a version identifier
20/// for the most-current serialized form, followed by the tree data.
21pub fn write_shard<H: HashSer, W: Write>(writer: &mut W, tree: &PrunableTree<H>) -> io::Result<()> {
22    fn write_inner<H: HashSer, W: Write>(
23        mut writer: &mut W,
24        tree: &PrunableTree<H>,
25    ) -> io::Result<()> {
26        match tree.deref() {
27            Node::Parent { ann, left, right } => {
28                writer.write_u8(PARENT_TAG)?;
29                Optional::write(&mut writer, ann.as_ref(), |w, h| {
30                    <H as HashSer>::write(h, w)
31                })?;
32                write_inner(writer, left)?;
33                write_inner(writer, right)?;
34                Ok(())
35            }
36            Node::Leaf { value } => {
37                writer.write_u8(LEAF_TAG)?;
38                value.0.write(&mut writer)?;
39                writer.write_u8(value.1.bits())?;
40                Ok(())
41            }
42            Node::Nil => {
43                writer.write_u8(NIL_TAG)?;
44                Ok(())
45            }
46        }
47    }
48
49    writer.write_u8(SER_V1)?;
50    write_inner(writer, tree)
51}
52
53fn read_shard_v1<H: HashSer, R: Read>(mut reader: &mut R) -> io::Result<PrunableTree<H>> {
54    match reader.read_u8()? {
55        PARENT_TAG => {
56            let ann = Optional::read(&mut reader, <H as HashSer>::read)?.map(Arc::new);
57            let left = read_shard_v1(reader)?;
58            let right = read_shard_v1(reader)?;
59            Ok(Tree::parent(ann, left, right))
60        }
61        LEAF_TAG => {
62            let value = <H as HashSer>::read(&mut reader)?;
63            let flags = reader.read_u8().and_then(|bits| {
64                RetentionFlags::from_bits(bits).ok_or_else(|| {
65                    io::Error::new(
66                        io::ErrorKind::InvalidData,
67                        format!(
68                            "Byte value {bits} does not correspond to a valid set of retention flags"
69                        ),
70                    )
71                })
72            })?;
73            Ok(Tree::leaf((value, flags)))
74        }
75        NIL_TAG => Ok(Tree::empty()),
76        other => Err(io::Error::new(
77            io::ErrorKind::InvalidData,
78            format!("Node tag not recognized: {other}"),
79        )),
80    }
81}
82
83/// Reads a [`PrunableTree`] from the provided [`Read`] instance.
84///
85/// This function operates by first parsing a 1-byte version identifier, and then dispatching to
86/// the correct deserialization function for the observed version, or returns an
87/// [`io::ErrorKind::InvalidData`] error in the case that the version is not recognized.
88pub fn read_shard<H: HashSer, R: Read>(mut reader: R) -> io::Result<PrunableTree<H>> {
89    match reader.read_u8()? {
90        SER_V1 => read_shard_v1(&mut reader),
91        other => Err(io::Error::new(
92            io::ErrorKind::InvalidData,
93            format!("Shard serialization version not recognized: {other}"),
94        )),
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use incrementalmerkletree::frontier::testing::{TestNode, arb_test_node};
101    use proptest::prelude::*;
102    use shardtree::testing::arb_prunable_tree;
103    use std::io::Cursor;
104
105    use super::{read_shard, write_shard};
106
107    proptest! {
108        #[test]
109        fn check_shard_roundtrip(
110            tree in arb_prunable_tree(arb_test_node(), 8, 32)
111        ) {
112            let mut tree_data = vec![];
113            write_shard(&mut tree_data, &tree).unwrap();
114            let cursor = Cursor::new(tree_data);
115            let tree_result = read_shard::<TestNode, _>(cursor).unwrap();
116            assert_eq!(tree, tree_result);
117        }
118    }
119}
zcash_client_backend/serialization/shardtree.rs

zcash_client_backend/serialization/
shardtree.rs
