Add encode/decode node test

firewood/src/merkle.rs

@@ -1257,9 +1257,10 @@ pub fn from_nibbles(nibbles: &[u8]) -> impl Iterator<Item = u8> + '_ {
 #[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<Vec<u8>>; 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);
+        }
+    }
 }

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<u8>);
 
@@ -53,6 +58,12 @@ impl<T: DeserializeOwned + AsRef<[u8]>> Encoded<T> {
     }
 }
 
+#[derive(Debug, Error)]
+pub enum Error {
+    #[error("decoding error")]
+    Decode(#[from] bincode::Error),
+}
+
 #[derive(PartialEq, Eq, Clone)]
 pub struct BranchNode {
     pub(super) chd: [Option<DiskAddress>; NBRANCH],
@@ -101,6 +112,26 @@ impl BranchNode {
         (only_chd, has_chd)
     }
 
+    pub fn decode(buf: &[u8]) -> Result<Self, Error> {
+        let mut items: Vec<Encoded<Vec<u8>>> = 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<Vec<u8>>; 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<S: ShaleStore<Node>>(&self, store: &S) -> Vec<u8> {
         // let mut stream = rlp::RlpStream::new_list(NBRANCH + 1);
         let mut list = <[Encoded<Vec<u8>>; NBRANCH + 1]>::default();
@@ -267,15 +298,15 @@ impl ExtNode {
     fn calc_eth_rlp<S: ShaleStore<Node>>(&self, store: &S) -> Vec<u8> {
         // let mut stream = rlp::RlpStream::new_list(2);
         let mut list = <[Encoded<Vec<u8>>; 2]>::default();
+        list[0] = Encoded::Data(
+            bincode::DefaultOptions::new()
+                .serialize(&from_nibbles(&self.0.encode(false)).collect::<Vec<_>>())
+                .unwrap(),
+        );
 
         if !self.1.is_null() {
             let mut r = store.get_item(self.1).unwrap();
             // stream.append(&from_nibbles(&self.0.encode(false)).collect::<Vec<_>>());
-            list[0] = Encoded::Data(
-                bincode::DefaultOptions::new()
-                    .serialize(&from_nibbles(&self.0.encode(false)).collect::<Vec<_>>())
-                    .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<S: ShaleStore<Node>>(&self, store: &S) -> Vec<u8> {
+    pub fn calc_eth_rlp<S: ShaleStore<Node>>(&self, store: &S) -> Vec<u8> {
         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<NodeType, Error> {
+        let items: Vec<Encoded<Vec<u8>>> = 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<u8> = 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<u8> = 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 {