diff --git a/common/append_merkle/src/lib.rs b/common/append_merkle/src/lib.rs index ccb7f6d..ee2ddfc 100644 --- a/common/append_merkle/src/lib.rs +++ b/common/append_merkle/src/lib.rs @@ -16,12 +16,63 @@ use tracing::{trace, warn}; use crate::merkle_tree::MerkleTreeWrite; pub use crate::merkle_tree::{ - Algorithm, HashElement, MerkleTreeInitialData, MerkleTreeRead, ZERO_HASHES, + Algorithm, HashElement, MerkleTreeInitialData, MerkleTreeRead, OptionalHash, ZERO_HASHES, }; pub use crate::node_manager::{EmptyNodeDatabase, NodeDatabase, NodeManager, NodeTransaction}; pub use proof::{Proof, RangeProof}; pub use sha3::Sha3Algorithm; +// Helper functions for converting between H256 and OptionalHash types +use ethereum_types::H256; + +impl AppendMerkleTree { + /// Convert a proof of OptionalHash to a proof of H256 + pub fn convert_proof_to_h256( + proof: &Proof, + ) -> Result, anyhow::Error> { + let lemma: Result, anyhow::Error> = proof + .lemma() + .iter() + .map(|oh| { + oh.0.ok_or_else(|| anyhow::anyhow!("Cannot convert null OptionalHash to H256")) + }) + .collect(); + + Proof::new(lemma?, proof.path().to_vec()) + } + + /// Convert a range proof of OptionalHash to a range proof of H256 + pub fn convert_range_proof_to_h256( + proof: &RangeProof, + ) -> Result, anyhow::Error> { + Ok(RangeProof { + left_proof: Self::convert_proof_to_h256(&proof.left_proof)?, + right_proof: Self::convert_proof_to_h256(&proof.right_proof)?, + }) + } + + /// Generate a proof and convert it to H256 + pub fn gen_proof_h256(&self, leaf_index: usize) -> Result, anyhow::Error> { + let proof = self.gen_proof(leaf_index)?; + Self::convert_proof_to_h256(&proof) + } + + /// Generate a range proof and convert it to H256 + pub fn gen_range_proof_h256( + &self, + start_index: usize, + end_index: usize, + ) -> Result, anyhow::Error> { + let proof = self.gen_range_proof(start_index, end_index)?; + Self::convert_range_proof_to_h256(&proof) + } + + /// Get the root as H256 (unwraps the OptionalHash) + pub fn root_h256(&self) -> H256 { + self.root().unwrap() + } +} + pub struct AppendMerkleTree> { /// Keep all the nodes in the latest version. `layers[0]` is the layer of leaves. node_manager: NodeManager, @@ -148,7 +199,7 @@ impl> AppendMerkleTree { pub fn append(&mut self, new_leaf: E) { let start_time = Instant::now(); - if new_leaf == E::null() { + if new_leaf.is_null() { // appending null is not allowed. return; } @@ -162,7 +213,7 @@ impl> AppendMerkleTree { pub fn append_list(&mut self, leaf_list: Vec) { let start_time = Instant::now(); - if leaf_list.contains(&E::null()) { + if leaf_list.iter().any(|leaf| leaf.is_null()) { // appending null is not allowed. return; } @@ -181,7 +232,7 @@ impl> AppendMerkleTree { /// TODO: Optimize to avoid storing the `null` nodes? pub fn append_subtree(&mut self, subtree_depth: usize, subtree_root: E) -> Result<()> { let start_time = Instant::now(); - if subtree_root == E::null() { + if subtree_root.is_null() { // appending null is not allowed. bail!("subtree_root is null"); } @@ -197,7 +248,7 @@ impl> AppendMerkleTree { pub fn append_subtree_list(&mut self, subtree_list: Vec<(usize, E)>) -> Result<()> { let start_time = Instant::now(); - if subtree_list.iter().any(|(_, root)| root == &E::null()) { + if subtree_list.iter().any(|(_, root)| root.is_null()) { // appending null is not allowed. bail!("subtree_list contains null"); } @@ -217,7 +268,7 @@ impl> AppendMerkleTree { /// This is needed if our merkle-tree in memory only keeps intermediate nodes instead of real leaves. pub fn update_last(&mut self, updated_leaf: E) { let start_time = Instant::now(); - if updated_leaf == E::null() { + if updated_leaf.is_null() { // updating to null is not allowed. return; } @@ -237,9 +288,9 @@ impl> AppendMerkleTree { /// Panics if the leaf is already set and different or the index is out of range. /// TODO: Batch computing intermediate nodes. pub fn fill_leaf(&mut self, index: usize, leaf: E) { - if leaf == E::null() { + if leaf.is_null() { // fill leaf with null is not allowed. - } else if self.node(0, index) == E::null() { + } else if self.node(0, index).is_null() { self.node_manager.start_transaction(); self.update_node(0, index, leaf); self.recompute_after_fill_leaves(index, index + 1); @@ -332,7 +383,7 @@ impl> AppendMerkleTree { // skip padding node. continue; } - if self.node(i, position) == E::null() { + if self.node(i, position).is_null() { self.update_node(i, position, data.clone()); updated_nodes.push((i, position, data)) } else if self.node(i, position) != data { @@ -357,7 +408,7 @@ impl> AppendMerkleTree { if position >= self.leaves() { bail!("Out of bound: position={} end={}", position, self.leaves()); } - if self.node(0, position) != E::null() { + if !self.node(0, position).is_null() { Ok(Some(self.node(0, position))) } else { // The leaf hash is unknown. @@ -472,7 +523,7 @@ impl> AppendMerkleTree { // Note that if we are recompute a range of an existing tree, // we do not need to keep these possibly null parent. This is only saved // for the case of constructing a new tree from the leaves. - let parent = if *left == E::null() || *right == E::null() { + let parent = if left.is_null() || right.is_null() { E::null() } else { A::parent(left, right) @@ -483,7 +534,7 @@ impl> AppendMerkleTree { assert_eq!(chunk.len(), 1); let r = &chunk[0]; // Same as above. - let parent = if *r == E::null() { + let parent = if r.is_null() { E::null() } else { A::parent_single(r, height + self.leaf_height) @@ -501,8 +552,8 @@ impl> AppendMerkleTree { match parent_index.cmp(&self.layer_len(height + 1)) { Ordering::Less => { // We do not overwrite with null. - if parent != E::null() { - if self.node(height + 1, parent_index) == E::null() + if !parent.is_null() { + if self.node(height + 1, parent_index).is_null() // The last node in a layer can be updated. || (self.node(height + 1, parent_index) != parent && parent_index == self.layer_len(height + 1) - 1) @@ -741,7 +792,7 @@ impl<'a, E: HashElement> MerkleTreeRead for HistoryTree<'a, E> { type E = E; fn node(&self, layer: usize, index: usize) -> Self::E { match self.delta_nodes.get(layer, index).expect("range checked") { - Some(node) if *node != E::null() => node.clone(), + Some(node) if !node.is_null() => node.clone(), _ => self .node_manager .get_node(layer, index) @@ -798,7 +849,7 @@ macro_rules! ensure_eq { #[cfg(test)] mod tests { - use crate::merkle_tree::MerkleTreeRead; + use crate::merkle_tree::{MerkleTreeRead, OptionalHash}; use crate::sha3::Sha3Algorithm; use crate::AppendMerkleTree; @@ -812,21 +863,30 @@ mod tests { for _ in 0..entry_len { data.push(H256::random()); } - let mut merkle = - AppendMerkleTree::::new(vec![H256::zero()], 0, None); - merkle.append_list(data.clone()); + let mut merkle = AppendMerkleTree::::new( + vec![OptionalHash::some(H256::zero())], + 0, + None, + ); + merkle.append_list(data.clone().into_iter().map(OptionalHash::some).collect()); merkle.commit(Some(0)); verify(&data, &mut merkle); data.push(H256::random()); - merkle.append(*data.last().unwrap()); + merkle.append(OptionalHash::some(*data.last().unwrap())); merkle.commit(Some(1)); verify(&data, &mut merkle); for _ in 0..6 { data.push(H256::random()); } - merkle.append_list(data[data.len() - 6..].to_vec()); + merkle.append_list( + data[data.len() - 6..] + .iter() + .copied() + .map(OptionalHash::some) + .collect(), + ); merkle.commit(Some(2)); verify(&data, &mut merkle); } @@ -840,9 +900,12 @@ mod tests { for _ in 0..entry_len { data.push(H256::random()); } - let mut merkle = - AppendMerkleTree::::new(vec![H256::zero()], 0, None); - merkle.append_list(data.clone()); + let mut merkle = AppendMerkleTree::::new( + vec![OptionalHash::some(H256::zero())], + 0, + None, + ); + merkle.append_list(data.clone().into_iter().map(OptionalHash::some).collect()); merkle.commit(Some(0)); for i in (0..data.len()).step_by(6) { @@ -850,12 +913,17 @@ mod tests { let range_proof = merkle.gen_range_proof(i + 1, end + 1).unwrap(); let mut new_data = Vec::new(); for _ in 0..3 { - new_data.push(H256::random()); + new_data.push(OptionalHash::some(H256::random())); } merkle.append_list(new_data); let seq = i as u64 / 6 + 1; merkle.commit(Some(seq)); - let r = range_proof.validate::(&data[i..end], i + 1); + let optional_data: Vec = data[i..end] + .iter() + .copied() + .map(OptionalHash::some) + .collect(); + let r = range_proof.validate::(&optional_data, i + 1); assert!(r.is_ok(), "{:?}", r); merkle.fill_with_range_proof(range_proof).unwrap(); } @@ -865,7 +933,11 @@ mod tests { #[test] fn test_proof_at_version() { let n = [2, 255, 256, 257]; - let mut merkle = AppendMerkleTree::::new(vec![H256::zero()], 0, None); + let mut merkle = AppendMerkleTree::::new( + vec![OptionalHash::some(H256::zero())], + 0, + None, + ); let mut start_pos = 0; for (tx_seq, &entry_len) in n.iter().enumerate() { @@ -873,7 +945,7 @@ mod tests { for _ in 0..entry_len { data.push(H256::random()); } - merkle.append_list(data.clone()); + merkle.append_list(data.clone().into_iter().map(OptionalHash::some).collect()); merkle.commit(Some(tx_seq as u64)); for i in (0..data.len()).step_by(6) { let end = std::cmp::min(start_pos + i + 3, data.len()); @@ -882,7 +954,12 @@ mod tests { .unwrap() .gen_range_proof(start_pos + i + 1, start_pos + end + 1) .unwrap(); - let r = range_proof.validate::(&data[i..end], start_pos + i + 1); + let optional_data: Vec = data[i..end] + .iter() + .copied() + .map(OptionalHash::some) + .collect(); + let r = range_proof.validate::(&optional_data, start_pos + i + 1); assert!(r.is_ok(), "{:?}", r); merkle.fill_with_range_proof(range_proof).unwrap(); } @@ -891,16 +968,21 @@ mod tests { } } - fn verify(data: &[H256], merkle: &mut AppendMerkleTree) { + fn verify(data: &[H256], merkle: &mut AppendMerkleTree) { for (i, item) in data.iter().enumerate() { let proof = merkle.gen_proof(i + 1).unwrap(); - let r = merkle.validate(&proof, item, i + 1); + let r = merkle.validate(&proof, &OptionalHash::some(*item), i + 1); assert!(matches!(r, Ok(true)), "{:?}", r); } for i in (0..data.len()).step_by(6) { let end = std::cmp::min(i + 3, data.len()); let range_proof = merkle.gen_range_proof(i + 1, end + 1).unwrap(); - let r = range_proof.validate::(&data[i..end], i + 1); + let optional_data: Vec = data[i..end] + .iter() + .copied() + .map(OptionalHash::some) + .collect(); + let r = range_proof.validate::(&optional_data, i + 1); assert!(r.is_ok(), "{:?}", r); merkle.fill_with_range_proof(range_proof).unwrap(); } diff --git a/common/append_merkle/src/merkle_tree.rs b/common/append_merkle/src/merkle_tree.rs index eac083f..9ed93fd 100644 --- a/common/append_merkle/src/merkle_tree.rs +++ b/common/append_merkle/src/merkle_tree.rs @@ -8,6 +8,175 @@ use std::fmt::Debug; use std::hash::Hash; use tracing::trace; +/// A wrapper around Option that properly handles null hashes +#[derive(Clone, Debug, Eq, PartialEq, Hash)] +pub struct OptionalHash(pub Option); + +impl OptionalHash { + pub fn some(hash: H256) -> Self { + OptionalHash(Some(hash)) + } + + pub fn none() -> Self { + OptionalHash(None) + } + + pub fn is_some(&self) -> bool { + self.0.is_some() + } + + pub fn is_none(&self) -> bool { + self.0.is_none() + } + + pub fn unwrap(&self) -> H256 { + self.0.unwrap() + } + + pub fn unwrap_or(&self, default: H256) -> H256 { + self.0.unwrap_or(default) + } + + pub fn as_ref(&self) -> Option<&H256> { + self.0.as_ref() + } + + pub fn to_h256_or_zero(&self) -> H256 { + self.0.unwrap_or_else(H256::zero) + } + + /// Convert a Proof to Proof + pub fn convert_proof_to_h256(proof: Proof) -> Result, anyhow::Error> { + let lemma: Vec = proof + .lemma() + .iter() + .map(|oh| oh.to_h256_or_zero()) + .collect(); + let path = proof.path().to_vec(); + Proof::new(lemma, path) + } + + /// Convert a Proof to Proof + pub fn convert_proof_from_h256( + proof: Proof, + ) -> Result, anyhow::Error> { + let lemma: Vec = proof + .lemma() + .iter() + .map(|h| OptionalHash::some(*h)) + .collect(); + let path = proof.path().to_vec(); + Proof::new(lemma, path) + } + + /// Convert a RangeProof to RangeProof + pub fn convert_range_proof_from_h256( + range_proof: RangeProof, + ) -> Result, anyhow::Error> { + Ok(RangeProof { + left_proof: Self::convert_proof_from_h256(range_proof.left_proof)?, + right_proof: Self::convert_proof_from_h256(range_proof.right_proof)?, + }) + } +} + +// Add From conversions for easier usage +impl From for OptionalHash { + fn from(hash: H256) -> Self { + OptionalHash::some(hash) + } +} + +impl From> for OptionalHash { + fn from(opt: Option) -> Self { + OptionalHash(opt) + } +} + +impl From for Option { + fn from(opt_hash: OptionalHash) -> Self { + opt_hash.0 + } +} + +impl AsRef<[u8]> for OptionalHash { + fn as_ref(&self) -> &[u8] { + match &self.0 { + Some(hash) => hash.as_ref(), + None => &[0u8; 32], // Return zeros for null hash + } + } +} + +impl AsMut<[u8]> for OptionalHash { + fn as_mut(&mut self) -> &mut [u8] { + if self.0.is_none() { + self.0 = Some(H256::zero()); + } + self.0.as_mut().unwrap().as_mut() + } +} + +impl Encode for OptionalHash { + fn is_ssz_fixed_len() -> bool { + true + } + + fn ssz_fixed_len() -> usize { + 33 // 1 byte for Some/None + 32 bytes for hash + } + + fn ssz_bytes_len(&self) -> usize { + 33 + } + + fn ssz_append(&self, buf: &mut Vec) { + match &self.0 { + Some(hash) => { + buf.push(1); // Some discriminant + hash.ssz_append(buf); + } + None => { + buf.push(0); // None discriminant + buf.extend_from_slice(&[0u8; 32]); + } + } + } +} + +impl Decode for OptionalHash { + fn is_ssz_fixed_len() -> bool { + true + } + + fn ssz_fixed_len() -> usize { + 33 + } + + fn from_ssz_bytes(bytes: &[u8]) -> Result { + if bytes.len() != 33 { + return Err(ssz::DecodeError::InvalidByteLength { + len: bytes.len(), + expected: 33, + }); + } + + match bytes[0] { + 0 => Ok(OptionalHash::none()), + 1 => { + let hash = H256::from_ssz_bytes(&bytes[1..])?; + Ok(OptionalHash::some(hash)) + } + _ => Err(ssz::DecodeError::BytesInvalid( + "Invalid discriminant for OptionalHash".to_string(), + )), + } + } +} + +unsafe impl Send for OptionalHash {} +unsafe impl Sync for OptionalHash {} + pub trait HashElement: Clone + Debug + Eq + Hash + AsRef<[u8]> + AsMut<[u8]> + Decode + Encode + Send + Sync { @@ -18,13 +187,28 @@ pub trait HashElement: } } +impl HashElement for OptionalHash { + fn end_pad(height: usize) -> Self { + OptionalHash::some(ZERO_HASHES[height]) + } + + fn null() -> Self { + OptionalHash::none() + } + + fn is_null(&self) -> bool { + self.is_none() + } +} + +// Keep the H256 implementation for backward compatibility impl HashElement for H256 { fn end_pad(height: usize) -> Self { ZERO_HASHES[height] } fn null() -> Self { - H256::repeat_byte(1) + H256::zero() // Use all zeros instead of 0x0101... to avoid collision } } @@ -70,7 +254,7 @@ pub trait MerkleTreeRead { self.leaves() ); } - if self.node(0, leaf_index) == Self::E::null() { + if self.node(0, leaf_index).is_null() { bail!("Not ready to generate proof for leaf_index={}", leaf_index); } if self.height() == 1 { @@ -102,7 +286,7 @@ pub trait MerkleTreeRead { index_in_layer >>= 1; } lemma.push(self.root()); - if lemma.contains(&Self::E::null()) { + if lemma.iter().any(|e| e.is_null()) { bail!( "Not enough data to generate proof, lemma={:?} path={:?}", lemma, diff --git a/common/append_merkle/src/sha3.rs b/common/append_merkle/src/sha3.rs index c1008ec..f5d38b7 100644 --- a/common/append_merkle/src/sha3.rs +++ b/common/append_merkle/src/sha3.rs @@ -1,4 +1,4 @@ -use crate::merkle_tree::ZERO_HASHES; +use crate::merkle_tree::{OptionalHash, ZERO_HASHES}; use crate::{Algorithm, HashElement}; use ethereum_types::H256; use once_cell::sync::Lazy; @@ -50,3 +50,16 @@ impl Algorithm for Sha3Algorithm { Self::leaf_raw(data) } } + +impl Algorithm for Sha3Algorithm { + fn parent(left: &OptionalHash, right: &OptionalHash) -> OptionalHash { + match (&left.0, &right.0) { + (Some(l), Some(r)) => OptionalHash::some(Self::parent(l, r)), + _ => OptionalHash::none(), + } + } + + fn leaf(data: &[u8]) -> OptionalHash { + OptionalHash::some(Self::leaf(data)) + } +} diff --git a/node/shared_types/src/lib.rs b/node/shared_types/src/lib.rs index 54f773d..dbae705 100644 --- a/node/shared_types/src/lib.rs +++ b/node/shared_types/src/lib.rs @@ -2,7 +2,7 @@ mod proof; use anyhow::{anyhow, bail, Error}; use append_merkle::{ - AppendMerkleTree, Proof as RawProof, RangeProof as RawRangeProof, Sha3Algorithm, + AppendMerkleTree, OptionalHash, Proof as RawProof, RangeProof as RawRangeProof, Sha3Algorithm, }; use ethereum_types::{Address, H256, U256}; use merkle_light::merkle::MerkleTree; @@ -32,7 +32,7 @@ pub type DataRoot = H256; pub type FlowProof = RawProof; pub type FlowRangeProof = RawRangeProof; -pub type Merkle = AppendMerkleTree; +pub type Merkle = AppendMerkleTree; // Each chunk is 32 bytes. pub const CHUNK_SIZE: usize = 256; diff --git a/node/storage/src/log_store/flow_store.rs b/node/storage/src/log_store/flow_store.rs index d630e2c..2ea5fde 100644 --- a/node/storage/src/log_store/flow_store.rs +++ b/node/storage/src/log_store/flow_store.rs @@ -12,7 +12,7 @@ use crate::log_store::{ use crate::{try_option, ZgsKeyValueDB}; use any::Any; use anyhow::{anyhow, bail, Result}; -use append_merkle::{MerkleTreeRead, NodeDatabase, NodeTransaction}; +use append_merkle::{MerkleTreeRead, NodeDatabase, NodeTransaction, OptionalHash}; use itertools::Itertools; use kvdb::DBTransaction; use parking_lot::RwLock; @@ -72,7 +72,8 @@ impl FlowStore { batch_index ) })?; - merkle.gen_proof(sector_index) + let optional_proof = merkle.gen_proof(sector_index)?; + OptionalHash::convert_proof_to_h256(optional_proof) } pub fn delete_batch_list(&self, batch_list: &[u64]) -> Result<()> { @@ -649,3 +650,64 @@ impl NodeTransaction for NodeDBTransaction { self } } + +// Adapter implementation for OptionalHash that delegates to the H256 implementation +impl NodeDatabase for FlowDBStore { + fn get_node(&self, layer: usize, pos: usize) -> Result> { + Ok(self.get_node(layer, pos)?.map(OptionalHash::some)) + } + + fn get_layer_size(&self, layer: usize) -> Result> { + // Layer size is the same regardless of hash type + >::get_layer_size(self, layer) + } + + fn start_transaction(&self) -> Box> { + Box::new(OptionalHashNodeDBTransaction(self.start_transaction())) + } + + fn commit(&self, tx: Box>) -> Result<()> { + let h256_tx = tx + .into_any() + .downcast::() + .map_err(|_| anyhow::anyhow!("Failed to downcast OptionalHashNodeDBTransaction"))?; + self.commit(h256_tx.0) + } +} + +// Wrapper for NodeTransaction that delegates to NodeTransaction +pub struct OptionalHashNodeDBTransaction(Box>); + +impl NodeTransaction for OptionalHashNodeDBTransaction { + fn save_node(&mut self, layer: usize, pos: usize, node: &OptionalHash) { + self.0.save_node(layer, pos, &node.to_h256_or_zero()); + } + + fn save_node_list(&mut self, nodes: &[(usize, usize, &OptionalHash)]) { + let h256_nodes: Vec<(usize, usize, DataRoot)> = nodes + .iter() + .map(|(layer, pos, oh)| (*layer, *pos, oh.to_h256_or_zero())) + .collect(); + let h256_node_refs: Vec<(usize, usize, &DataRoot)> = h256_nodes + .iter() + .map(|(layer, pos, h)| (*layer, *pos, h)) + .collect(); + self.0.save_node_list(&h256_node_refs); + } + + fn remove_node_list(&mut self, nodes: &[(usize, usize)]) { + self.0.remove_node_list(nodes); + } + + fn save_layer_size(&mut self, layer: usize, size: usize) { + self.0.save_layer_size(layer, size); + } + + fn remove_layer_size(&mut self, layer: usize) { + self.0.remove_layer_size(layer); + } + + fn into_any(self: Box) -> Box { + Box::new(self) + } +} diff --git a/node/storage/src/log_store/load_chunk/mod.rs b/node/storage/src/log_store/load_chunk/mod.rs index 9d60d25..95c002c 100644 --- a/node/storage/src/log_store/load_chunk/mod.rs +++ b/node/storage/src/log_store/load_chunk/mod.rs @@ -204,9 +204,9 @@ impl EntryBatch { } } } - Ok(Some( - try_option!(self.to_merkle_tree(is_first_chunk)?).root(), - )) + Ok(try_option!(self.to_merkle_tree(is_first_chunk)?) + .root() + .into()) } pub fn submit_seal_result(&mut self, answer: SealAnswer) -> Result<()> { @@ -243,7 +243,7 @@ impl EntryBatch { pub fn to_merkle_tree(&self, is_first_chunk: bool) -> Result> { let initial_leaves = if is_first_chunk { - vec![H256::zero()] + vec![H256::zero().into()] } else { vec![] }; @@ -256,7 +256,7 @@ impl EntryBatch { ); merkle.append_list(data_to_merkle_leaves(&leaf_data).expect("aligned")); } - merkle.append_subtree(subtree.subtree_height, subtree.root)?; + merkle.append_subtree(subtree.subtree_height, subtree.root.into())?; } if merkle.leaves() != SECTORS_PER_LOAD { let leaf_data = try_option!( diff --git a/node/storage/src/log_store/log_manager.rs b/node/storage/src/log_store/log_manager.rs index 516092e..e9d837f 100644 --- a/node/storage/src/log_store/log_manager.rs +++ b/node/storage/src/log_store/log_manager.rs @@ -9,7 +9,7 @@ use crate::log_store::{ }; use crate::{try_option, ZgsKeyValueDB}; use anyhow::{anyhow, bail, Result}; -use append_merkle::{Algorithm, MerkleTreeRead, Sha3Algorithm}; +use append_merkle::{Algorithm, MerkleTreeRead, OptionalHash, Sha3Algorithm}; use ethereum_types::H256; use kvdb_rocksdb::{Database, DatabaseConfig}; use merkle_light::merkle::{log2_pow2, MerkleTree}; @@ -55,13 +55,11 @@ const PAD_DELAY: Duration = Duration::from_secs(2); // Process at most 1M entries (256MB) pad data at a time. const PAD_MAX_SIZE: usize = 1 << 20; -static PAD_SEGMENT_ROOT: Lazy = Lazy::new(|| { - Merkle::new( - data_to_merkle_leaves(&[0; ENTRY_SIZE * PORA_CHUNK_SIZE]).unwrap(), - 0, - None, - ) - .root() +static PAD_SEGMENT_ROOT: Lazy = Lazy::new(|| { + let h256_leaves = data_to_merkle_leaves(&[0; ENTRY_SIZE * PORA_CHUNK_SIZE]).unwrap(); + let optional_leaves: Vec = + h256_leaves.into_iter().map(OptionalHash::from).collect(); + Merkle::new(optional_leaves, 0, None).root() }); pub struct UpdateFlowMessage { pub pad_data: usize, @@ -130,7 +128,8 @@ impl MerkleManager { fn try_initialize(&mut self, flow_store: &FlowStore) -> Result<()> { if self.pora_chunks_merkle.leaves() == 0 && self.last_chunk_merkle.leaves() == 0 { - self.last_chunk_merkle.append(H256::zero()); + self.last_chunk_merkle + .append(OptionalHash::some(H256::zero())); self.pora_chunks_merkle .update_last(self.last_chunk_merkle.root()); } else if self.last_chunk_merkle.leaves() != 0 { @@ -222,9 +221,17 @@ impl LogStoreChunkWrite for LogManager { self.append_entries(flow_entry_array, &mut merkle)?; if let Some(file_proof) = maybe_file_proof { + // Convert H256 proof to OptionalHash proof + let optional_proof = OptionalHash::convert_proof_from_h256(file_proof)?; + // Convert H256 merkle nodes to OptionalHash merkle nodes + let optional_nodes: Vec<(usize, OptionalHash)> = tx + .merkle_nodes + .into_iter() + .map(|(depth, hash)| (depth, OptionalHash::some(hash))) + .collect(); merkle.pora_chunks_merkle.fill_with_file_proof( - file_proof, - tx.merkle_nodes, + optional_proof, + optional_nodes, tx.start_entry_index, )?; } @@ -402,9 +409,9 @@ impl LogStoreWrite for LogManager { // `merkle` is used in `validate_range_proof`. let mut merkle = self.merkle.write(); if valid { - merkle - .pora_chunks_merkle - .fill_with_range_proof(data.proof.clone())?; + merkle.pora_chunks_merkle.fill_with_range_proof( + OptionalHash::convert_range_proof_from_h256(data.proof.clone())?, + )?; } Ok(valid) } @@ -601,7 +608,7 @@ impl LogStoreRead for LogManager { let tx = self .get_tx_by_seq_number(tx_seq)? .ok_or_else(|| anyhow!("tx missing"))?; - let leaves = data_to_merkle_leaves(&data.chunks.data)?; + let leaves = data_to_merkle_leaves_h256(&data.chunks.data)?; data.proof.validate::( &leaves, (data.chunks.start_index + tx.start_entry_index) as usize, @@ -610,7 +617,7 @@ impl LogStoreRead for LogManager { .merkle .read_recursive() .pora_chunks_merkle - .check_root(&data.proof.root())) + .check_root(&data.proof.root().into())) } fn get_sync_progress(&self) -> Result> { @@ -650,7 +657,7 @@ impl LogStoreRead for LogManager { fn get_context(&self) -> crate::error::Result<(DataRoot, u64)> { let merkle = self.merkle.read_recursive(); Ok(( - merkle.pora_chunks_merkle.root(), + merkle.pora_chunks_merkle.root().to_h256_or_zero(), merkle.last_chunk_start_index() + merkle.last_chunk_merkle.leaves() as u64, )) } @@ -835,7 +842,9 @@ impl LogManager { None => self.gen_proof_at_version(flow_index, None), Some(root) => { let merkle = self.merkle.read_recursive(); - let tx_seq = merkle.pora_chunks_merkle.tx_seq_at_root(&root)?; + let tx_seq = merkle + .pora_chunks_merkle + .tx_seq_at_root(&OptionalHash::from(root))?; self.gen_proof_at_version(flow_index, Some(tx_seq)) } } @@ -849,11 +858,15 @@ impl LogManager { let merkle = self.merkle.read_recursive(); let seg_index = sector_to_segment(flow_index); let top_proof = match maybe_tx_seq { - None => merkle.pora_chunks_merkle.gen_proof(seg_index)?, - Some(tx_seq) => merkle - .pora_chunks_merkle - .at_version(tx_seq)? - .gen_proof(seg_index)?, + None => OptionalHash::convert_proof_to_h256( + merkle.pora_chunks_merkle.gen_proof(seg_index)?, + )?, + Some(tx_seq) => OptionalHash::convert_proof_to_h256( + merkle + .pora_chunks_merkle + .at_version(tx_seq)? + .gen_proof(seg_index)?, + )?, }; // TODO(zz): Maybe we can decide that all proofs are at the PoRA chunk level, so @@ -870,13 +883,17 @@ impl LogManager { .gen_proof_in_batch(seg_index, flow_index as usize % PORA_CHUNK_SIZE)? } else { match maybe_tx_seq { - None => merkle - .last_chunk_merkle - .gen_proof(flow_index as usize % PORA_CHUNK_SIZE)?, - Some(tx_version) => merkle - .last_chunk_merkle - .at_version(tx_version)? - .gen_proof(flow_index as usize % PORA_CHUNK_SIZE)?, + None => OptionalHash::convert_proof_to_h256( + merkle + .last_chunk_merkle + .gen_proof(flow_index as usize % PORA_CHUNK_SIZE)?, + )?, + Some(tx_version) => OptionalHash::convert_proof_to_h256( + merkle + .last_chunk_merkle + .at_version(tx_version)? + .gen_proof(flow_index as usize % PORA_CHUNK_SIZE)?, + )?, } }; entry_proof(&top_proof, &sub_proof) @@ -902,7 +919,7 @@ impl LogManager { if merkle.last_chunk_merkle.leaves() + subtree_size <= PORA_CHUNK_SIZE { merkle .last_chunk_merkle - .append_subtree(subtree_depth, subtree_root)?; + .append_subtree(subtree_depth, OptionalHash::some(subtree_root))?; if merkle.last_chunk_merkle.leaves() == subtree_size { // `last_chunk_merkle` was empty, so this is a new leaf in the top_tree. merkle @@ -924,9 +941,10 @@ impl LogManager { // the chunks boundary. assert_eq!(merkle.last_chunk_merkle.leaves(), 0); assert!(subtree_size >= PORA_CHUNK_SIZE); - merkle - .pora_chunks_merkle - .append_subtree(subtree_depth - log2_pow2(PORA_CHUNK_SIZE), subtree_root)?; + merkle.pora_chunks_merkle.append_subtree( + subtree_depth - log2_pow2(PORA_CHUNK_SIZE), + subtree_root.into(), + )?; } } @@ -961,9 +979,11 @@ impl LogManager { let mut completed_chunk_index = None; if pad_data.len() < last_chunk_pad { is_full_empty = false; - merkle - .last_chunk_merkle - .append_list(data_to_merkle_leaves(&pad_data)?); + let pad_leaves: Vec = data_to_merkle_leaves(&pad_data)? + .into_iter() + .map(OptionalHash::from) + .collect(); + merkle.last_chunk_merkle.append_list(pad_leaves); merkle .pora_chunks_merkle .update_last(merkle.last_chunk_merkle.root()); @@ -971,9 +991,12 @@ impl LogManager { if last_chunk_pad != 0 { is_full_empty = false; // Pad the last chunk. - merkle - .last_chunk_merkle - .append_list(data_to_merkle_leaves(&pad_data[..last_chunk_pad])?); + let last_chunk_leaves: Vec = + data_to_merkle_leaves(&pad_data[..last_chunk_pad])? + .into_iter() + .map(OptionalHash::from) + .collect(); + merkle.last_chunk_merkle.append_list(last_chunk_leaves); merkle .pora_chunks_merkle .update_last(merkle.last_chunk_merkle.root()); @@ -983,7 +1006,7 @@ impl LogManager { // Pad with more complete chunks. let mut start_index = last_chunk_pad / ENTRY_SIZE; while pad_data.len() >= (start_index + PORA_CHUNK_SIZE) * ENTRY_SIZE { - merkle.pora_chunks_merkle.append(*PAD_SEGMENT_ROOT); + merkle.pora_chunks_merkle.append(PAD_SEGMENT_ROOT.clone()); start_index += PORA_CHUNK_SIZE; } assert_eq!(pad_data.len(), start_index * ENTRY_SIZE); @@ -1068,7 +1091,7 @@ impl LogManager { if chunk_index < merkle.pora_chunks_merkle.leaves() as u64 { merkle .pora_chunks_merkle - .fill_leaf(chunk_index as usize, chunk_root); + .fill_leaf(chunk_index as usize, OptionalHash::some(chunk_root)); } else { // TODO(zz): This assumption may be false in the future. unreachable!("We always insert tx nodes before put_chunks"); @@ -1216,7 +1239,7 @@ impl LogManager { let mut to_insert_subtrees = Vec::new(); let mut start_index = 0; for (subtree_height, root) in subtree_list { - to_insert_subtrees.push((start_index, subtree_height, root)); + to_insert_subtrees.push((start_index, subtree_height, root.to_h256_or_zero())); start_index += 1 << (subtree_height - 1); } self.flow_store @@ -1264,14 +1287,14 @@ macro_rules! try_option { /// This should be called with input checked. pub fn sub_merkle_tree(leaf_data: &[u8]) -> Result { Ok(FileMerkleTree::new( - data_to_merkle_leaves(leaf_data)? + data_to_merkle_leaves_h256(leaf_data)? .into_iter() .map(|h| h.0) .collect::>(), )) } -pub fn data_to_merkle_leaves(leaf_data: &[u8]) -> Result> { +pub fn data_to_merkle_leaves(leaf_data: &[u8]) -> Result> { let start_time = Instant::now(); if leaf_data.len() % ENTRY_SIZE != 0 { bail!("merkle_tree: mismatched data size"); @@ -1294,6 +1317,20 @@ pub fn data_to_merkle_leaves(leaf_data: &[u8]) -> Result> { Ok(r) } +/// Convert Vec to Vec for compatibility with existing code +pub fn optional_hash_to_h256_vec(optional_hashes: Vec) -> Vec { + optional_hashes + .into_iter() + .map(|oh| oh.to_h256_or_zero()) + .collect() +} + +/// Convenience function that combines data_to_merkle_leaves and conversion to H256 +pub fn data_to_merkle_leaves_h256(leaf_data: &[u8]) -> Result> { + let optional_hashes = data_to_merkle_leaves(leaf_data)?; + Ok(optional_hash_to_h256_vec(optional_hashes)) +} + pub fn bytes_to_entries(size_bytes: u64) -> u64 { if size_bytes % ENTRY_SIZE as u64 == 0 { size_bytes / ENTRY_SIZE as u64 diff --git a/node/storage/src/log_store/tests.rs b/node/storage/src/log_store/tests.rs index 5b9f42c..f3f8e0c 100644 --- a/node/storage/src/log_store/tests.rs +++ b/node/storage/src/log_store/tests.rs @@ -1,9 +1,9 @@ use crate::log_store::log_manager::{ - data_to_merkle_leaves, sub_merkle_tree, tx_subtree_root_list_padded, LogConfig, LogManager, - PORA_CHUNK_SIZE, + data_to_merkle_leaves, data_to_merkle_leaves_h256, sub_merkle_tree, + tx_subtree_root_list_padded, LogConfig, LogManager, PORA_CHUNK_SIZE, }; use crate::log_store::{LogStoreChunkRead, LogStoreChunkWrite, LogStoreRead, LogStoreWrite}; -use append_merkle::{Algorithm, AppendMerkleTree, MerkleTreeRead, Sha3Algorithm}; +use append_merkle::{Algorithm, AppendMerkleTree, MerkleTreeRead, OptionalHash, Sha3Algorithm}; use ethereum_types::H256; use rand::random; use shared_types::{compute_padded_chunk_size, ChunkArray, Transaction, CHUNK_SIZE}; @@ -22,11 +22,26 @@ fn test_put_get() { data[i * CHUNK_SIZE] = random(); } let (padded_chunks, _) = compute_padded_chunk_size(data_size); - let mut merkle = AppendMerkleTree::::new(vec![H256::zero()], 0, None); - merkle.append_list(data_to_merkle_leaves(&LogManager::padding_raw(start_offset - 1)).unwrap()); + let mut merkle = AppendMerkleTree::::new( + vec![OptionalHash::some(H256::zero())], + 0, + None, + ); + let padding_leaves: Vec = + data_to_merkle_leaves(&LogManager::padding_raw(start_offset - 1)) + .unwrap() + .into_iter() + .map(OptionalHash::from) + .collect(); + merkle.append_list(padding_leaves); let mut data_padded = data.clone(); data_padded.append(&mut vec![0u8; CHUNK_SIZE]); - merkle.append_list(data_to_merkle_leaves(&data_padded).unwrap()); + let data_leaves: Vec = data_to_merkle_leaves(&data_padded) + .unwrap() + .into_iter() + .map(OptionalHash::from) + .collect(); + merkle.append_list(data_leaves); merkle.commit(Some(0)); let tx_merkle = sub_merkle_tree(&data).unwrap(); let tx = Transaction { @@ -78,16 +93,25 @@ fn test_put_get() { .unwrap() .unwrap(); assert_eq!(chunk_with_proof.chunk, chunk_array.chunk_at(i).unwrap()); + + // Create an H256 merkle tree for proof comparison + let mut h256_merkle = AppendMerkleTree::::new(vec![], 0, None); + h256_merkle.append_list( + data_to_merkle_leaves(&LogManager::padding_raw(start_offset - 1)).unwrap(), + ); + h256_merkle.append_list(data_to_merkle_leaves(&data_padded).unwrap()); + assert_eq!( chunk_with_proof.proof, - merkle.gen_proof(i + start_offset).unwrap() + OptionalHash::convert_proof_to_h256(h256_merkle.gen_proof(i + start_offset).unwrap()) + .unwrap() ); let r = chunk_with_proof.proof.validate::( &Sha3Algorithm::leaf(&chunk_with_proof.chunk.0), i + start_offset, ); assert!(r.is_ok(), "proof={:?} \n r={:?}", chunk_with_proof.proof, r); - assert!(merkle.check_root(&chunk_with_proof.proof.root())); + assert!(h256_merkle.check_root(&chunk_with_proof.proof.root().into())); } for i in (0..chunk_count).step_by(PORA_CHUNK_SIZE / 3) { let end = std::cmp::min(i + PORA_CHUNK_SIZE, chunk_count); @@ -102,7 +126,7 @@ fn test_put_get() { assert!(chunk_array_with_proof .proof .validate::( - &data_to_merkle_leaves(&chunk_array_with_proof.chunks.data).unwrap(), + &data_to_merkle_leaves_h256(&chunk_array_with_proof.chunks.data).unwrap(), i + start_offset ) .is_ok()); @@ -119,12 +143,16 @@ fn test_root() { } let mt = sub_merkle_tree(&data).unwrap(); println!("{:?} {}", mt.root(), hex::encode(mt.root())); - let append_mt = AppendMerkleTree::::new( - data_to_merkle_leaves(&data).unwrap(), + let append_mt = AppendMerkleTree::::new( + data_to_merkle_leaves(&data) + .unwrap() + .into_iter() + .map(OptionalHash::from) + .collect(), 0, None, ); - assert_eq!(mt.root(), append_mt.root().0); + assert_eq!(mt.root(), append_mt.root().to_h256_or_zero().0); } } diff --git a/node/storage/src/log_store/tx_store.rs b/node/storage/src/log_store/tx_store.rs index 45cb7b5..ee2b3c4 100644 --- a/node/storage/src/log_store/tx_store.rs +++ b/node/storage/src/log_store/tx_store.rs @@ -6,7 +6,7 @@ use crate::log_store::log_manager::{ use crate::log_store::metrics; use crate::{try_option, LogManager, ZgsKeyValueDB}; use anyhow::{anyhow, Result}; -use append_merkle::{AppendMerkleTree, MerkleTreeRead, Sha3Algorithm}; +use append_merkle::{AppendMerkleTree, MerkleTreeRead, OptionalHash, Sha3Algorithm}; use ethereum_types::H256; use merkle_light::merkle::log2_pow2; use shared_types::{DataRoot, Transaction}; @@ -329,7 +329,7 @@ impl TransactionStore { &self, pora_chunk_index: usize, mut tx_seq: u64, - ) -> Result> { + ) -> Result> { let last_chunk_start_index = pora_chunk_index as u64 * PORA_CHUNK_SIZE as u64; let mut tx_list = Vec::new(); // Find the first tx within the last chunk. @@ -384,9 +384,13 @@ impl TransactionStore { } let mut merkle = if last_chunk_start_index == 0 { // The first entry hash is initialized as zero. - AppendMerkleTree::::new_with_depth(vec![H256::zero()], 1, None) + AppendMerkleTree::::new_with_depth( + vec![H256::zero().into()], + 1, + None, + ) } else { - AppendMerkleTree::::new_with_depth( + AppendMerkleTree::::new_with_depth( vec![], log2_pow2(PORA_CHUNK_SIZE) + 1, None, @@ -400,9 +404,12 @@ impl TransactionStore { cmp::min(first_subtree, PORA_CHUNK_SIZE) - (merkle.leaves() % first_subtree); merkle.append_list(data_to_merkle_leaves(&LogManager::padding_raw(pad_len))?); } - // Since we are building the last merkle with a given last tx_seq, it's ensured - // that appending subtrees will not go beyond the max size. - merkle.append_subtree_list(subtree_list)?; + // Convert H256 to OptionalHash for append_subtree_list + let subtree_list_optional_hash = subtree_list + .into_iter() + .map(|(depth, hash)| (depth, hash.into())) + .collect(); + merkle.append_subtree_list(subtree_list_optional_hash)?; merkle.commit(Some(tx_seq)); } Ok(merkle)