diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 8601b60..2617787 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -38,6 +38,9 @@ pub use store::MerkleStore; mod node; pub use node::InnerNodeInfo; +mod partial_mt; +pub use partial_mt::PartialMerkleTree; + // ERRORS // ================================================================================================ diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs new file mode 100644 index 0000000..3807de3 --- /dev/null +++ b/src/merkle/partial_mt/mod.rs @@ -0,0 +1,329 @@ +use super::{ + BTreeMap, BTreeSet, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, ValuePath, Vec, + Word, EMPTY_WORD, +}; +use crate::utils::{format, string::String, word_to_hex}; +use core::fmt; + +#[cfg(test)] +mod tests; + +// CONSTANTS +// ================================================================================================ + +/// Index of the root node. +const ROOT_INDEX: NodeIndex = NodeIndex::root(); + +/// An RpoDigest consisting of 4 ZERO elements. +const EMPTY_DIGEST: RpoDigest = RpoDigest::new(EMPTY_WORD); + +// PARTIAL MERKLE TREE +// ================================================================================================ + +/// A partial Merkle tree with NodeIndex keys and 4-element RpoDigest leaf values. Partial Merkle +/// Tree allows to create Merkle Tree by providing Merkle paths of different lengths. +/// +/// The root of the tree is recomputed on each new leaf update. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PartialMerkleTree { + max_depth: u8, + nodes: BTreeMap, + leaves: BTreeSet, +} + +impl Default for PartialMerkleTree { + fn default() -> Self { + Self::new() + } +} + +impl PartialMerkleTree { + // CONSTANTS + // -------------------------------------------------------------------------------------------- + + /// Minimum supported depth. + pub const MIN_DEPTH: u8 = 1; + + /// Maximum supported depth. + pub const MAX_DEPTH: u8 = 64; + + // CONSTRUCTORS + // -------------------------------------------------------------------------------------------- + + /// Returns a new emply [PartialMerkleTree]. + pub fn new() -> Self { + PartialMerkleTree { + max_depth: 0, + nodes: BTreeMap::new(), + leaves: BTreeSet::new(), + } + } + + /// Appends the provided paths iterator into the set. + /// + /// Analogous to [Self::add_path]. + pub fn with_paths(paths: I) -> Result + where + I: IntoIterator, + { + // create an empty tree + let tree = PartialMerkleTree::new(); + + paths.into_iter().try_fold(tree, |mut tree, (index, value, path)| { + tree.add_path(index, value, path)?; + Ok(tree) + }) + } + + // PUBLIC ACCESSORS + // -------------------------------------------------------------------------------------------- + + /// Returns the root of this Merkle tree. + pub fn root(&self) -> RpoDigest { + self.nodes.get(&ROOT_INDEX).cloned().unwrap_or(EMPTY_DIGEST) + } + + /// Returns the depth of this Merkle tree. + pub fn max_depth(&self) -> u8 { + self.max_depth + } + + /// Returns a node at the specified NodeIndex. + /// + /// # Errors + /// Returns an error if the specified NodeIndex is not contained in the nodes map. + pub fn get_node(&self, index: NodeIndex) -> Result { + self.nodes.get(&index).ok_or(MerkleError::NodeNotInSet(index)).map(|hash| *hash) + } + + /// Returns true if provided index contains in the leaves set, false otherwise. + pub fn is_leaf(&self, index: NodeIndex) -> bool { + self.leaves.contains(&index) + } + + /// Returns a vector of paths from every leaf to the root. + pub fn paths(&self) -> Vec<(NodeIndex, ValuePath)> { + let mut paths = Vec::new(); + self.leaves.iter().for_each(|&leaf| { + paths.push(( + leaf, + ValuePath { + value: *self.get_node(leaf).expect("Failed to get leaf node"), + path: self.get_path(leaf).expect("Failed to get path"), + }, + )); + }); + paths + } + + /// Returns a Merkle path from the node at the specified index to the root. + /// + /// The node itself is not included in the path. + /// + /// # Errors + /// Returns an error if: + /// - the specified index has depth set to 0 or the depth is greater than the depth of this + /// Merkle tree. + /// - the specified index is not contained in the nodes map. + pub fn get_path(&self, mut index: NodeIndex) -> Result { + if index.is_root() { + return Err(MerkleError::DepthTooSmall(index.depth())); + } else if index.depth() > self.max_depth() { + return Err(MerkleError::DepthTooBig(index.depth() as u64)); + } + + if !self.nodes.contains_key(&index) { + return Err(MerkleError::NodeNotInSet(index)); + } + + let mut path = Vec::new(); + for _ in 0..index.depth() { + let sibling_index = index.sibling(); + index.move_up(); + let sibling = + self.nodes.get(&sibling_index).cloned().expect("Sibling node not in the map"); + path.push(Word::from(sibling)); + } + Ok(MerklePath::new(path)) + } + + // ITERATORS + // -------------------------------------------------------------------------------------------- + + /// Returns an iterator over the leaves of this [PartialMerkleTree]. + pub fn leaves(&self) -> impl Iterator + '_ { + self.leaves.iter().map(|&leaf| { + ( + leaf, + self.get_node(leaf).expect(&format!("Leaf with {leaf} is not in the nodes map")), + ) + }) + } + + // STATE MUTATORS + // -------------------------------------------------------------------------------------------- + + /// Adds the nodes of the specified Merkle path to this [PartialMerkleTree]. The `index_value` + /// and `value` parameters specify the leaf node at which the path starts. + /// + /// # Errors + /// Returns an error if: + /// - The depth of the specified node_index is greater than 64 or smaller than 1. + /// - The specified path is not consistent with other paths in the set (i.e., resolves to a + /// different root). + pub fn add_path( + &mut self, + index_value: u64, + value: RpoDigest, + path: MerklePath, + ) -> Result<(), MerkleError> { + let index_value = NodeIndex::new(path.len() as u8, index_value)?; + + Self::check_depth(index_value.depth())?; + self.update_depth(index_value.depth()); + + // add provided node and its sibling to the leaves set + self.leaves.insert(index_value); + let sibling_node_index = index_value.sibling(); + self.leaves.insert(sibling_node_index); + + // add provided node and its sibling to the nodes map + self.nodes.insert(index_value, value); + self.nodes.insert(sibling_node_index, path[0].into()); + + // traverse to the root, updating the nodes + let mut index_value = index_value; + let node = Rpo256::merge(&index_value.build_node(value, path[0].into())); + let root = path.iter().skip(1).copied().fold(node, |node, hash| { + index_value.move_up(); + // insert calculated node to the nodes map + self.nodes.insert(index_value, node); + + // if the calculated node was a leaf, remove it from leaves set. + self.leaves.remove(&index_value); + + let sibling_node = index_value.sibling(); + + // Insert node from Merkle path to the nodes map. This sibling node becomes a leaf only + // if it is a new node (it wasn't in nodes map). + // Node can be in 3 states: internal node, leaf of the tree and not a tree node at all. + // - Internal node can only stay in this state -- addition of a new path can't make it + // a leaf or remove it from the tree. + // - Leaf node can stay in the same state (remain a leaf) or can become an internal + // node. In the first case we don't need to do anything, and the second case is handled + // by the call of `self.leaves.remove(&index_value);` + // - New node can be a calculated node or a "sibling" node from a Merkle Path: + // --- Calculated node, obviously, never can be a leaf. + // --- Sibling node can be only a leaf, because otherwise it is not a new node. + if self.nodes.insert(sibling_node, hash.into()).is_none() { + self.leaves.insert(sibling_node); + } + + Rpo256::merge(&index_value.build_node(node, hash.into())) + }); + + // if the path set is empty (the root is all ZEROs), set the root to the root of the added + // path; otherwise, the root of the added path must be identical to the current root + if self.root() == EMPTY_DIGEST { + self.nodes.insert(ROOT_INDEX, root); + } else if self.root() != root { + return Err(MerkleError::ConflictingRoots([*self.root(), *root].to_vec())); + } + + Ok(()) + } + + /// Updates value of the leaf at the specified index returning the old leaf value. + /// + /// This also recomputes all hashes between the leaf and the root, updating the root itself. + /// + /// # Errors + /// Returns an error if: + /// - The depth of the specified node_index is greater than 64 or smaller than 1. + /// - The specified node index is not corresponding to the leaf. + pub fn update_leaf( + &mut self, + node_index: NodeIndex, + value: RpoDigest, + ) -> Result { + // check correctness of the depth and update it + Self::check_depth(node_index.depth())?; + self.update_depth(node_index.depth()); + + // insert NodeIndex to the leaves Set + self.leaves.insert(node_index); + + // add node value to the nodes Map + let old_value = self + .nodes + .insert(node_index, value) + .ok_or(MerkleError::NodeNotInSet(node_index))?; + + // if the old value and new value are the same, there is nothing to update + if value == old_value { + return Ok(old_value); + } + + let mut node_index = node_index; + let mut value = value; + for _ in 0..node_index.depth() { + let sibling = self.nodes.get(&node_index.sibling()).expect("sibling should exist"); + value = Rpo256::merge(&node_index.build_node(value, *sibling)); + node_index.move_up(); + self.nodes.insert(node_index, value); + } + + Ok(old_value) + } + + // UTILITY FUNCTIONS + // -------------------------------------------------------------------------------------------- + + /// Utility to visualize a [PartialMerkleTree] in text. + pub fn print(&self) -> Result { + let indent = " "; + let mut s = String::new(); + s.push_str("root: "); + s.push_str(&word_to_hex(&self.root())?); + s.push('\n'); + for d in 1..=self.max_depth() { + let entries = 2u64.pow(d.into()); + for i in 0..entries { + let index = NodeIndex::new(d, i).expect("The index must always be valid"); + let node = self.get_node(index); + let node = match node { + Err(_) => continue, + Ok(node) => node, + }; + + for _ in 0..d { + s.push_str(indent); + } + s.push_str(&format!("({}, {}): ", index.depth(), index.value())); + s.push_str(&word_to_hex(&node)?); + s.push('\n'); + } + } + + Ok(s) + } + + // HELPER METHODS + // -------------------------------------------------------------------------------------------- + + /// Updates depth value with the maximum of current and provided depth. + fn update_depth(&mut self, new_depth: u8) { + self.max_depth = new_depth.max(self.max_depth); + } + + /// Returns an error if the depth is 0 or is greater than 64. + fn check_depth(depth: u8) -> Result<(), MerkleError> { + // validate the range of the depth. + if depth < Self::MIN_DEPTH { + return Err(MerkleError::DepthTooSmall(depth)); + } else if Self::MAX_DEPTH < depth { + return Err(MerkleError::DepthTooBig(depth as u64)); + } + Ok(()) + } +} diff --git a/src/merkle/partial_mt/tests.rs b/src/merkle/partial_mt/tests.rs new file mode 100644 index 0000000..ba62896 --- /dev/null +++ b/src/merkle/partial_mt/tests.rs @@ -0,0 +1,313 @@ +use super::{ + super::{int_to_node, MerkleStore, MerkleTree, NodeIndex, PartialMerkleTree}, + ValuePath, Vec, Word, +}; + +// TEST DATA +// ================================================================================================ + +const NODE10: NodeIndex = NodeIndex::new_unchecked(1, 0); +const NODE11: NodeIndex = NodeIndex::new_unchecked(1, 1); + +const NODE20: NodeIndex = NodeIndex::new_unchecked(2, 0); +const NODE22: NodeIndex = NodeIndex::new_unchecked(2, 2); +const NODE23: NodeIndex = NodeIndex::new_unchecked(2, 3); + +const NODE30: NodeIndex = NodeIndex::new_unchecked(3, 0); +const NODE31: NodeIndex = NodeIndex::new_unchecked(3, 1); +const NODE32: NodeIndex = NodeIndex::new_unchecked(3, 2); +const NODE33: NodeIndex = NodeIndex::new_unchecked(3, 3); + +const VALUES8: [Word; 8] = [ + int_to_node(30), + int_to_node(31), + int_to_node(32), + int_to_node(33), + int_to_node(34), + int_to_node(35), + int_to_node(36), + int_to_node(37), +]; + +// TESTS +// ================================================================================================ + +// For the Partial Merkle Tree tests we will use parts of the Merkle Tree which full form is +// illustrated below: +// +// __________ root __________ +// / \ +// ____ 10 ____ ____ 11 ____ +// / \ / \ +// 20 21 22 23 +// / \ / \ / \ / \ +// (30) (31) (32) (33) (34) (35) (36) (37) +// +// Where node number is a concatenation of its depth and index. For example, node with +// NodeIndex(3, 5) will be labled as `35`. Leaves of the tree are shown as nodes with parenthesis +// (33). + +/// Checks that root returned by `root()` function is equal to the expected one. +#[test] +fn get_root() { + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let expected_root = mt.root(); + + let ms = MerkleStore::from(&mt); + + let path33 = ms.get_path(expected_root, NODE33).unwrap(); + + let pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap(); + + assert_eq!(pmt.root(), expected_root.into()); +} + +/// This test checks correctness of the `add_path()` and `get_path()` functions. First it creates a +/// PMT using `add_path()` by adding Merkle Paths from node 33 and node 22 to the empty PMT. Then +/// it checks that paths returned by `get_path()` function are equal to the expected ones. +#[test] +fn add_and_get_paths() { + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let expected_root = mt.root(); + + let ms = MerkleStore::from(&mt); + + let expected_path33 = ms.get_path(expected_root, NODE33).unwrap(); + let expected_path22 = ms.get_path(expected_root, NODE22).unwrap(); + + let mut pmt = PartialMerkleTree::new(); + pmt.add_path(3, expected_path33.value.into(), expected_path33.path.clone()) + .unwrap(); + pmt.add_path(2, expected_path22.value.into(), expected_path22.path.clone()) + .unwrap(); + + let path33 = pmt.get_path(NODE33).unwrap(); + let path22 = pmt.get_path(NODE22).unwrap(); + let actual_root = pmt.root(); + + assert_eq!(expected_path33.path, path33); + assert_eq!(expected_path22.path, path22); + assert_eq!(expected_root, *actual_root); +} + +/// Checks that function `get_node` used on nodes 10 and 32 returns expected values. +#[test] +fn get_node() { + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let expected_root = mt.root(); + + let ms = MerkleStore::from(&mt); + + let path33 = ms.get_path(expected_root, NODE33).unwrap(); + + let pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap(); + + assert_eq!(ms.get_node(expected_root, NODE32).unwrap(), *pmt.get_node(NODE32).unwrap()); + assert_eq!(ms.get_node(expected_root, NODE10).unwrap(), *pmt.get_node(NODE10).unwrap()); +} + +/// Updates leaves of the PMT using `update_leaf()` function and checks that new root of the tree +/// is equal to the expected one. +#[test] +fn update_leaf() { + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let root = mt.root(); + + let mut ms = MerkleStore::from(&mt); + let path33 = ms.get_path(root, NODE33).unwrap(); + + let mut pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap(); + + let new_value32 = int_to_node(132); + let expected_root = ms.set_node(root, NODE32, new_value32).unwrap().root; + + pmt.update_leaf(NODE32, new_value32.into()).unwrap(); + let actual_root = pmt.root(); + + assert_eq!(expected_root, *actual_root); + + let new_value20 = int_to_node(120); + let expected_root = ms.set_node(expected_root, NODE20, new_value20).unwrap().root; + + pmt.update_leaf(NODE20, new_value20.into()).unwrap(); + let actual_root = pmt.root(); + + assert_eq!(expected_root, *actual_root); +} + +/// Checks that paths of the PMT returned by `paths()` function are equal to the expected ones. +#[test] +fn get_paths() { + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let expected_root = mt.root(); + + let ms = MerkleStore::from(&mt); + + let path33 = ms.get_path(expected_root, NODE33).unwrap(); + let path22 = ms.get_path(expected_root, NODE22).unwrap(); + + let mut pmt = PartialMerkleTree::new(); + pmt.add_path(3, path33.value.into(), path33.path.clone()).unwrap(); + pmt.add_path(2, path22.value.into(), path22.path.clone()).unwrap(); + // After PMT creation with path33 (33; 32, 20, 11) and path22 (22; 23, 10) we will have this + // tree: + // + // ______root______ + // / \ + // ___10___ ___11___ + // / \ / \ + // (20) 21 (22) (23) + // / \ + // (32) (33) + // + // Which have leaf nodes 20, 22, 23, 32 and 33. Hence overall we will have 5 paths -- one path + // for each leaf. + + let leaves = vec![NODE20, NODE22, NODE23, NODE32, NODE33]; + let expected_paths: Vec<(NodeIndex, ValuePath)> = leaves + .iter() + .map(|&leaf| { + ( + leaf, + ValuePath { + value: mt.get_node(leaf).unwrap().into(), + path: mt.get_path(leaf).unwrap(), + }, + ) + }) + .collect(); + + let actual_paths = pmt.paths(); + + assert_eq!(expected_paths, actual_paths); +} + +// Checks correctness of leaves determination when using the `leaves()` function. +#[test] +fn leaves() { + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let expected_root = mt.root(); + + let ms = MerkleStore::from(&mt); + + let path33 = ms.get_path(expected_root, NODE33).unwrap(); + let path22 = ms.get_path(expected_root, NODE22).unwrap(); + + let mut pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap(); + // After PMT creation with path33 (33; 32, 20, 11) we will have this tree: + // + // ______root______ + // / \ + // ___10___ (11) + // / \ + // (20) 21 + // / \ + // (32) (33) + // + // Which have leaf nodes 11, 20, 32 and 33. + + let value11 = mt.get_node(NODE11).unwrap().into(); + let value20 = mt.get_node(NODE20).unwrap().into(); + let value32 = mt.get_node(NODE32).unwrap().into(); + let value33 = mt.get_node(NODE33).unwrap().into(); + + let leaves = vec![(NODE11, value11), (NODE20, value20), (NODE32, value32), (NODE33, value33)]; + + let expected_leaves = leaves.iter().map(|&tuple| tuple); + assert!(expected_leaves.eq(pmt.leaves())); + + pmt.add_path(2, path22.value.into(), path22.path).unwrap(); + // After adding the path22 (22; 23, 10) to the existing PMT we will have this tree: + // + // ______root______ + // / \ + // ___10___ ___11___ + // / \ / \ + // (20) 21 (22) (23) + // / \ + // (32) (33) + // + // Which have leaf nodes 20, 22, 23, 32 and 33. + + let value20 = mt.get_node(NODE20).unwrap().into(); + let value22 = mt.get_node(NODE22).unwrap().into(); + let value23 = mt.get_node(NODE23).unwrap().into(); + let value32 = mt.get_node(NODE32).unwrap().into(); + let value33 = mt.get_node(NODE33).unwrap().into(); + + let leaves = vec![ + (NODE20, value20), + (NODE22, value22), + (NODE23, value23), + (NODE32, value32), + (NODE33, value33), + ]; + + let expected_leaves = leaves.iter().map(|&tuple| tuple); + assert!(expected_leaves.eq(pmt.leaves())); +} + +/// Checks that addition of the path with different root will cause an error. +#[test] +fn err_add_path() { + let path33 = vec![int_to_node(1), int_to_node(2), int_to_node(3)].into(); + let path22 = vec![int_to_node(4), int_to_node(5)].into(); + + let mut pmt = PartialMerkleTree::new(); + pmt.add_path(3, int_to_node(6).into(), path33).unwrap(); + + assert!(pmt.add_path(2, int_to_node(7).into(), path22).is_err()); +} + +/// Checks that the request of the node which is not in the PMT will cause an error. +#[test] +fn err_get_node() { + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let expected_root = mt.root(); + + let ms = MerkleStore::from(&mt); + + let path33 = ms.get_path(expected_root, NODE33).unwrap(); + + let pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap(); + + assert!(pmt.get_node(NODE22).is_err()); + assert!(pmt.get_node(NODE23).is_err()); + assert!(pmt.get_node(NODE30).is_err()); + assert!(pmt.get_node(NODE31).is_err()); +} + +/// Checks that the request of the path from the leaf which is not in the PMT will cause an error. +#[test] +fn err_get_path() { + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let expected_root = mt.root(); + + let ms = MerkleStore::from(&mt); + + let path33 = ms.get_path(expected_root, NODE33).unwrap(); + + let pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap(); + + assert!(pmt.get_path(NODE22).is_err()); + assert!(pmt.get_path(NODE23).is_err()); + assert!(pmt.get_path(NODE30).is_err()); + assert!(pmt.get_path(NODE31).is_err()); +} + +#[test] +fn err_update_leaf() { + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let expected_root = mt.root(); + + let ms = MerkleStore::from(&mt); + + let path33 = ms.get_path(expected_root, NODE33).unwrap(); + + let mut pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap(); + + assert!(pmt.update_leaf(NODE22, int_to_node(22).into()).is_err()); + assert!(pmt.update_leaf(NODE23, int_to_node(23).into()).is_err()); + assert!(pmt.update_leaf(NODE30, int_to_node(30).into()).is_err()); + assert!(pmt.update_leaf(NODE31, int_to_node(31).into()).is_err()); +} diff --git a/src/utils.rs b/src/utils.rs index 522d9ea..b9b7849 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -2,6 +2,12 @@ use super::Word; use crate::utils::string::String; use core::fmt::{self, Write}; +#[cfg(not(feature = "std"))] +pub use alloc::format; + +#[cfg(feature = "std")] +pub use std::format; + // RE-EXPORTS // ================================================================================================ pub use winter_utils::{