diff --git a/Cargo.lock b/Cargo.lock index 575ad306fac..8b54878260f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1699,6 +1699,7 @@ dependencies = [ "sha2 0.10.9", "sha3", "strum 0.27.2", + "tailcall", "tempfile", "tezos-smart-rollup-constants", "tezos-smart-rollup-utils", @@ -1722,6 +1723,8 @@ dependencies = [ "blake3", "derive_more", "hex", + "proptest", + "rand 0.9.2", "thiserror 2.0.17", ] @@ -2615,6 +2618,26 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "tailcall" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "695cb8728beab065fe29d03ae105a23c7cdd6e8d9ed5fd717a0bbbd8c56ab0bb" +dependencies = [ + "tailcall-impl", +] + +[[package]] +name = "tailcall-impl" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10b7bf224f4cbce8706b34f6fd1db63bc9b1fa933a3b9e6eee2e78a660416e13" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "tap" version = "1.0.1" diff --git a/Cargo.toml b/Cargo.toml index a5806bcfa5a..5d2a39f9fa5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -67,6 +67,7 @@ rustc_apfloat = "0.2.3" serde_json = "1.0.143" sha2 = "0.10.9" sha3 = "0.10.8" +tailcall = "1.0.1" tempfile = "3.22.0" thiserror = "2.0.17" trait-set = "0.3.0" diff --git a/data/Cargo.toml b/data/Cargo.toml index 97743117a85..f3c83ee6a3d 100644 --- a/data/Cargo.toml +++ b/data/Cargo.toml @@ -13,3 +13,7 @@ blake3.workspace = true derive_more.workspace = true hex.workspace = true thiserror.workspace = true + +[dev-dependencies] +proptest.workspace = true +rand.workspace = true diff --git a/data/src/compressed_merkle_tree.rs b/data/src/compressed_merkle_tree.rs new file mode 100644 index 00000000000..8a985798e05 --- /dev/null +++ b/data/src/compressed_merkle_tree.rs @@ -0,0 +1,694 @@ +// SPDX-FileCopyrightText: 2025 TriliTech +// SPDX-License-Identifier: MIT + +use std::num::NonZeroUsize; + +use crate::hash::Hash; +use crate::merkle_proof::proof_tree::MerkleProof; +use crate::merkle_proof::proof_tree::MerkleProofLeaf; +use crate::merkle_tree::MerkleTree; +use crate::merkle_tree::MerkleTreeLeafData; +use crate::merkle_tree::MerkleTreeNodeData; +use crate::tree::Tree; + +// TODO RV-322: Choose optimal Merkleisation parameters for main memory. +/// Size of the Merkle leaf used for Merkleising [`DynArrays`]. +/// +/// [`DynArrays`]: [`crate::state_backend::layout::DynArray`] +pub const MERKLE_LEAF_SIZE: NonZeroUsize = NonZeroUsize::new(4096).unwrap(); + +/// Type of access associated with leaves in a [`CompressedMerkleTree`]. +/// +/// If a subtree only contains leaves which have not been accessed, it can be +/// compressed into a blinded leaf. Leaves which have been accessed also hold +/// the leaf data. +#[derive(Debug, Clone, PartialEq)] +pub enum CompressedAccessInfo { + /// A leaf which has not been accessed + NoAccess, + /// A leaf which has been accessed + ReadWrite(Vec), +} + +impl CompressedAccessInfo { + pub fn from_access_info(access_info: bool, data: Vec) -> Self { + if access_info { + Self::ReadWrite(data) + } else { + Self::NoAccess + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct CompressedMerkleTreeLeafData { + pub hash: Hash, + pub compressed_access_info: CompressedAccessInfo, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct CompressedMerkleTreeNodeData { + pub hash: Hash, +} + +/// Intermediary representation obtained when compressing a [`MerkleTree`]. +/// +/// For the compressed tree, we only care about the data in the non-blinded leaves. +pub type CompressedMerkleTree = Tree; + +impl CompressedMerkleTree { + pub fn to_proof(self) -> MerkleProof { + match self { + CompressedMerkleTree::Leaf { + data: + CompressedMerkleTreeLeafData { + hash, + compressed_access_info, + }, + } => match compressed_access_info { + CompressedAccessInfo::NoAccess => Tree::Leaf { + data: MerkleProofLeaf::Blind(hash), + }, + CompressedAccessInfo::ReadWrite(data) => Tree::Leaf { + data: MerkleProofLeaf::Read(data), + }, + }, + CompressedMerkleTree::Node { children, .. } => Tree::Node { + data: Default::default(), + children: children.into_iter().map(|child| child.to_proof()).collect(), + }, + } + } +} + +#[derive(Debug, Clone)] +pub struct MerkleTreeCompressionError; + +fn maybe_compress_node(hash: Hash, children: Vec) -> CompressedMerkleTree { + let mut none_accessed = true; + let mut hasher = blake3::Hasher::new(); + for child in children.iter() { + match child { + CompressedMerkleTree::Leaf { + data: + CompressedMerkleTreeLeafData { + hash: leaf_hash, + compressed_access_info, + }, + } => { + hasher.update(leaf_hash.as_ref()); + match compressed_access_info { + CompressedAccessInfo::ReadWrite(_) => none_accessed = false, + _ => {} + }; + } + CompressedMerkleTree::Node { + data: CompressedMerkleTreeNodeData { hash: node_hash }, + .. + } => { + none_accessed = false; + hasher.update(node_hash.as_ref()); + } + } + } + if none_accessed { + CompressedMerkleTree::Leaf { + data: CompressedMerkleTreeLeafData { + hash: Hash::new(hasher.finalize().into()), + compressed_access_info: CompressedAccessInfo::NoAccess, + }, + } + } else { + CompressedMerkleTree::Node { + data: CompressedMerkleTreeNodeData { hash }, + children, + } + } +} + +fn merkle_child_to_compressed_child( + child: MerkleTree, +) -> Result { + if let MerkleTree::Leaf { + data: + MerkleTreeLeafData { + hash, + access_info, + data, + }, + } = child + { + if access_info { + Ok(CompressedMerkleTree::Leaf { + data: CompressedMerkleTreeLeafData { + hash, + compressed_access_info: CompressedAccessInfo::ReadWrite(data), + }, + }) + } else { + Ok(CompressedMerkleTree::Leaf { + data: CompressedMerkleTreeLeafData { + hash, + compressed_access_info: CompressedAccessInfo::NoAccess, + }, + }) + } + } else { + Err(MerkleTreeCompressionError) + } +} + +/// Turns a [MerkleTree] into a [CompressedMerkleTree] +pub fn merkle_tree_to_compressed_merkle_tree( + merkle_tree: MerkleTree, +) -> Result { + let mut nodes: Vec<(MerkleTree, usize)> = vec![(merkle_tree, 0)]; + let mut compressed_nodes: Vec<(CompressedMerkleTree, usize)> = vec![]; + + while let Some((node, parent_index)) = nodes.pop() { + match node { + leaf @ MerkleTree::Leaf { .. } => { + compressed_nodes.push((merkle_child_to_compressed_child(leaf)?, parent_index)); + } + MerkleTree::Node { + data: MerkleTreeNodeData { hash }, + children, + } => { + compressed_nodes.push(( + CompressedMerkleTree::Node { + data: CompressedMerkleTreeNodeData { hash }, + children: vec![], + }, + parent_index, + )); + let new_parent_index = compressed_nodes.len() - 1; + for child in children { + nodes.push((child, new_parent_index)); + } + } + } + } + + while compressed_nodes.len() > 1 { + let (compressed_node, parent_index) = compressed_nodes.pop().unwrap(); + if let (CompressedMerkleTree::Node { children, .. }, _) = + &mut compressed_nodes[parent_index] + { + match compressed_node { + leaf @ CompressedMerkleTree::Leaf { .. } => children.push(leaf), + CompressedMerkleTree::Node { + data: CompressedMerkleTreeNodeData { hash }, + children: node_children, + } => children.push(maybe_compress_node(hash, node_children)), + } + } else { + panic!("This should not happen"); + } + } + + Ok(compressed_nodes.pop().unwrap().0) +} + +/// Turns a [MerkleTree] into a [MerkleProof] +pub fn merkle_tree_to_merkle_proof(merkle_tree: MerkleTree) -> MerkleProof { + merkle_tree_to_compressed_merkle_tree(merkle_tree) + .expect("This conversion should not fail") + .to_proof() +} + +/// Helper function which allows iterating over chunks of a dynamic array +/// and writing them to a writer. The last chunk may be smaller than the +/// Merkle leaf size. The implementations of [`HashState`] and +/// [`ProofLayout`] both use it, ensuring consistency between the two. +/// +/// [`HashState`]: octez_riscv_data::hash::HashState +/// [`ProofLayout`]: crate::state_backend::proof_layout::ProofLayout +pub fn chunks_to_writer [u8; MERKLE_LEAF_SIZE.get()]>( + writer: &mut T, + len: usize, + read: F, +) -> Result<(), std::io::Error> { + let merkle_leaf_size = MERKLE_LEAF_SIZE.get(); + assert!(len >= merkle_leaf_size); + + let mut address = 0; + + while address + merkle_leaf_size <= len { + writer.write_all(read(address).as_slice())?; + address += merkle_leaf_size; + } + + // When the last chunk is smaller than `MERKLE_LEAF_SIZE`, + // read the last `MERKLE_LEAF_SIZE` bytes and pass a subslice containing + // only the bytes not previously read to the writer. + if address != len { + address += merkle_leaf_size; + let buffer = read(len.saturating_sub(merkle_leaf_size)); + writer.write_all(&buffer[address.saturating_sub(len)..])?; + }; + + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::io::Cursor; + + use proptest::prelude::*; + + use super::CompressedAccessInfo; + use super::CompressedMerkleTree; + use super::CompressedMerkleTreeLeafData; + use super::CompressedMerkleTreeNodeData; + use super::MERKLE_LEAF_SIZE; + use super::chunks_to_writer; + use super::merkle_tree_to_compressed_merkle_tree; + use super::merkle_tree_to_merkle_proof; + use crate::hash::Hash; + use crate::hash::HashError; + use crate::merkle_proof::proof_tree::MerkleProof; + use crate::merkle_proof::proof_tree::MerkleProofLeaf; + use crate::merkle_tree::MerkleTree; + use crate::merkle_tree::MerkleTreeLeafData; + + impl CompressedMerkleTree { + /// Get the root hash of a compressed Merkle tree + fn root_hash(&self) -> Hash { + match self { + Self::Node { + data: CompressedMerkleTreeNodeData { hash }, + .. + } => *hash, + Self::Leaf { + data: CompressedMerkleTreeLeafData { hash, .. }, + } => *hash, + } + } + + /// Check the validity of the Merkle root by recomputing all hashes + fn check_root_hash(&self) -> bool { + let mut deque = std::collections::VecDeque::new(); + deque.push_back(self); + + while let Some(node) = deque.pop_front() { + let is_valid_hash = match node { + Self::Leaf { + data: + CompressedMerkleTreeLeafData { + hash, + compressed_access_info, + }, + } => match compressed_access_info { + CompressedAccessInfo::NoAccess => true, + CompressedAccessInfo::ReadWrite(data) => { + &Hash::blake3_hash_bytes(data) == hash + } + }, + Self::Node { + data: CompressedMerkleTreeNodeData { hash }, + children, + } => { + let children_hashes = children.iter().map(|child| { + deque.push_back(child); + child.root_hash() + }); + &Hash::combine(children_hashes) == hash + } + }; + if !is_valid_hash { + return false; + } + } + true + } + } + + fn m_l(data: &[u8], access: bool) -> Result { + let hash = Hash::blake3_hash_bytes(data); + Ok(MerkleTree::Leaf { + data: MerkleTreeLeafData { + hash, + access_info: access, + data: data.to_vec(), + }, + }) + } + + fn m_t(left: MerkleTree, right: MerkleTree) -> MerkleTree { + MerkleTree::make_merkle_node(vec![left, right]) + } + + #[test] + fn test_compression() { + let test = |l: Vec>| -> Result<_, HashError> { + // The LHS leaf will be blinded + let single_leaves_t = m_t(m_l(&l[0], false)?, m_l(&l[1], true)?); + + // The whole subtree will be blinded and compressed + let no_access_t = m_t( + m_l(&l[2], false)?, + m_t(m_l(&l[3], false)?, m_l(&l[4], false)?), + ); + + // No leaf will be blinded, the tree will not be compressed + let read_write_3_t = m_t(m_t(m_l(&l[5], true)?, m_l(&l[6], true)?), m_l(&l[7], true)?); + + // No leaf will be blinded, the tree will not be compressed + let read_write_4_t = m_t( + m_t(m_l(&l[8], true)?, m_l(&l[9], true)?), + m_t(m_l(&l[10], true)?, m_l(&l[11], true)?), + ); + + let combine_isolated_t = m_t(m_t(no_access_t, read_write_3_t), read_write_4_t); + + let mix_t = m_t( + // The whole subtree will be compressed + m_t( + m_l(&l[12], false)?, + m_t( + m_l(&l[13], false)?, + m_t(m_l(&l[14], false)?, m_l(&l[15], false)?), + ), + ), + m_t( + m_t( + m_l(&l[16], true)?, + // Only the non-accessed leaves will be compressed + m_t(m_l(&l[17], false)?, m_l(&l[18], false)?), + ), + m_l(&l[19], true)?, + ), + ); + + let merkle_tree = m_t(single_leaves_t, m_t(combine_isolated_t, mix_t)); + + let merkle_proof_leaf = + |data: &Vec, access: bool| -> Result { + let hash = Hash::blake3_hash_bytes(data); + Ok(MerkleProof::Leaf { + data: if access { + MerkleProofLeaf::Read(data.clone()) + } else { + MerkleProofLeaf::Blind(hash) + }, + }) + }; + + let proof_single_leaves = MerkleProof::Node { + data: Default::default(), + children: vec![ + merkle_proof_leaf(&l[0], false)?, + merkle_proof_leaf(&l[1], true)?, + ], + }; + + // The structure of the original subtree is compressed into a single leaf. + let proof_no_access = MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(Hash::combine([ + Hash::blake3_hash_bytes(&l[2]), + Hash::combine([ + Hash::blake3_hash_bytes(&l[3]), + Hash::blake3_hash_bytes(&l[4]), + ]), + ])), + }; + + let proof_read_write_3 = MerkleProof::Node { + data: Default::default(), + children: vec![ + MerkleProof::Node { + data: Default::default(), + children: vec![ + merkle_proof_leaf(&l[5], true)?, + merkle_proof_leaf(&l[6], true)?, + ], + }, + merkle_proof_leaf(&l[7], true)?, + ], + }; + + let proof_read_write_4 = MerkleProof::Node { + data: Default::default(), + children: vec![ + MerkleProof::Node { + data: Default::default(), + children: vec![ + merkle_proof_leaf(&l[8], true)?, + merkle_proof_leaf(&l[9], true)?, + ], + }, + MerkleProof::Node { + data: Default::default(), + children: vec![ + merkle_proof_leaf(&l[10], true)?, + merkle_proof_leaf(&l[11], true)?, + ], + }, + ], + }; + + let proof_combine_isolated = MerkleProof::Node { + data: Default::default(), + children: vec![ + MerkleProof::Node { + data: Default::default(), + children: vec![proof_no_access, proof_read_write_3], + }, + proof_read_write_4, + ], + }; + + let proof_mix = MerkleProof::Node { + data: Default::default(), + children: vec![ + // The structure of the original subtree is compressed into a single leaf. + MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(Hash::combine([ + Hash::blake3_hash_bytes(&l[12]), + Hash::combine([ + Hash::blake3_hash_bytes(&l[13]), + Hash::combine([ + Hash::blake3_hash_bytes(&l[14]), + Hash::blake3_hash_bytes(&l[15]), + ]), + ]), + ])), + }, + MerkleProof::Node { + data: Default::default(), + children: vec![ + MerkleProof::Node { + data: Default::default(), + children: vec![ + merkle_proof_leaf(&l[16], true)?, + MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(Hash::combine([ + Hash::blake3_hash_bytes(&l[17]), + Hash::blake3_hash_bytes(&l[18]), + ])), + }, + ], + }, + merkle_proof_leaf(&l[19], true)?, + ], + }, + ], + }; + + let proof = MerkleProof::Node { + data: Default::default(), + children: vec![proof_single_leaves, MerkleProof::Node { + data: Default::default(), + children: vec![proof_combine_isolated, proof_mix], + }], + }; + + let merkle_tree_root_hash = merkle_tree.root_hash(); + + let compressed_merkle_tree = merkle_tree_to_compressed_merkle_tree(merkle_tree.clone()); + assert!(compressed_merkle_tree.check_root_hash()); + assert_eq!(compressed_merkle_tree.root_hash(), merkle_tree_root_hash); + + let compressed_merkle_proof = compressed_merkle_tree.clone().to_proof(); + assert_eq!(compressed_merkle_proof, proof); + + assert_eq!(merkle_tree_to_merkle_proof(merkle_tree), proof); + assert_eq!(compressed_merkle_proof.root_hash(), merkle_tree_root_hash); + + Ok(()) + }; + + // this whole proptest macro delegates to a pure rust function in order to have easy access to formatting + proptest!(|(l in prop::collection::vec( + prop::collection::vec(0u8..255, 0..100), + 20 + ))| { + test(l).expect("Unexpected Hashing error"); + }); + } + + fn check(compressed_merkle_tree: CompressedMerkleTree, merkle_proof: MerkleProof) { + let proof_from_compressed_merkle_tree = compressed_merkle_tree.to_proof(); + assert_eq!(proof_from_compressed_merkle_tree, merkle_proof); + } + + #[test] + fn transform_compressed_merkle_tree_to_proof() { + use CompressedAccessInfo::*; + + let gen_hash_data = || { + let data = rand::random::<[u8; 12]>().to_vec(); + let hash = Hash::blake3_hash_bytes(&data); + (data, hash) + }; + + let (data, hash) = gen_hash_data(); + + // Check leaves + check( + CompressedMerkleTree::Leaf { + data: CompressedMerkleTreeLeafData { + hash, + compressed_access_info: NoAccess, + }, + }, + MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(hash), + }, + ); + check( + CompressedMerkleTree::Leaf { + data: CompressedMerkleTreeLeafData { + hash, + compressed_access_info: ReadWrite(data.clone()), + }, + }, + MerkleProof::Leaf { + data: MerkleProofLeaf::Read(data.clone()), + }, + ); + + // Check nodes + let [d0, d1, d2, d3, d4, d5, d6, d7, d8] = [0; 9].map(|_| gen_hash_data()); + let l0 = MerkleProof::Leaf { + data: MerkleProofLeaf::Read(d0.0.clone()), + }; + let l1 = MerkleProof::Leaf { + data: MerkleProofLeaf::Read(d1.0.clone()), + }; + let l2 = MerkleProof::Leaf { + data: MerkleProofLeaf::Read(d2.0.clone()), + }; + let l3 = MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(d3.1), + }; + let l4 = MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(d4.1), + }; + let l5 = MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(d5.1), + }; + + let n6 = MerkleProof::Node { + data: Default::default(), + children: vec![l0, l1, l3], + }; + let n7 = MerkleProof::Node { + data: Default::default(), + children: vec![l4, l2, l5], + }; + let root = MerkleProof::Node { + data: Default::default(), + children: vec![n6, n7], + }; + + let t0 = CompressedMerkleTree::Leaf { + data: CompressedMerkleTreeLeafData { + hash: d0.1, + compressed_access_info: ReadWrite(d0.0), + }, + }; + let t1 = CompressedMerkleTree::Leaf { + data: CompressedMerkleTreeLeafData { + hash: d1.1, + compressed_access_info: ReadWrite(d1.0), + }, + }; + let t2 = CompressedMerkleTree::Leaf { + data: CompressedMerkleTreeLeafData { + hash: d2.1, + compressed_access_info: ReadWrite(d2.0), + }, + }; + let t3 = CompressedMerkleTree::Leaf { + data: CompressedMerkleTreeLeafData { + hash: d3.1, + compressed_access_info: NoAccess, + }, + }; + let t4 = CompressedMerkleTree::Leaf { + data: CompressedMerkleTreeLeafData { + hash: d4.1, + compressed_access_info: NoAccess, + }, + }; + let t5 = CompressedMerkleTree::Leaf { + data: CompressedMerkleTreeLeafData { + hash: d5.1, + compressed_access_info: NoAccess, + }, + }; + + let t6 = CompressedMerkleTree::Node { + data: CompressedMerkleTreeNodeData { hash: d6.1 }, + children: vec![t0, t1, t3], + }; + let t7 = CompressedMerkleTree::Node { + data: CompressedMerkleTreeNodeData { hash: d7.1 }, + children: vec![t4, t2, t5], + }; + let t_root = CompressedMerkleTree::Node { + data: CompressedMerkleTreeNodeData { hash: d8.1 }, + children: vec![t6, t7], + }; + + check(t_root, root); + } + + const LENS: [usize; 4] = [0, 1, 535, 4095]; + const _: () = { + if MERKLE_LEAF_SIZE.get() != 4096 { + panic!( + "Test values in `LENS` assume a specific MERKLE_LEAF_SIZE, change them accordingly" + ); + } + }; + + // Test `chunks_to_writer` with a variety of lengths + macro_rules! generate_test_chunks_to_writer { + ( $name:ident, $i:expr ) => { + proptest! { + #[test] + fn $name( + data in proptest::collection::vec(any::(), 4 * MERKLE_LEAF_SIZE.get() + LENS[$i]) + ) { + const LEN: usize = 4 * MERKLE_LEAF_SIZE.get() + LENS[$i]; + + let read = |pos: usize| { + assert!(pos + MERKLE_LEAF_SIZE.get() <= LEN); + data[pos..pos + MERKLE_LEAF_SIZE.get()].try_into().unwrap() + }; + + let mut writer = Cursor::new(Vec::new()); + chunks_to_writer::< _, _>(&mut writer, LEN, read).unwrap(); + assert_eq!(writer.into_inner(), data); + } + } + } + } + + generate_test_chunks_to_writer!(test_chunks_to_writer_0, 0); + generate_test_chunks_to_writer!(test_chunks_to_writer_1, 1); + generate_test_chunks_to_writer!(test_chunks_to_writer_2, 2); + generate_test_chunks_to_writer!(test_chunks_to_writer_3, 3); +} diff --git a/data/src/hash.rs b/data/src/hash.rs index 4325bdbbc1d..d959b218729 100644 --- a/data/src/hash.rs +++ b/data/src/hash.rs @@ -52,10 +52,14 @@ pub const DIGEST_SIZE: usize = 32; )] #[debug("{}", self)] pub struct Hash { - digest: [u8; DIGEST_SIZE], + pub digest: [u8; DIGEST_SIZE], } impl Hash { + pub fn new(digest: [u8; 32]) -> Self { + Hash { digest } + } + /// Hash a slice of bytes pub fn blake3_hash_bytes(bytes: &[u8]) -> Self { let digest = blake3::hash(bytes).into(); diff --git a/data/src/lib.rs b/data/src/lib.rs index eabc6e87478..e6c8835110b 100644 --- a/data/src/lib.rs +++ b/data/src/lib.rs @@ -4,6 +4,7 @@ // SPDX-License-Identifier: MIT pub mod clone; +pub mod compressed_merkle_tree; pub mod foldable; pub mod hash; pub mod merkle_proof; diff --git a/data/src/merkle_proof/proof_tree.rs b/data/src/merkle_proof/proof_tree.rs index 9b3f5ffb0ae..ae3a42d10c4 100644 --- a/data/src/merkle_proof/proof_tree.rs +++ b/data/src/merkle_proof/proof_tree.rs @@ -5,11 +5,12 @@ use bincode::enc::write::Writer; use super::tag::LeafTag; use super::tag::Tag; -use super::transform::ModifyResult; -use super::transform::impl_modify_map_collect; use crate::hash::Hash; use crate::tree::Tree; +#[derive(Clone, Debug, PartialEq, Default)] +pub struct MerkleProofNodeData; + /// Merkle proof tree structure. /// /// Leaves can be read and/or written to. @@ -23,7 +24,7 @@ use crate::tree::Tree; /// do not need to be stored in either the proof or its encoding. /// /// [`MerkleTree`]: crate::merkle_tree::MerkleTree -pub type MerkleProof = Tree; +pub type MerkleProof = Tree; impl bincode::Encode for MerkleProof { fn encode( @@ -34,25 +35,27 @@ impl bincode::Encode for MerkleProof { while let Some(node) = nodes.pop() { match node { - Self::Node(trees) => { + Self::Node { children, .. } => { Tag::Node.encode(encoder)?; // We add the children in reverse order so that when we pop them from the // `nodes` stack, they are in the original order. - nodes.extend(trees.iter().rev()); - } - - Self::Leaf(MerkleProofLeaf::Read(data)) => { - Tag::Leaf(LeafTag::Read).encode(encoder)?; - - // We want to write the raw data, and avoid the bincode length prefix. The decoder - // will know how many bytes to read. - encoder.writer().write(data.as_slice())?; + nodes.extend(children.iter().rev()); } - - Self::Leaf(MerkleProofLeaf::Blind(hash)) => { - Tag::Leaf(LeafTag::Blind).encode(encoder)?; - hash.encode(encoder)?; + Self::Leaf { data } => { + match data { + MerkleProofLeaf::Read(read_data) => { + Tag::Leaf(LeafTag::Read).encode(encoder)?; + + // We want to write the raw data, and avoid the bincode length prefix. The decoder + // will know how many bytes to read. + encoder.writer().write(read_data.as_slice())?; + } + MerkleProofLeaf::Blind(hash) => { + Tag::Leaf(LeafTag::Blind).encode(encoder)?; + hash.encode(encoder)?; + } + } } } } @@ -77,40 +80,126 @@ pub enum MerkleProofLeaf { Read(Vec), } +enum NodeLeaf { + Node, + Leaf, +} + +struct HashState { + node_leaf: NodeLeaf, + parent_index: usize, + digests: Vec<[u8; 32]>, +} + +impl HashState { + pub fn new(node_leaf: NodeLeaf, parent_index: usize) -> Self { + Self { + node_leaf, + parent_index, + digests: vec![], + } + } + + pub fn new_with_digest(node_leaf: NodeLeaf, parent_index: usize, digest: [u8; 32]) -> Self { + Self { + node_leaf, + parent_index, + digests: vec![digest], + } + } + + pub fn push(&mut self, digest: [u8; 32]) { + self.digests.push(digest); + } + + pub fn get_digest(&self) -> [u8; 32] { + match self.node_leaf { + NodeLeaf::Node => { + let mut hasher = blake3::Hasher::new(); + + for digest in self.digests.iter() { + hasher.update(digest); + } + + hasher.finalize().into() + } + NodeLeaf::Leaf => self.digests[0], + } + } + + pub fn get_parent_index(&self) -> usize { + self.parent_index + } +} + impl MerkleProof { /// Create a new Merkle proof as a read leaf. pub fn leaf_read(data: Vec) -> Self { - MerkleProof::Leaf(MerkleProofLeaf::Read(data)) + MerkleProof::Leaf { + data: MerkleProofLeaf::Read(data), + } } /// Create a new Merkle proof as a blind leaf. pub fn leaf_blind(hash: Hash) -> Self { - MerkleProof::Leaf(MerkleProofLeaf::Blind(hash)) + MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(hash), + } } - /// Compute the root hash of the Merkle proof. pub fn root_hash(&self) -> Hash { - impl_modify_map_collect( - self, - |subtree| match subtree { - Tree::Node(vec) => ModifyResult::NodeContinue((), vec.iter().collect()), - Tree::Leaf(data) => ModifyResult::LeafStop(data), - }, - |leaf| match leaf { - MerkleProofLeaf::Blind(hash) => *hash, - MerkleProofLeaf::Read(data) => Hash::blake3_hash_bytes(data.as_slice()), - }, - |(), leaves| Hash::combine(leaves), - ) + let mut nodes: Vec<(&MerkleProof, usize)> = vec![(self, 0)]; + let mut hashes: Vec = vec![]; + + while let Some((node, parent_index)) = nodes.pop() { + match node { + Tree::Leaf { + data: MerkleProofLeaf::Blind(hash), + } => { + hashes.push(HashState::new_with_digest( + NodeLeaf::Leaf, + parent_index, + hash.digest, + )); + } + Tree::Leaf { + data: MerkleProofLeaf::Read(data), + } => { + hashes.push(HashState::new_with_digest( + NodeLeaf::Leaf, + parent_index, + Hash::blake3_hash_bytes(data.as_slice()).digest, + )); + } + Tree::Node { children, .. } => { + hashes.push(HashState::new(NodeLeaf::Node, parent_index)); + let new_parent_index = hashes.len() - 1; + for child in children.iter() { + nodes.push((&child, new_parent_index)); + } + } + } + } + + while hashes.len() > 1 { + let hash_state = hashes.pop().unwrap(); + hashes[hash_state.get_parent_index()].push(hash_state.get_digest()); + } + + Hash::new(hashes[0].get_digest()) } } impl From<&MerkleProof> for Tag { fn from(value: &MerkleProof) -> Self { match value { - MerkleProof::Node(_) => Tag::Node, - MerkleProof::Leaf(MerkleProofLeaf::Blind(_)) => Tag::Leaf(LeafTag::Blind), - MerkleProof::Leaf(MerkleProofLeaf::Read(_)) => Tag::Leaf(LeafTag::Read), + MerkleProof::Node { .. } => Tag::Node, + MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(_), + } => Tag::Leaf(LeafTag::Blind), + MerkleProof::Leaf { + data: MerkleProofLeaf::Read(_), + } => Tag::Leaf(LeafTag::Read), } } } @@ -128,13 +217,14 @@ mod tests { let merkle_proofs = [ MerkleProof::leaf_read([1, 2, 3].to_vec()), MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[1, 3, 4])), - Tree::Node( - [ + Tree::Node { + data: Default::default(), + children: [ MerkleProof::leaf_read([1, 2, 3].to_vec()), MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[1, 3, 4])), ] .to_vec(), - ), + }, ]; for merkle_proof in merkle_proofs.iter() { bincode::encode_to_vec(merkle_proof, get_bincode_config()) @@ -144,13 +234,14 @@ mod tests { #[test] fn we_can_take_the_merkle_proof_of_the_root_hash() { - let node = Tree::Node( - [ + let node = Tree::Node { + data: Default::default(), + children: [ MerkleProof::leaf_read([1, 2, 3].to_vec()), MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[1, 3, 4])), ] .to_vec(), - ); + }; let _ = node.root_hash(); } } diff --git a/data/src/merkle_tree.rs b/data/src/merkle_tree.rs index cde74e96d59..e882ee1176f 100644 --- a/data/src/merkle_tree.rs +++ b/data/src/merkle_tree.rs @@ -7,17 +7,26 @@ use crate::foldable::Fold; use crate::foldable::Foldable; use crate::foldable::NodeFold; use crate::hash::Hash; +use crate::tree::Tree; + +#[derive(Debug, Clone, PartialEq)] +pub struct MerkleTreeLeafData { + pub hash: Hash, + pub access_info: bool, + pub data: Vec, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct MerkleTreeNodeData { + pub hash: Hash, +} /// A variable-width Merkle tree with access metadata for leaves. /// /// Values of this type are produced by the proof-generating backend to capture /// a snapshot of the machine state along with access information for leaves /// which hold data that was used in a particular evaluation step. -#[derive(Debug, Clone)] -pub enum MerkleTree { - Leaf(Hash, bool, Vec), - Node(Hash, Vec), -} +pub type MerkleTree = Tree; impl MerkleTree { /// Returns the precalculated root hash of the node. @@ -37,23 +46,37 @@ impl MerkleTree { /// ``` pub fn root_hash(&self) -> Hash { match self { - Self::Node(hash, _) => *hash, - Self::Leaf(hash, _, _) => *hash, + Self::Node { + data: MerkleTreeNodeData { hash }, + .. + } => *hash, + Self::Leaf { + data: MerkleTreeLeafData { hash, .. }, + } => *hash, } } /// Creates a merkle tree which is a single leaf pub fn make_merkle_leaf(data: Vec, access_info: bool) -> Self { let hash = Hash::blake3_hash_bytes(&data); - MerkleTree::Leaf(hash, access_info, data) + MerkleTree::Leaf { + data: MerkleTreeLeafData { + hash, + access_info, + data, + }, + } } /// Takes a list of children nodes and creates a /// new parent node on top of them. pub fn make_merkle_node(children: Vec) -> Self { let children_hashes = children.iter().map(|t| t.root_hash()); - let node_hash = Hash::combine(children_hashes); - MerkleTree::Node(node_hash, children) + let hash = Hash::combine(children_hashes); + MerkleTree::Node { + data: MerkleTreeNodeData { hash }, + children, + } } /// Recomputes the hashes for the whole tree @@ -65,8 +88,18 @@ impl MerkleTree { while let Some(node) = deque.pop_front() { let is_valid_hash = match node { - Self::Leaf(hash, _, data) => &Hash::blake3_hash_bytes(data) == hash, - Self::Node(hash, children) => { + Self::Leaf { + data: + MerkleTreeLeafData { + hash, + data: node_data, + .. + }, + } => &Hash::blake3_hash_bytes(node_data) == hash, + Self::Node { + data: MerkleTreeNodeData { hash }, + children, + } => { let children_hashes: Vec = children .iter() .map(|child| { diff --git a/data/src/tree.rs b/data/src/tree.rs index 94c1c94bb9f..3a0e6510a6b 100644 --- a/data/src/tree.rs +++ b/data/src/tree.rs @@ -4,14 +4,14 @@ /// Generic tree structure used to model the [`crate::merkle_proof::proof_tree::MerkleProof`], /// as well as the full & partial shapes of a [`crate::merkle_tree::MerkleTree`]. #[derive(Clone, Debug, PartialEq)] -pub enum Tree { - Node(Vec), - Leaf(A), +pub enum Tree { + Node { data: NodeData, children: Vec }, + Leaf { data: LeafData }, } /// Used in [`crate::merkle_proof::transform::impl_modify_map_collect`] -impl From for Tree { +impl From for Tree { fn from(value: D) -> Self { - Tree::Leaf(value) + Tree::Leaf { data: value } } } diff --git a/src/riscv/lib/Cargo.toml b/src/riscv/lib/Cargo.toml index 9eb4779ab8e..f493e376c60 100644 --- a/src/riscv/lib/Cargo.toml +++ b/src/riscv/lib/Cargo.toml @@ -44,6 +44,7 @@ zerocopy.workspace = true zerocopy-derive.workspace = true octez-riscv-data.workspace = true unty.workspace = true +tailcall.workspace = true [dependencies.__tracing_do_not_use_directly] workspace = true diff --git a/src/riscv/lib/src/pvm/common.rs b/src/riscv/lib/src/pvm/common.rs index 515086acb48..6aa56fb599d 100644 --- a/src/riscv/lib/src/pvm/common.rs +++ b/src/riscv/lib/src/pvm/common.rs @@ -10,6 +10,8 @@ use std::ops::ControlFlow; use bincode::Decode; use bincode::Encode; use octez_riscv_data::clone::CloneState; +use octez_riscv_data::compressed_merkle_tree::CompressedMerkleTree; +use octez_riscv_data::compressed_merkle_tree::merkle_tree_to_compressed_merkle_tree; use octez_riscv_data::foldable::Fold; use octez_riscv_data::foldable::Foldable; use octez_riscv_data::foldable::NodeFold; @@ -47,7 +49,6 @@ use crate::state_backend::ManagerBase; use crate::state_backend::ManagerClone; use crate::state_backend::ProofTree; use crate::state_backend::proof_backend::ProofWrapper; -use crate::state_backend::proof_backend::merkle::merkle_tree_to_merkle_proof; use crate::state_backend::proof_backend::proof::Proof; use crate::state_backend::proof_backend::proof::deserialise_owned; use crate::struct_layout; @@ -396,7 +397,9 @@ where let _ = self.input_request(); let merkle_tree = MerkleTree::from_foldable(self); - let merkle_proof = merkle_tree_to_merkle_proof(merkle_tree); + let compressed_merkle_tree: CompressedMerkleTree = + merkle_tree_to_compressed_merkle_tree(merkle_tree).unwrap(); + let merkle_proof: MerkleProof = compressed_merkle_tree.to_proof(); let final_hash = Hash::from_foldable(self); let proof = Proof::new(merkle_proof, final_hash); diff --git a/src/riscv/lib/src/state_backend/layout.rs b/src/riscv/lib/src/state_backend/layout.rs index 9d123e2840a..3cfac154328 100644 --- a/src/riscv/lib/src/state_backend/layout.rs +++ b/src/riscv/lib/src/state_backend/layout.rs @@ -295,6 +295,8 @@ where #[cfg(test)] mod tests { + use octez_riscv_data::compressed_merkle_tree::CompressedMerkleTree; + use octez_riscv_data::compressed_merkle_tree::merkle_tree_to_compressed_merkle_tree; use octez_riscv_data::hash::Hash; use octez_riscv_data::merkle_tree::MerkleTree; use octez_riscv_data::mode::Normal; @@ -309,7 +311,6 @@ mod tests { use crate::state_backend::ProofLayout; use crate::state_backend::ProofPart; use crate::state_backend::proof_backend::ProofWrapper; - use crate::state_backend::proof_backend::merkle::merkle_tree_to_merkle_proof; use crate::state_backend::proof_backend::proof::deserialise_owned; use crate::state_backend::verify_backend::handle_stepper_panics; @@ -387,7 +388,10 @@ mod tests { assert_eq!(hash, tree_root_hash); // Produce a proof - let proof = merkle_tree_to_merkle_proof(tree); + let compressed_merkle_tree: CompressedMerkleTree = + merkle_tree_to_compressed_merkle_tree(tree) + .expect("This conversion should not fail"); + let proof = compressed_merkle_tree.to_proof(); let proof_hash = proof.root_hash(); assert_eq!(hash, proof_hash); diff --git a/src/riscv/lib/src/state_backend/proof_backend.rs b/src/riscv/lib/src/state_backend/proof_backend.rs index 20ba4458ea0..a7ce09157c7 100644 --- a/src/riscv/lib/src/state_backend/proof_backend.rs +++ b/src/riscv/lib/src/state_backend/proof_backend.rs @@ -467,15 +467,16 @@ impl<'normal> FnManager<'normal, Normal> for ProofWrapper { mod tests { use std::collections::VecDeque; + use octez_riscv_data::compressed_merkle_tree::MERKLE_LEAF_SIZE; use octez_riscv_data::hash::Hash; use octez_riscv_data::merkle_tree::MerkleTree; + use octez_riscv_data::merkle_tree::MerkleTreeLeafData; use octez_riscv_data::mode::Normal; use proptest::array; use proptest::prop_assert; use proptest::prop_assert_eq; use proptest::proptest; - use super::merkle::MERKLE_LEAF_SIZE; use super::*; use crate::state_backend::Cells; use crate::state_backend::DynCells; @@ -547,7 +548,13 @@ mod tests { let merkle_tree = MerkleTree::from_foldable(&proof_cells); merkle_tree.check_root_hash(); match merkle_tree { - MerkleTree::Leaf(hash, access_info, _) => { + MerkleTree::Leaf{ + data: MerkleTreeLeafData { + hash, + access_info, + .. + } + } => { prop_assert_eq!(hash, initial_root_hash); prop_assert!(access_info); } @@ -678,8 +685,8 @@ mod tests { let mut queue = VecDeque::with_capacity(LEAVES + 1); let pages_tree = match merkle_tree { - MerkleTree::Leaf(_, _, _) => panic!("Did not expect leaf"), - MerkleTree::Node(_, mut children) => { + MerkleTree::Leaf{ .. } => panic!("Did not expect leaf"), + MerkleTree::Node{ mut children, ..} => { // The node for the pages is the second child. children.remove(1) }, @@ -689,8 +696,10 @@ mod tests { let mut leaf: usize = 0; while let Some(node) = queue.pop_front() { match node { - MerkleTree::Node(_, children) => queue.extend(children), - MerkleTree::Leaf(_, access_info, _) => { + MerkleTree::Node{ children, ..} => queue.extend(children), + MerkleTree::Leaf{ + data: MerkleTreeLeafData { access_info, .. } + } => { prop_assert_eq!( access_info, read_leaves.contains(&leaf) || diff --git a/src/riscv/lib/src/state_backend/proof_backend/merkle.rs b/src/riscv/lib/src/state_backend/proof_backend/merkle.rs index 44055ec83d6..55376ecb770 100644 --- a/src/riscv/lib/src/state_backend/proof_backend/merkle.rs +++ b/src/riscv/lib/src/state_backend/proof_backend/merkle.rs @@ -4,139 +4,22 @@ //! Merkle trees used for proof generation by the PVM -use std::num::NonZeroUsize; - +use octez_riscv_data::compressed_merkle_tree::MERKLE_LEAF_SIZE; use octez_riscv_data::hash::Hash; use octez_riscv_data::hash::HashError; use octez_riscv_data::merkle_proof::proof_tree::MerkleProof; use octez_riscv_data::merkle_proof::proof_tree::MerkleProofLeaf; -use octez_riscv_data::merkle_proof::transform::ModifyResult; -use octez_riscv_data::merkle_proof::transform::impl_modify_map_collect; use octez_riscv_data::merkle_tree::MerkleTree; +use octez_riscv_data::tree::Tree; use super::DynAccess; -// TODO RV-322: Choose optimal Merkleisation parameters for main memory. -/// Size of the Merkle leaf used for Merkleising [`DynArrays`]. -/// -/// [`DynArrays`]: [`crate::state_backend::layout::DynArray`] -pub const MERKLE_LEAF_SIZE: NonZeroUsize = NonZeroUsize::new(4096).unwrap(); - // TODO RV-322: Choose optimal Merkleisation parameters for main memory. /// Arity of the Merkle tree used for Merkleising [`DynArrays`]. /// /// [`DynArrays`]: [`crate::state_backend::layout::DynArray`] pub const MERKLE_ARITY: usize = 4; -/// Turns a [MerkleTree] into a [CompressedMerkleTree] -fn merkle_tree_to_compressed_merkle_tree(merkle_tree: MerkleTree) -> CompressedMerkleTree { - use CompressedMerkleTree::Leaf as CompressedLeaf; - use CompressedMerkleTree::Node as CompressedNode; - - impl_modify_map_collect( - merkle_tree, - |subtree| match subtree { - MerkleTree::Leaf(hash, access_info, data) => { - ModifyResult::LeafStop((hash, access_info, data)) - } - MerkleTree::Node(hash, children) => ModifyResult::NodeContinue(hash, children), - }, - |(hash, access, data)| (hash, CompressedAccessInfo::from_access_info(access, data)), - |hash, compact_children| { - let (hashes, compressions) = compact_children - .iter() - .map(|child| { - use CompressedAccessInfo::*; - match child { - CompressedLeaf(hash, access_info) => (hash, match access_info { - NoAccess => Some(access_info.clone()), - ReadWrite(_) => None, - }), - CompressedNode(hash, _) => (hash, None), - } - }) - .unzip::<_, _, Vec, Vec<_>>(); - - // Obtain the compression type of all children - // if all children have been compressed successfully to the same type of leaf. - let compression = compressions - .into_iter() - .reduce(|a, b| (a == b).then_some(a?)) - .flatten(); - - match compression { - Some(access) => CompressedLeaf(Hash::combine(hashes), access), - None => CompressedNode(hash, compact_children), - } - }, - ) -} - -/// Turns a [MerkleTree] into a [MerkleProof] -pub(crate) fn merkle_tree_to_merkle_proof(merkle_tree: MerkleTree) -> MerkleProof { - merkle_tree_to_compressed_merkle_tree(merkle_tree).into() -} - -/// Intermediary representation obtained when compressing a [`MerkleTree`]. -/// -/// For the compressed tree, we only care about the data in the non-blinded leaves. -#[derive(Debug, Clone, PartialEq)] -enum CompressedMerkleTree { - Leaf(Hash, CompressedAccessInfo), - Node(Hash, Vec), -} - -impl From for MerkleProof { - fn from(value: CompressedMerkleTree) -> Self { - // Explicitly stating error type to enforce the infallible error type - impl_modify_map_collect( - value, - |subtree| match subtree { - CompressedMerkleTree::Leaf(hash, access) => ModifyResult::LeafStop((hash, access)), - CompressedMerkleTree::Node(_, children) => ModifyResult::NodeContinue((), children), - }, - |(hash, access)| { - use CompressedAccessInfo::*; - match access { - NoAccess => MerkleProofLeaf::Blind(hash), - ReadWrite(data) => MerkleProofLeaf::Read(data), - } - }, - |(), children| MerkleProof::Node(children), - ) - } -} - -/// Used in [`impl_modify_map_collect`] -impl From<(Hash, CompressedAccessInfo)> for CompressedMerkleTree { - fn from(data: (Hash, CompressedAccessInfo)) -> CompressedMerkleTree { - CompressedMerkleTree::Leaf(data.0, data.1) - } -} - -/// Type of access associated with leaves in a [`CompressedMerkleTree`]. -/// -/// If a subtree only contains leaves which have not been accessed, it can be -/// compressed into a blinded leaf. Leaves which have been accessed also hold -/// the leaf data. -#[derive(Debug, Clone, PartialEq)] -enum CompressedAccessInfo { - /// A leaf which has not been accessed - NoAccess, - /// A leaf which has been accessed - ReadWrite(Vec), -} - -impl CompressedAccessInfo { - fn from_access_info(access_info: bool, data: Vec) -> Self { - if access_info { - Self::ReadWrite(data) - } else { - Self::NoAccess - } - } -} - pub trait AccessInfoAggregatable { /// Aggregate the access information of the Merkle tree described by /// the layout of the given data, without constructing the tree. @@ -309,305 +192,3 @@ pub(crate) fn build_custom_merkle_tree( ) })) } - -#[cfg(test)] -mod tests { - use std::io::Cursor; - - use octez_riscv_data::hash::Hash; - use octez_riscv_data::hash::HashError; - use octez_riscv_data::merkle_proof::proof_tree::MerkleProof; - use octez_riscv_data::merkle_proof::proof_tree::MerkleProofLeaf; - use proptest::prelude::*; - - use super::CompressedAccessInfo; - use super::CompressedMerkleTree; - use super::MERKLE_LEAF_SIZE; - use super::MerkleTree; - use super::chunks_to_writer; - use crate::state_backend::proof_backend::merkle::merkle_tree_to_compressed_merkle_tree; - use crate::state_backend::proof_backend::merkle::merkle_tree_to_merkle_proof; - - impl CompressedMerkleTree { - /// Get the root hash of a compressed Merkle tree - fn root_hash(&self) -> Hash { - match self { - Self::Node(hash, _) => *hash, - Self::Leaf(hash, _) => *hash, - } - } - - /// Check the validity of the Merkle root by recomputing all hashes - fn check_root_hash(&self) -> bool { - let mut deque = std::collections::VecDeque::new(); - deque.push_back(self); - - while let Some(node) = deque.pop_front() { - let is_valid_hash = match node { - Self::Leaf(hash, access_info) => match access_info { - CompressedAccessInfo::NoAccess => true, - CompressedAccessInfo::ReadWrite(data) => { - &Hash::blake3_hash_bytes(data) == hash - } - }, - Self::Node(hash, children) => { - let children_hashes = children.iter().map(|child| { - deque.push_back(child); - child.root_hash() - }); - &Hash::combine(children_hashes) == hash - } - }; - if !is_valid_hash { - return false; - } - } - true - } - } - - fn m_l(data: &[u8], access: bool) -> Result { - let hash = Hash::blake3_hash_bytes(data); - Ok(MerkleTree::Leaf(hash, access, data.to_vec())) - } - - fn m_t(left: MerkleTree, right: MerkleTree) -> MerkleTree { - MerkleTree::make_merkle_node(vec![left, right]) - } - - #[test] - fn test_compression() { - let test = |l: Vec>| -> Result<_, HashError> { - // The LHS leaf will be blinded - let single_leaves_t = m_t(m_l(&l[0], false)?, m_l(&l[1], true)?); - - // The whole subtree will be blinded and compressed - let no_access_t = m_t( - m_l(&l[2], false)?, - m_t(m_l(&l[3], false)?, m_l(&l[4], false)?), - ); - - // No leaf will be blinded, the tree will not be compressed - let read_write_3_t = m_t(m_t(m_l(&l[5], true)?, m_l(&l[6], true)?), m_l(&l[7], true)?); - - // No leaf will be blinded, the tree will not be compressed - let read_write_4_t = m_t( - m_t(m_l(&l[8], true)?, m_l(&l[9], true)?), - m_t(m_l(&l[10], true)?, m_l(&l[11], true)?), - ); - - let combine_isolated_t = m_t(m_t(no_access_t, read_write_3_t), read_write_4_t); - - let mix_t = m_t( - // The whole subtree will be compressed - m_t( - m_l(&l[12], false)?, - m_t( - m_l(&l[13], false)?, - m_t(m_l(&l[14], false)?, m_l(&l[15], false)?), - ), - ), - m_t( - m_t( - m_l(&l[16], true)?, - // Only the non-accessed leaves will be compressed - m_t(m_l(&l[17], false)?, m_l(&l[18], false)?), - ), - m_l(&l[19], true)?, - ), - ); - - let merkle_tree = m_t(single_leaves_t, m_t(combine_isolated_t, mix_t)); - - let merkle_proof_leaf = - |data: &Vec, access: bool| -> Result { - let hash = Hash::blake3_hash_bytes(data); - Ok(MerkleProof::Leaf(if access { - MerkleProofLeaf::Read(data.clone()) - } else { - MerkleProofLeaf::Blind(hash) - })) - }; - - let proof_single_leaves = MerkleProof::Node(vec![ - merkle_proof_leaf(&l[0], false)?, - merkle_proof_leaf(&l[1], true)?, - ]); - - // The structure of the original subtree is compressed into a single leaf. - let proof_no_access = MerkleProof::Leaf(MerkleProofLeaf::Blind(Hash::combine([ - Hash::blake3_hash_bytes(&l[2]), - Hash::combine([ - Hash::blake3_hash_bytes(&l[3]), - Hash::blake3_hash_bytes(&l[4]), - ]), - ]))); - - let proof_read_write_3 = MerkleProof::Node(vec![ - MerkleProof::Node(vec![ - merkle_proof_leaf(&l[5], true)?, - merkle_proof_leaf(&l[6], true)?, - ]), - merkle_proof_leaf(&l[7], true)?, - ]); - - let proof_read_write_4 = MerkleProof::Node(vec![ - MerkleProof::Node(vec![ - merkle_proof_leaf(&l[8], true)?, - merkle_proof_leaf(&l[9], true)?, - ]), - MerkleProof::Node(vec![ - merkle_proof_leaf(&l[10], true)?, - merkle_proof_leaf(&l[11], true)?, - ]), - ]); - - let proof_combine_isolated = MerkleProof::Node(vec![ - MerkleProof::Node(vec![proof_no_access, proof_read_write_3]), - proof_read_write_4, - ]); - - let proof_mix = MerkleProof::Node(vec![ - // The structure of the original subtree is compressed into a single leaf. - MerkleProof::Leaf(MerkleProofLeaf::Blind(Hash::combine([ - Hash::blake3_hash_bytes(&l[12]), - Hash::combine([ - Hash::blake3_hash_bytes(&l[13]), - Hash::combine([ - Hash::blake3_hash_bytes(&l[14]), - Hash::blake3_hash_bytes(&l[15]), - ]), - ]), - ]))), - MerkleProof::Node(vec![ - MerkleProof::Node(vec![ - merkle_proof_leaf(&l[16], true)?, - MerkleProof::Leaf(MerkleProofLeaf::Blind(Hash::combine([ - Hash::blake3_hash_bytes(&l[17]), - Hash::blake3_hash_bytes(&l[18]), - ]))), - ]), - merkle_proof_leaf(&l[19], true)?, - ]), - ]); - - let proof = MerkleProof::Node(vec![ - proof_single_leaves, - MerkleProof::Node(vec![proof_combine_isolated, proof_mix]), - ]); - - let merkle_tree_root_hash = merkle_tree.root_hash(); - - let compressed_merkle_tree = merkle_tree_to_compressed_merkle_tree(merkle_tree.clone()); - assert!(compressed_merkle_tree.check_root_hash()); - assert_eq!(compressed_merkle_tree.root_hash(), merkle_tree_root_hash); - - let compressed_merkle_proof: MerkleProof = compressed_merkle_tree.into(); - assert_eq!(compressed_merkle_proof, proof); - - assert_eq!(merkle_tree_to_merkle_proof(merkle_tree), proof); - assert_eq!(compressed_merkle_proof.root_hash(), merkle_tree_root_hash); - - Ok(()) - }; - - // this whole proptest macro delegates to a pure rust function in order to have easy access to formatting - proptest!(|(l in prop::collection::vec( - prop::collection::vec(0u8..255, 0..100), - 20 - ))| { - test(l).expect("Unexpected Hashing error"); - }); - } - - #[test] - fn transform_compressed_merkle_tree_to_proof() { - use CompressedAccessInfo::*; - - let check = |merkle_tree, merkle_proof| { - let proof_from_merkle = MerkleProof::from(merkle_tree); - assert_eq!(proof_from_merkle, merkle_proof); - }; - - let gen_hash_data = || { - let data = rand::random::<[u8; 12]>().to_vec(); - let hash = Hash::blake3_hash_bytes(&data); - (data, hash) - }; - - let (data, hash) = gen_hash_data(); - - // Check leaves - check( - CompressedMerkleTree::Leaf(hash, NoAccess), - MerkleProof::Leaf(MerkleProofLeaf::Blind(hash)), - ); - check( - CompressedMerkleTree::Leaf(hash, ReadWrite(data.clone())), - MerkleProof::Leaf(MerkleProofLeaf::Read(data.clone())), - ); - - // Check nodes - let [d0, d1, d2, d3, d4, d5, d6, d7, d8] = [0; 9].map(|_| gen_hash_data()); - let l0 = MerkleProof::Leaf(MerkleProofLeaf::Read(d0.0.clone())); - let l1 = MerkleProof::Leaf(MerkleProofLeaf::Read(d1.0.clone())); - let l2 = MerkleProof::Leaf(MerkleProofLeaf::Read(d2.0.clone())); - let l3 = MerkleProof::Leaf(MerkleProofLeaf::Blind(d3.1)); - let l4 = MerkleProof::Leaf(MerkleProofLeaf::Blind(d4.1)); - let l5 = MerkleProof::Leaf(MerkleProofLeaf::Blind(d5.1)); - - let n6 = MerkleProof::Node(vec![l0, l1, l3]); - let n7 = MerkleProof::Node(vec![l4, l2, l5]); - let root = MerkleProof::Node(vec![n6, n7]); - - let t0 = CompressedMerkleTree::Leaf(d0.1, ReadWrite(d0.0)); - let t1 = CompressedMerkleTree::Leaf(d1.1, ReadWrite(d1.0)); - let t2 = CompressedMerkleTree::Leaf(d2.1, ReadWrite(d2.0)); - let t3 = CompressedMerkleTree::Leaf(d3.1, NoAccess); - let t4 = CompressedMerkleTree::Leaf(d4.1, NoAccess); - let t5 = CompressedMerkleTree::Leaf(d5.1, NoAccess); - - let t6 = CompressedMerkleTree::Node(d6.1, vec![t0, t1, t3]); - let t7 = CompressedMerkleTree::Node(d7.1, vec![t4, t2, t5]); - let t_root = CompressedMerkleTree::Node(d8.1, vec![t6, t7]); - - check(t_root, root); - } - - const LENS: [usize; 4] = [0, 1, 535, 4095]; - const _: () = { - if MERKLE_LEAF_SIZE.get() != 4096 { - panic!( - "Test values in `LENS` assume a specific MERKLE_LEAF_SIZE, change them accordingly" - ); - } - }; - - // Test `chunks_to_writer` with a variety of lengths - macro_rules! generate_test_chunks_to_writer { - ( $name:ident, $i:expr ) => { - proptest! { - #[test] - fn $name( - data in proptest::collection::vec(any::(), 4 * MERKLE_LEAF_SIZE.get() + LENS[$i]) - ) { - const LEN: usize = 4 * MERKLE_LEAF_SIZE.get() + LENS[$i]; - - let read = |pos: usize| { - assert!(pos + MERKLE_LEAF_SIZE.get() <= LEN); - data[pos..pos + MERKLE_LEAF_SIZE.get()].try_into().unwrap() - }; - - let mut writer = Cursor::new(Vec::new()); - chunks_to_writer::< _, _>(&mut writer, LEN, read).unwrap(); - assert_eq!(writer.into_inner(), data); - } - } - } - } - - generate_test_chunks_to_writer!(test_chunks_to_writer_0, 0); - generate_test_chunks_to_writer!(test_chunks_to_writer_1, 1); - generate_test_chunks_to_writer!(test_chunks_to_writer_2, 2); - generate_test_chunks_to_writer!(test_chunks_to_writer_3, 3); -} diff --git a/src/riscv/lib/src/state_backend/proof_backend/proof.rs b/src/riscv/lib/src/state_backend/proof_backend/proof.rs index ddd05847f99..24e8b219bcf 100644 --- a/src/riscv/lib/src/state_backend/proof_backend/proof.rs +++ b/src/riscv/lib/src/state_backend/proof_backend/proof.rs @@ -179,12 +179,16 @@ mod tests { /// Generate the serialisation bound for a [`MerkleProof`] leaf. fn from_merkle_leaf(leaf: &MerkleProof) -> Self { match leaf { - MerkleProof::Node(_) => panic!("Expected a Merkle proof leaf"), - MerkleProof::Leaf(MerkleProofLeaf::Blind(hash)) => Self { + MerkleProof::Node { .. } => panic!("Expected a Merkle proof leaf"), + MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(hash), + } => Self { nodes_count: 1, content_size: hash.as_ref().len() as u64, }, - MerkleProof::Leaf(MerkleProofLeaf::Read(data)) => Self { + MerkleProof::Leaf { + data: MerkleProofLeaf::Read(data), + } => Self { nodes_count: 1, content_size: data.len() as u64, }, @@ -201,8 +205,12 @@ mod tests { let blind_hash: Hash = Hash::blake3_hash_bytes(&raw_array); match is_leaf_read { - true => MerkleProof::Leaf(MerkleProofLeaf::Read(raw_array)), - false => MerkleProof::Leaf(MerkleProofLeaf::Blind(blind_hash)), + true => MerkleProof::Leaf { + data: MerkleProofLeaf::Read(raw_array), + }, + false => MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(blind_hash), + }, } } @@ -223,12 +231,16 @@ mod tests { let raw_array: [u8; 10] = rand::random(); - let rleaf = MerkleProof::Leaf(MerkleProofLeaf::Read(raw_array.to_vec())); + let rleaf = MerkleProof::Leaf { + data: MerkleProofLeaf::Read(raw_array.to_vec()), + }; check_serialisation(rleaf, &[&[TAG_READ], raw_array.as_slice()].concat()); let hash = Hash::blake3_hash_bytes(&raw_array); check_serialisation( - MerkleProof::Leaf(MerkleProofLeaf::Blind(hash)), + MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(hash), + }, &[&[TAG_BLIND], hash.as_ref()].concat(), ); } @@ -240,15 +252,29 @@ mod tests { let h1 = Hash::blake3_hash_bytes(&[1, 2, 3]); let h2 = Hash::blake3_hash_bytes(&[20, 30, 1, 5, 6]); - let n1 = MerkleProof::Leaf(MerkleProofLeaf::Read(vec![12, 15, 30, 40])); - let n2 = MerkleProof::Leaf(MerkleProofLeaf::Blind(h1)); - let n3 = MerkleProof::Leaf(MerkleProofLeaf::Blind(h2)); - let n4 = MerkleProof::Leaf(MerkleProofLeaf::Read(vec![123, 234, 42, 1, 2, 3])); - - let root = MerkleProof::Node(vec![n1.clone()]); + let n1 = MerkleProof::Leaf { + data: MerkleProofLeaf::Read(vec![12, 15, 30, 40]), + }; + let n2 = MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(h1), + }; + let n3 = MerkleProof::Leaf { + data: MerkleProofLeaf::Blind(h2), + }; + let n4 = MerkleProof::Leaf { + data: MerkleProofLeaf::Read(vec![123, 234, 42, 1, 2, 3]), + }; + + let root = MerkleProof::Node { + data: Default::default(), + children: vec![n1.clone()], + }; check_serialisation(root, &[TAG_NODE, TAG_READ, 12, 15, 30, 40]); - let root = MerkleProof::Node(vec![n1.clone(), n2.clone()]); + let root = MerkleProof::Node { + data: Default::default(), + children: vec![n1.clone(), n2.clone()], + }; check_serialisation( root, &[ @@ -261,7 +287,10 @@ mod tests { .concat(), ); - let root = MerkleProof::Node(vec![n1.clone(), n2.clone(), n3.clone()]); + let root = MerkleProof::Node { + data: Default::default(), + children: vec![n1.clone(), n2.clone(), n3.clone()], + }; check_serialisation( root, &[ @@ -276,7 +305,10 @@ mod tests { .concat(), ); - let root = MerkleProof::Node(vec![n1.clone(), n2.clone(), n4.clone(), n3.clone()]); + let root = MerkleProof::Node { + data: Default::default(), + children: vec![n1.clone(), n2.clone(), n4.clone(), n3.clone()], + }; check_serialisation( root, &[ @@ -313,7 +345,10 @@ mod tests { .map(SerialisationBound::from_merkle_leaf) .collect(); - let root = MerkleProof::Node(children); + let root = MerkleProof::Node { + data: Default::default(), + children: children, + }; let bound = SerialisationBound::node_combine(bounds); check_bounds(root, &bound); @@ -356,7 +391,7 @@ mod tests { nothing_taken = new_bounds.is_empty(); if !nothing_taken { - let node = MerkleProof::Node(new_children); + let node = MerkleProof::Node{ data: Default::default(), children: new_children}; let bound = SerialisationBound::node_combine(new_bounds); new_nodes.push((node, bound)); diff --git a/src/riscv/lib/src/state_backend/proof_backend/proof/deserialise_owned.rs b/src/riscv/lib/src/state_backend/proof_backend/proof/deserialise_owned.rs index 12b3e76f292..84b07ef69e1 100644 --- a/src/riscv/lib/src/state_backend/proof_backend/proof/deserialise_owned.rs +++ b/src/riscv/lib/src/state_backend/proof_backend/proof/deserialise_owned.rs @@ -73,13 +73,11 @@ impl ProofTreeDeserialiser<'_> { pub fn deserialise_as_leaf(self) -> Result>> { match self.0 { ProofPart::Absent => Ok(Partial::Absent), - ProofPart::Present(Tree::Node(_)) => Err(ProofError::UnexpectedNode), - ProofPart::Present(Tree::Leaf(MerkleProofLeaf::Blind(hash))) => { - Ok(Partial::Blinded(*hash)) - } - ProofPart::Present(Tree::Leaf(MerkleProofLeaf::Read(items))) => { - Ok(Partial::Present(items.clone())) - } + ProofPart::Present(Tree::Node { .. }) => Err(ProofError::UnexpectedNode), + ProofPart::Present(Tree::Leaf { data }) => match data { + MerkleProofLeaf::Blind(hash) => Ok(Partial::Blinded(*hash)), + MerkleProofLeaf::Read(items) => Ok(Partial::Present(items.clone())), + }, } } @@ -87,14 +85,12 @@ impl ProofTreeDeserialiser<'_> { pub fn deserialise_as_node(self) -> Result>> { match self.0 { ProofPart::Absent => Ok(Partial::Absent), - ProofPart::Present(Tree::Leaf(MerkleProofLeaf::Blind(hash))) => { - Ok(Partial::Blinded(*hash)) - } - ProofPart::Present(Tree::Leaf(MerkleProofLeaf::Read(_))) => { - Err(ProofError::UnexpectedLeaf) - } - ProofPart::Present(Tree::Node(trees)) => Ok(Partial::Present( - trees + ProofPart::Present(Tree::Leaf { data }) => match data { + MerkleProofLeaf::Blind(hash) => Ok(Partial::Blinded(*hash)), + MerkleProofLeaf::Read(_) => Err(ProofError::UnexpectedLeaf), + }, + ProofPart::Present(Tree::Node { children, .. }) => Ok(Partial::Present( + children .iter() .map(ProofPart::Present) .map(ProofTreeDeserialiser) diff --git a/src/riscv/lib/src/state_backend/proof_backend/proof/deserialiser.rs b/src/riscv/lib/src/state_backend/proof_backend/proof/deserialiser.rs index c6b2606da51..7646517adbc 100644 --- a/src/riscv/lib/src/state_backend/proof_backend/proof/deserialiser.rs +++ b/src/riscv/lib/src/state_backend/proof_backend/proof/deserialiser.rs @@ -144,10 +144,13 @@ mod tests { assert_eq!(comp_fn.into_result(), 0); // We expect to get the Absent case since the father of the nested node is blinded - let merkle_proof = MerkleProof::Node(vec![ - MerkleProof::leaf_read(Hash::blake3_hash_bytes(&[0, 1, 2]).as_ref().to_vec()), - MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[3, 4, 5])), - ]); + let merkle_proof = MerkleProof::Node { + data: Default::default(), + children: vec![ + MerkleProof::leaf_read(Hash::blake3_hash_bytes(&[0, 1, 2]).as_ref().to_vec()), + MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[3, 4, 5])), + ], + }; let proof: ProofTreeDeserialiser = ProofTree::Present(&merkle_proof).into(); let comp_fn = computation_i16(proof).unwrap(); assert_eq!(comp_fn.into_result(), 0); @@ -236,10 +239,16 @@ mod tests { } // the same test for the OwnedDeserialiser - let merkle_proof = MerkleProof::Node(vec![ - MerkleProof::leaf_read(hash_read[0..5].to_vec()), - MerkleProof::Node(vec![MerkleProof::leaf_read(bool_read.to_vec())]), - ]); + let merkle_proof = MerkleProof::Node { + data: Default::default(), + children: vec![ + MerkleProof::leaf_read(hash_read[0..5].to_vec()), + MerkleProof::Node { + data: Default::default(), + children: vec![MerkleProof::leaf_read(bool_read.to_vec())], + }, + ], + }; let res = run_owned_deserialiser(computation_bool, &merkle_proof); @@ -273,10 +282,16 @@ mod tests { assert!(matches!(res, Err(ProofError::Deserialise(_)))); - let merkle_proof = MerkleProof::Node(vec![ - MerkleProof::leaf_read(hash_read.to_vec()), - MerkleProof::Node(vec![MerkleProof::leaf_read(bad_bool_bincode.to_vec())]), - ]); + let merkle_proof = MerkleProof::Node { + data: Default::default(), + children: vec![ + MerkleProof::leaf_read(hash_read.to_vec()), + MerkleProof::Node { + data: Default::default(), + children: vec![MerkleProof::leaf_read(bad_bool_bincode.to_vec())], + }, + ], + }; let res = run_owned_deserialiser(computation_bool, &merkle_proof); eprintln!("Result: {res:?}"); assert!(matches!(res, Err(ProofError::Deserialise(_)))); @@ -306,12 +321,16 @@ mod tests { #[test] fn test_blind_computation() { // The nested leaf is blinded - let absent_shape = MerkleProof::Node(vec![ - MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[0, 1, 2])), - MerkleProof::Node(vec![MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[ - 0, 1, 2, - ]))]), - ]); + let absent_shape = MerkleProof::Node { + data: Default::default(), + children: vec![ + MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[0, 1, 2])), + MerkleProof::Node { + data: Default::default(), + children: vec![MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[0, 1, 2]))], + }, + ], + }; let comp_fn = computation_i16::(ProofTree::Present(&absent_shape).into()); @@ -354,22 +373,46 @@ mod tests { #[test] fn test_bad_structure() { - let bad_shape_1 = MerkleProof::Node(vec![]); - let bad_shape_2 = MerkleProof::Node(vec![ - MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[0, 1, 2])), - MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[0, 1, 2])), - MerkleProof::Node(vec![]), - MerkleProof::Node(vec![]), - MerkleProof::Node(vec![]), - ]); - let bad_shape_3 = MerkleProof::Node(vec![ - MerkleProof::Node(vec![]), - MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[0, 1, 2])), - ]); - let bad_shape_4 = MerkleProof::Node(vec![ - MerkleProof::leaf_read([42_u8; 32].to_vec()), - MerkleProof::leaf_read(100_i32.to_le_bytes().to_vec()), - ]); + let bad_shape_1 = MerkleProof::Node { + data: Default::default(), + children: vec![], + }; + let bad_shape_2 = MerkleProof::Node { + data: Default::default(), + children: vec![ + MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[0, 1, 2])), + MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[0, 1, 2])), + MerkleProof::Node { + data: Default::default(), + children: vec![], + }, + MerkleProof::Node { + data: Default::default(), + children: vec![], + }, + MerkleProof::Node { + data: Default::default(), + children: vec![], + }, + ], + }; + let bad_shape_3 = MerkleProof::Node { + data: Default::default(), + children: vec![ + MerkleProof::Node { + data: Default::default(), + children: vec![], + }, + MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[0, 1, 2])), + ], + }; + let bad_shape_4 = MerkleProof::Node { + data: Default::default(), + children: vec![ + MerkleProof::leaf_read([42_u8; 32].to_vec()), + MerkleProof::leaf_read(100_i32.to_le_bytes().to_vec()), + ], + }; // Tree is missing branches let comp_fn = @@ -451,12 +494,15 @@ mod tests { #[test] fn test_valid_computation() { - let merkleproof = MerkleProof::Node(vec![ - MerkleProof::leaf_read(0x140A_0000_i32.to_le_bytes().to_vec()), - MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[3, 4, 5])), - MerkleProof::leaf_read(0xC0005_i32.to_le_bytes().to_vec()), - MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[9, 10, 11])), - ]); + let merkleproof = MerkleProof::Node { + data: Default::default(), + children: vec![ + MerkleProof::leaf_read(0x140A_0000_i32.to_le_bytes().to_vec()), + MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[3, 4, 5])), + MerkleProof::leaf_read(0xC0005_i32.to_le_bytes().to_vec()), + MerkleProof::leaf_blind(Hash::blake3_hash_bytes(&[9, 10, 11])), + ], + }; let proof: ProofTreeDeserialiser = ProofTree::Present(&merkleproof).into(); let comp_fn = computation_leaves(proof).unwrap(); diff --git a/src/riscv/lib/src/state_backend/proof_layout.rs b/src/riscv/lib/src/state_backend/proof_layout.rs index 1a4aa562252..83a6aa6c2dc 100644 --- a/src/riscv/lib/src/state_backend/proof_layout.rs +++ b/src/riscv/lib/src/state_backend/proof_layout.rs @@ -10,6 +10,7 @@ use bincode::Decode; use bincode::Encode; use bincode::error::DecodeError; use bincode::error::EncodeError; +use octez_riscv_data::compressed_merkle_tree::MERKLE_LEAF_SIZE; use octez_riscv_data::hash::Hash; use octez_riscv_data::merkle_proof::Deserialiser; use octez_riscv_data::merkle_proof::DeserialiserError; @@ -33,7 +34,6 @@ use super::Layout; use super::Many; use super::RefVerifyAlloc; use super::proof_backend::merkle::MERKLE_ARITY; -use super::proof_backend::merkle::MERKLE_LEAF_SIZE; use super::proof_backend::proof::deserialiser::Result; use super::verify_backend; use super::verify_backend::PartialState; @@ -153,11 +153,11 @@ impl<'a> ProofTree<'a> { }; match proof { - Tree::Node(branches) => { + Tree::Node { children, .. } => { let branches: &[MerkleProof; LEN] = - branches.as_slice().try_into().map_err(|_| { + children.as_slice().try_into().map_err(|_| { ProofError::BadNumberOfBranches { - got: branches.len(), + got: children.len(), expected: LEN, } })?; @@ -174,7 +174,7 @@ impl<'a> ProofTree<'a> { .unwrap()) } - Tree::Leaf(leaf) => match leaf { + Tree::Leaf { data } => match data { MerkleProofLeaf::Blind(_hash) => Ok(boxed_array![ProofTree::Absent; LEN]), _ => Err(ProofError::UnexpectedLeaf)?, }, @@ -185,8 +185,8 @@ impl<'a> ProofTree<'a> { pub fn into_leaf(self) -> Result> { if let ProofTree::Present(proof) = self { match proof { - Tree::Node(_) => Err(ProofError::UnexpectedNode), - Tree::Leaf(leaf) => match leaf { + Tree::Node { .. } => Err(ProofError::UnexpectedNode), + Tree::Leaf { data } => match data { MerkleProofLeaf::Blind(_) => Ok(ProofPart::Absent), MerkleProofLeaf::Read(data) => Ok(ProofPart::Present(data.as_slice())), }, @@ -204,11 +204,11 @@ impl<'a> ProofTree<'a> { return Err(PartialHashError::PotentiallyRecoverable); }; - let Tree::Leaf(leaf) = proof else { + let Tree::Leaf { data } = proof else { return Err(ProofError::UnexpectedNode.into()); }; - let hash = match leaf { + let hash = match data { MerkleProofLeaf::Blind(hash) => *hash, MerkleProofLeaf::Read(data) => Hash::blake3_hash_bytes(data), }; @@ -232,14 +232,14 @@ impl<'a> ProofTree<'a> { }; match proof { - Tree::Node(branches) if branches.len() != LEN => Err(PartialHashError::FromProof( - ProofError::BadNumberOfBranches { - got: branches.len(), + Tree::Node { children, .. } if children.len() != LEN => Err( + PartialHashError::FromProof(ProofError::BadNumberOfBranches { + got: children.len(), expected: LEN, - }, - )), - Tree::Node(branches) => Ok(( - branches + }), + ), + Tree::Node { children, .. } => Ok(( + children .iter() .map(ProofTree::Present) .collect::>() @@ -248,7 +248,7 @@ impl<'a> ProofTree<'a> { .map_err(|_| PartialHashError::Fatal)?, None, )), - Tree::Leaf(leaf) => match leaf { + Tree::Leaf { data: leaf } => match leaf { MerkleProofLeaf::Blind(hash) => { Ok((boxed_array![ProofTree::Absent; LEN], Some(*hash))) } @@ -282,7 +282,10 @@ impl OwnedProofPart { match partial { Partial::Absent => OwnedProofPart::Absent, Partial::Blinded(hash) => OwnedProofPart::Present(MerkleProof::leaf_blind(hash)), - Partial::Present(children) => OwnedProofPart::Present(MerkleProof::Node(children)), + Partial::Present(children) => OwnedProofPart::Present(MerkleProof::Node { + data: Default::default(), + children, + }), } } @@ -309,7 +312,10 @@ impl OwnedProofPart { } } - OwnedProofPart::Present(MerkleProof::Node(partial_children)) + OwnedProofPart::Present(MerkleProof::Node { + data: Default::default(), + children: partial_children, + }) } } @@ -507,32 +513,34 @@ impl ProofLayout for DynArray { }; let children = match proof { - Tree::Leaf(MerkleProofLeaf::Blind(hash)) => { - if !region.is_completely_absent() { - // The verifier contains data, but the proof is not in the right shape for us - // to insert back the data and obtain a root hash. This indicates an invalid - // proof. - // - // From a partial Merkle tree perspective, technically you can blind a node in - // the proof and then write all the data necessary to expand the node fully - // after the verification is done. In the context of a dynamic region that - // would entail setting the length and writing all pages. However, dynamic - // regions cannot be re-created hence you cannot "write" the length. Any - // modification requires at least the length node to be present. Hence, - // blinding the pages node can be paired with modifications, but you cannot - // blind the dynamic region's root node if there are modifications. - return Err(PartialHashError::Fatal); - } - - // We already know that the dynamic region has been untouched at this point. - return Ok(*hash); - } + Tree::Leaf { data } => { + match data { + MerkleProofLeaf::Blind(hash) => { + if !region.is_completely_absent() { + // The verifier contains data, but the proof is not in the right shape for us + // to insert back the data and obtain a root hash. This indicates an invalid + // proof. + // + // From a partial Merkle tree perspective, technically you can blind a node in + // the proof and then write all the data necessary to expand the node fully + // after the verification is done. In the context of a dynamic region that + // would entail setting the length and writing all pages. However, dynamic + // regions cannot be re-created hence you cannot "write" the length. Any + // modification requires at least the length node to be present. Hence, + // blinding the pages node can be paired with modifications, but you cannot + // blind the dynamic region's root node if there are modifications. + return Err(PartialHashError::Fatal); + } - Tree::Leaf(MerkleProofLeaf::Read(_)) => { - return Err(PartialHashError::FromProof(ProofError::UnexpectedLeaf)); + // We already know that the dynamic region has been untouched at this point. + return Ok(*hash); + } + MerkleProofLeaf::Read(_) => { + return Err(PartialHashError::FromProof(ProofError::UnexpectedLeaf)); + } + } } - - Tree::Node(children) => children.as_slice(), + Tree::Node { children, .. } => children.as_slice(), }; let [length_tree, pages_tree] = children else { @@ -544,12 +552,15 @@ impl ProofLayout for DynArray { )); }; - let length = if let Tree::Leaf(MerkleProofLeaf::Read(data)) = length_tree { + let length = if let Tree::Leaf { + data: MerkleProofLeaf::Read(read_data), + } = length_tree + { // The length data must be present if the node is present. Practically, if there is any // pages data to be dealt with we require the length. Or if there is no pages data, // then the only reason the parent node would be present is if the length was to be // read during verification. - serialisation::deserialise::(data).map_err(ProofError::Deserialise)? as usize + serialisation::deserialise::(read_data).map_err(ProofError::Deserialise)? as usize } else { return Err(PartialHashError::Fatal); }; @@ -1064,6 +1075,9 @@ fn push_work_items_for_branches<'a, const CHILDREN: usize>( #[cfg(test)] mod tests { + use octez_riscv_data::compressed_merkle_tree::CompressedMerkleTree; + use octez_riscv_data::compressed_merkle_tree::merkle_tree_to_compressed_merkle_tree; + use octez_riscv_data::compressed_merkle_tree::merkle_tree_to_merkle_proof; use octez_riscv_data::merkle_tree::MerkleTree; use octez_riscv_data::mode::Prove; use proptest::prop_assert; @@ -1078,7 +1092,6 @@ mod tests { use crate::state_backend::ManagerWrite; use crate::state_backend::proof_backend::ProofRegion; use crate::state_backend::proof_backend::ProofWrapper; - use crate::state_backend::proof_backend::merkle::merkle_tree_to_merkle_proof; use crate::state_backend::proof_backend::proof::deserialise_owned; const CELLS_SIZE: usize = 32; @@ -1167,7 +1180,9 @@ mod tests { let post_hash = Hash::from_foldable(&proof_cell); let tree = MerkleTree::from_foldable(&proof_cell); - let proof_tree = merkle_tree_to_merkle_proof(tree); + let compressed_merkle_tree: CompressedMerkleTree = + merkle_tree_to_compressed_merkle_tree(tree).expect("This conversion should not fail"); + let proof_tree = compressed_merkle_tree.to_proof(); assert_eq!(proof_tree.root_hash(), init_hash); // Instantiating the verifier state allows us to replay the computation and verify it does diff --git a/src/riscv/lib/src/state_backend/region.rs b/src/riscv/lib/src/state_backend/region.rs index 1591737d02d..60647bed396 100644 --- a/src/riscv/lib/src/state_backend/region.rs +++ b/src/riscv/lib/src/state_backend/region.rs @@ -13,6 +13,8 @@ use bincode::enc::Encoder; use bincode::error::DecodeError; use bincode::error::EncodeError; use octez_riscv_data::clone::CloneState; +use octez_riscv_data::compressed_merkle_tree::MERKLE_LEAF_SIZE; +use octez_riscv_data::compressed_merkle_tree::chunks_to_writer; use octez_riscv_data::foldable::Fold; use octez_riscv_data::foldable::Foldable; use octez_riscv_data::hash::Hash; @@ -43,9 +45,7 @@ use crate::state_backend::Elem; use crate::state_backend::RegionProj; use crate::state_backend::normal_backend::region_elem_offset; use crate::state_backend::proof_backend::merkle::MERKLE_ARITY; -use crate::state_backend::proof_backend::merkle::MERKLE_LEAF_SIZE; use crate::state_backend::proof_backend::merkle::MerkleWriter; -use crate::state_backend::proof_backend::merkle::chunks_to_writer; use crate::state_context::projection::ApplyCons; use crate::state_context::projection::CellCons; use crate::state_context::projection::CellsCons; diff --git a/src/riscv/lib/src/state_backend/verify_backend.rs b/src/riscv/lib/src/state_backend/verify_backend.rs index 43f6212ec0d..798ea8c5626 100644 --- a/src/riscv/lib/src/state_backend/verify_backend.rs +++ b/src/riscv/lib/src/state_backend/verify_backend.rs @@ -11,6 +11,7 @@ use std::panic::resume_unwind; use bincode::Encode; use bincode::enc::Encoder; use bincode::error::EncodeError; +use octez_riscv_data::compressed_merkle_tree::MERKLE_LEAF_SIZE; use octez_riscv_data::hash::Hash; use octez_riscv_data::mode::Normal; use octez_riscv_data::mode::Verify; @@ -26,7 +27,6 @@ use super::Ref; use crate::state_backend::Elem; use crate::state_backend::ProofError; use crate::state_backend::elem_bytes; -use crate::state_backend::proof_backend::merkle::MERKLE_LEAF_SIZE; /// Panic payload that is raised when a value isn't present when running in `Verify` mode. #[derive(Copy, Clone, Debug, Eq, PartialEq, derive_more::Display, thiserror::Error)] diff --git a/src/riscv/lib/tests/test_proofs.rs b/src/riscv/lib/tests/test_proofs.rs index f6888e9e326..5301982554f 100644 --- a/src/riscv/lib/tests/test_proofs.rs +++ b/src/riscv/lib/tests/test_proofs.rs @@ -250,11 +250,22 @@ mod proof_helpers { use octez_riscv_data::tree::Tree; pub fn fully_blinded(hash: Hash) -> Proof { - Proof::new(Tree::Leaf(MerkleProofLeaf::Blind(hash)), hash) + Proof::new( + Tree::Leaf { + data: MerkleProofLeaf::Blind(hash), + }, + hash, + ) } pub(crate) fn empty(hash: Hash) -> Proof { - Proof::new(Tree::Node(Vec::new()), hash) + Proof::new( + Tree::Node { + data: Default::default(), + children: Vec::new(), + }, + hash, + ) } pub(crate) fn with_final_hash(proof: &Proof, hash: Hash) -> Proof {