diff --git a/firewood/src/merkle.rs b/firewood/src/merkle.rs index 6b0eef84f..cee523654 100644 --- a/firewood/src/merkle.rs +++ b/firewood/src/merkle.rs @@ -1257,9 +1257,10 @@ pub fn from_nibbles(nibbles: &[u8]) -> impl Iterator + '_ { #[cfg(test)] mod test { use super::*; - use shale::cached::PlainMem; + use shale::cached::{DynamicMem, PlainMem}; use shale::{CachedStore, Storable}; use std::ops::Deref; + use std::sync::Arc; use test_case::test_case; #[test_case(vec![0x12, 0x34, 0x56], vec![0x1, 0x2, 0x3, 0x4, 0x5, 0x6])] @@ -1382,4 +1383,87 @@ mod test { check(node); } } + #[test] + fn test_encode() { + const RESERVED: usize = 0x1000; + + let mut dm = shale::cached::DynamicMem::new(0x10000, 0); + let compact_header = DiskAddress::null(); + dm.write( + compact_header.into(), + &shale::to_dehydrated(&shale::compact::CompactSpaceHeader::new( + std::num::NonZeroUsize::new(RESERVED).unwrap(), + std::num::NonZeroUsize::new(RESERVED).unwrap(), + )) + .unwrap(), + ); + let compact_header = shale::StoredView::ptr_to_obj( + &dm, + compact_header, + shale::compact::CompactHeader::MSIZE, + ) + .unwrap(); + let mem_meta = Arc::new(dm); + let mem_payload = Arc::new(DynamicMem::new(0x10000, 0x1)); + + let cache = shale::ObjCache::new(1); + let space = + shale::compact::CompactSpace::new(mem_meta, mem_payload, compact_header, cache, 10, 16) + .expect("CompactSpace init fail"); + + let store = Box::new(space); + let merkle = Merkle::new(store); + + { + let chd = Node::new(NodeType::Leaf(LeafNode( + PartialPath(vec![0x1, 0x2, 0x3]), + Data(vec![0x4, 0x5]), + ))); + let chd_ref = merkle.new_node(chd.clone()).unwrap(); + let chd_rlp = chd_ref.get_eth_rlp(merkle.store.as_ref()); + let new_chd = Node::new(NodeType::decode(chd_rlp).unwrap()); + let new_chd_rlp = new_chd.get_eth_rlp(merkle.store.as_ref()); + assert_eq!(chd_rlp, new_chd_rlp); + + let mut chd_eth_rlp: [Option>; NBRANCH] = Default::default(); + chd_eth_rlp[0] = Some(chd_rlp.to_vec()); + let node = Node::new(NodeType::Branch(BranchNode { + chd: [None; NBRANCH], + value: Some(Data("value1".as_bytes().to_vec())), + chd_eth_rlp, + })); + + let node_ref = merkle.new_node(node.clone()).unwrap(); + + let r = node_ref.get_eth_rlp(merkle.store.as_ref()); + let new_node = Node::new(NodeType::decode(r).unwrap()); + let new_rlp = new_node.get_eth_rlp(merkle.store.as_ref()); + assert_eq!(r, new_rlp); + } + + { + let chd = Node::new(NodeType::Branch(BranchNode { + chd: [None; NBRANCH], + value: Some(Data("value1".as_bytes().to_vec())), + chd_eth_rlp: Default::default(), + })); + let chd_ref = merkle.new_node(chd.clone()).unwrap(); + let chd_rlp = chd_ref.get_eth_rlp(merkle.store.as_ref()); + let new_chd = Node::new(NodeType::decode(chd_rlp).unwrap()); + let new_chd_rlp = new_chd.get_eth_rlp(merkle.store.as_ref()); + assert_eq!(chd_rlp, new_chd_rlp); + + let node = Node::new(NodeType::Extension(ExtNode( + PartialPath(vec![0x1, 0x2, 0x3]), + DiskAddress::null(), + Some(chd_rlp.to_vec()), + ))); + let node_ref = merkle.new_node(node.clone()).unwrap(); + + let r = node_ref.get_eth_rlp(merkle.store.as_ref()); + let new_node = Node::new(NodeType::decode(r).unwrap()); + let new_rlp = new_node.get_eth_rlp(merkle.store.as_ref()); + assert_eq!(r, new_rlp); + } + } } diff --git a/firewood/src/merkle/node.rs b/firewood/src/merkle/node.rs index 44444d552..f8c0d9f7d 100644 --- a/firewood/src/merkle/node.rs +++ b/firewood/src/merkle/node.rs @@ -16,11 +16,16 @@ use std::{ }; use crate::merkle::to_nibble_array; +use crate::nibbles::Nibbles; +use thiserror::Error; use super::{from_nibbles, PartialPath, TrieHash, TRIE_HASH_LEN}; pub const NBRANCH: usize = 16; +const EXT_NODE_SIZE: usize = 2; +const BRANCH_NODE_SIZE: usize = 17; + #[derive(Debug, PartialEq, Eq, Clone)] pub struct Data(pub(super) Vec); @@ -53,6 +58,12 @@ impl> Encoded { } } +#[derive(Debug, Error)] +pub enum Error { + #[error("decoding error")] + Decode(#[from] bincode::Error), +} + #[derive(PartialEq, Eq, Clone)] pub struct BranchNode { pub(super) chd: [Option; NBRANCH], @@ -101,6 +112,26 @@ impl BranchNode { (only_chd, has_chd) } + pub fn decode(buf: &[u8]) -> Result { + let mut items: Vec>> = bincode::DefaultOptions::new().deserialize(buf)?; + + // we've already validated the size, that's why we can safely unwrap + let data = items.pop().unwrap().decode()?; + // Extract the value of the branch node and set to None if it's an empty Vec + let value = Some(data).filter(|data| !data.is_empty()); + + // Record rlp values of all children. + let mut chd_eth_rlp: [Option>; NBRANCH] = Default::default(); + + // we popped the last element, so their should only be NBRANCH items left + for (i, chd) in items.into_iter().enumerate() { + let data = chd.decode()?; + chd_eth_rlp[i] = Some(data).filter(|data| !data.is_empty()); + } + + Ok(BranchNode::new([None; NBRANCH], value, chd_eth_rlp)) + } + fn calc_eth_rlp>(&self, store: &S) -> Vec { // let mut stream = rlp::RlpStream::new_list(NBRANCH + 1); let mut list = <[Encoded>; NBRANCH + 1]>::default(); @@ -267,15 +298,15 @@ impl ExtNode { fn calc_eth_rlp>(&self, store: &S) -> Vec { // let mut stream = rlp::RlpStream::new_list(2); let mut list = <[Encoded>; 2]>::default(); + list[0] = Encoded::Data( + bincode::DefaultOptions::new() + .serialize(&from_nibbles(&self.0.encode(false)).collect::>()) + .unwrap(), + ); if !self.1.is_null() { let mut r = store.get_item(self.1).unwrap(); // stream.append(&from_nibbles(&self.0.encode(false)).collect::>()); - list[0] = Encoded::Data( - bincode::DefaultOptions::new() - .serialize(&from_nibbles(&self.0.encode(false)).collect::>()) - .unwrap(), - ); if r.get_eth_rlp_long(store) { // stream.append(&&(*r.get_root_hash(store))[..]); @@ -394,13 +425,47 @@ pub enum NodeType { } impl NodeType { - fn calc_eth_rlp>(&self, store: &S) -> Vec { + pub fn calc_eth_rlp>(&self, store: &S) -> Vec { match &self { NodeType::Leaf(n) => n.calc_eth_rlp(), NodeType::Extension(n) => n.calc_eth_rlp(store), NodeType::Branch(n) => n.calc_eth_rlp(store), } } + + pub fn decode(buf: &[u8]) -> Result { + let items: Vec>> = bincode::DefaultOptions::new() + .deserialize(dbg!(buf)) + .map_err(Error::Decode)?; + + match items.len() { + EXT_NODE_SIZE => { + let mut items = items.into_iter(); + let decoded_key: Vec = items.next().unwrap().decode()?; + + let decoded_key_nibbles = Nibbles::<0>::new(&decoded_key); + + let (cur_key_path, term) = + PartialPath::from_nibbles(decoded_key_nibbles.into_iter()); + let cur_key = cur_key_path.into_inner(); + let data: Vec = items.next().unwrap().decode()?; + + if term { + Ok(NodeType::Leaf(LeafNode::new(cur_key, data))) + } else { + Ok(NodeType::Extension(ExtNode::new( + cur_key, + DiskAddress::null(), + Some(data), + ))) + } + } + BRANCH_NODE_SIZE => Ok(NodeType::Branch(BranchNode::decode(buf)?)), + _ => Err(Error::Decode(Box::new(bincode::ErrorKind::Custom( + String::from(""), + )))), + } + } } impl Node {