jmt/
writer.rs

1use alloc::collections::{BTreeMap, BTreeSet};
2
3use alloc::vec::Vec;
4use anyhow::Result;
5use borsh::{BorshDeserialize, BorshSerialize};
6#[cfg(any(test))]
7use proptest_derive::Arbitrary;
8
9use crate::{
10    node_type::{Node, NodeKey},
11    types::Version,
12    KeyHash, OwnedValue,
13};
14
15/// Defines the interface used to write a batch of updates from a
16/// [`JellyfishMerkleTree`](crate::JellyfishMerkleTree)
17/// to the underlying storage holding nodes.
18pub trait TreeWriter {
19    /// Writes a node batch into storage.
20    fn write_node_batch(&self, node_batch: &NodeBatch) -> Result<()>;
21}
22
23/// Node batch that will be written into db atomically with other batches.
24#[derive(Debug, Clone, PartialEq, Default, Eq, borsh::BorshSerialize, borsh::BorshDeserialize)]
25pub struct NodeBatch {
26    nodes: BTreeMap<NodeKey, Node>,
27    values: BTreeMap<(Version, KeyHash), Option<OwnedValue>>,
28}
29
30impl NodeBatch {
31    /// Creates a new node batch
32    pub fn new(
33        nodes: BTreeMap<NodeKey, Node>,
34        values: BTreeMap<(Version, KeyHash), Option<OwnedValue>>,
35    ) -> Self {
36        NodeBatch { nodes, values }
37    }
38
39    /// Reset a NodeBatch to its empty state.
40    pub fn clear(&mut self) {
41        self.nodes.clear();
42        self.values.clear()
43    }
44
45    /// Get a node by key.
46    pub fn get_node(&self, node_key: &NodeKey) -> Option<&Node> {
47        self.nodes.get(node_key)
48    }
49
50    /// Returns a reference to the current set of nodes.
51    pub fn nodes(&self) -> &BTreeMap<NodeKey, Node> {
52        &self.nodes
53    }
54
55    /// Insert a node into the batch.
56    pub fn insert_node(&mut self, node_key: NodeKey, node: Node) -> Option<Node> {
57        self.nodes.insert(node_key, node)
58    }
59
60    /// Insert a node into the batch.
61    pub fn insert_value(&mut self, version: Version, key_hash: KeyHash, value: OwnedValue) {
62        self.values.insert((version, key_hash), Some(value));
63    }
64
65    /// Returns a reference to the current set of nodes.
66    pub fn values(&self) -> &BTreeMap<(Version, KeyHash), core::option::Option<Vec<u8>>> {
67        &self.values
68    }
69
70    /// Extend a node batch.
71    pub fn extend(
72        &mut self,
73        nodes: impl IntoIterator<Item = (NodeKey, Node)>,
74        values: impl IntoIterator<Item = ((Version, KeyHash), Option<OwnedValue>)>,
75    ) {
76        self.nodes.extend(nodes);
77        self.values.extend(values);
78    }
79
80    /// Merge two NodeBatches into a single one.
81    pub fn merge(&mut self, rhs: Self) {
82        self.extend(rhs.nodes, rhs.values)
83    }
84
85    /// Check if the node batch contains any items.
86    pub fn is_empty(&self) -> bool {
87        self.nodes.is_empty() && self.values.is_empty()
88    }
89}
90/// [`StaleNodeIndex`](struct.StaleNodeIndex.html) batch that will be written into db atomically
91/// with other batches.
92pub type StaleNodeIndexBatch = BTreeSet<StaleNodeIndex>;
93
94#[derive(Clone, Debug, Default, Eq, PartialEq, borsh::BorshSerialize, borsh::BorshDeserialize)]
95pub struct NodeStats {
96    pub new_nodes: usize,
97    pub new_leaves: usize,
98    pub stale_nodes: usize,
99    pub stale_leaves: usize,
100}
101
102/// Indicates a node becomes stale since `stale_since_version`.
103#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, BorshDeserialize, BorshSerialize)]
104#[cfg_attr(any(test), derive(Arbitrary))]
105pub struct StaleNodeIndex {
106    /// The version since when the node is overwritten and becomes stale.
107    pub stale_since_version: Version,
108    /// The [`NodeKey`](node_type/struct.NodeKey.html) identifying the node associated with this
109    /// record.
110    pub node_key: NodeKey,
111}
112
113/// This is a wrapper of [`NodeBatch`](type.NodeBatch.html),
114/// [`StaleNodeIndexBatch`](type.StaleNodeIndexBatch.html) and some stats of nodes that represents
115/// the incremental updates of a tree and pruning indices after applying a write set,
116/// which is a vector of `hashed_account_address` and `new_value` pairs.
117#[derive(Clone, Debug, Default, Eq, PartialEq, BorshSerialize, BorshDeserialize)]
118pub struct TreeUpdateBatch {
119    pub node_batch: NodeBatch,
120    pub stale_node_index_batch: StaleNodeIndexBatch,
121    pub node_stats: Vec<NodeStats>,
122}
jmt/writer.rs

jmt/
writer.rs
