From 854e801c0eb0bb1b021171437c87bddc9c1cd9a8 Mon Sep 17 00:00:00 2001 From: Richard Pringle Date: Tue, 24 Oct 2023 12:52:40 -0400 Subject: [PATCH] Insert test (#328) Signed-off-by: Richard Pringle Co-authored-by: Ron Kuris --- firewood/src/merkle.rs | 277 +++++++++------------------- firewood/src/merkle/node.rs | 105 +++++++++-- firewood/src/merkle/partial_path.rs | 18 ++ firewood/src/merkle/trie_hash.rs | 37 ++-- firewood/src/proof.rs | 2 +- 5 files changed, 222 insertions(+), 217 deletions(-) diff --git a/firewood/src/merkle.rs b/firewood/src/merkle.rs index 3bacd663c..04d43f870 100644 --- a/firewood/src/merkle.rs +++ b/firewood/src/merkle.rs @@ -79,11 +79,11 @@ impl + Send + Sync> Merkle { pub fn init_root(&self) -> Result { self.store .put_item( - Node::new(NodeType::Branch(BranchNode { + Node::branch(BranchNode { children: [None; NBRANCH], value: None, children_encoded: Default::default(), - })), + }), Node::max_branch_node_size(), ) .map_err(MerkleError::Shale) @@ -219,7 +219,7 @@ impl + Send + Sync> Merkle { let new_branch_address = self.put_node(new_branch)?.as_ptr(); if idx > 0 { - self.put_node(Node::new(NodeType::Extension(ExtNode { + self.put_node(Node::from(NodeType::Extension(ExtNode { path: PartialPath(matching_path[..idx].to_vec()), child: new_branch_address, child_encoded: None, @@ -347,7 +347,7 @@ impl + Send + Sync> Merkle { .as_ptr(); if !prefix.is_empty() { - self.put_node(Node::new(NodeType::Extension(ExtNode { + self.put_node(Node::from(NodeType::Extension(ExtNode { path: PartialPath(prefix.to_vec()), child: branch_address, child_encoded: None, @@ -427,10 +427,7 @@ impl + Send + Sync> Merkle { // insert the leaf to the empty slot // create a new leaf let leaf_ptr = self - .put_node(Node::new(NodeType::Leaf(LeafNode( - PartialPath(key_nibbles.collect()), - Data(val), - ))))? + .put_node(Node::leaf(PartialPath(key_nibbles.collect()), Data(val)))? .as_ptr(); // set the current child to point to this leaf node.write(|u| { @@ -565,11 +562,11 @@ impl + Send + Sync> Merkle { chd[idx as usize] = Some(c_ptr); let branch = self - .put_node(Node::new(NodeType::Branch(BranchNode { + .put_node(Node::branch(BranchNode { children: chd, value: Some(Data(val)), children_encoded: Default::default(), - })))? + }))? .as_ptr(); set_parent(branch, &mut parents); @@ -609,10 +606,7 @@ impl + Send + Sync> Merkle { // from: [p: Branch] -> [b (v)]x -> [Leaf]x // to: [p: Branch] -> [Leaf (v)] let leaf = self - .put_node(Node::new(NodeType::Leaf(LeafNode( - PartialPath(Vec::new()), - val, - ))))? + .put_node(Node::leaf(PartialPath(Vec::new()), val))? .as_ptr(); p_ref .write(|p| { @@ -625,10 +619,7 @@ impl + Send + Sync> Merkle { // from: P -> [p: Ext]x -> [b (v)]x -> [leaf]x // to: P -> [Leaf (v)] let leaf = self - .put_node(Node::new(NodeType::Leaf(LeafNode( - PartialPath(n.path.clone().into_inner()), - val, - ))))? + .put_node(Node::leaf(PartialPath(n.path.clone().into_inner()), val))? .as_ptr(); deleted.push(p_ptr); set_parent(leaf, parents); @@ -649,7 +640,7 @@ impl + Send + Sync> Merkle { // \____[Leaf]x // to: [p: Branch] -> [Ext] -> [Branch] let ext = self - .put_node(Node::new(NodeType::Extension(ExtNode { + .put_node(Node::from(NodeType::Extension(ExtNode { path: PartialPath(vec![idx]), child: c_ptr, child_encoded: None, @@ -778,7 +769,7 @@ impl + Send + Sync> Merkle { // from: [Branch] -> [Branch]x -> [Branch] // to: [Branch] -> [Ext] -> [Branch] n.children[b_idx as usize] = Some( - self.put_node(Node::new(NodeType::Extension(ExtNode { + self.put_node(Node::from(NodeType::Extension(ExtNode { path: PartialPath(vec![idx]), child: c_ptr, child_encoded: None, @@ -1308,136 +1299,21 @@ pub fn from_nibbles(nibbles: &[u8]) -> impl Iterator + '_ { } #[cfg(test)] -mod test { +mod tests { use super::*; - use shale::cached::{DynamicMem, PlainMem}; - use shale::{CachedStore, Storable}; - use std::ops::Deref; + use node::tests::{extension, leaf}; + use shale::{cached::DynamicMem, compact::CompactSpace, CachedStore}; use std::sync::Arc; use test_case::test_case; #[test_case(vec![0x12, 0x34, 0x56], vec![0x1, 0x2, 0x3, 0x4, 0x5, 0x6])] #[test_case(vec![0xc0, 0xff], vec![0xc, 0x0, 0xf, 0xf])] - fn test_to_nibbles(bytes: Vec, nibbles: Vec) { + fn to_nibbles(bytes: Vec, nibbles: Vec) { let n: Vec<_> = bytes.into_iter().flat_map(to_nibble_array).collect(); assert_eq!(n, nibbles); } - const ZERO_HASH: TrieHash = TrieHash([0u8; TRIE_HASH_LEN]); - - #[test] - fn test_hash_len() { - assert_eq!(TRIE_HASH_LEN, ZERO_HASH.dehydrated_len() as usize); - } - #[test] - fn test_dehydrate() { - let mut to = [1u8; TRIE_HASH_LEN]; - assert_eq!( - { - ZERO_HASH.dehydrate(&mut to).unwrap(); - &to - }, - ZERO_HASH.deref() - ); - } - - #[test] - fn test_hydrate() { - let mut store = PlainMem::new(TRIE_HASH_LEN as u64, 0u8); - store.write(0, ZERO_HASH.deref()); - assert_eq!(TrieHash::hydrate(0, &store).unwrap(), ZERO_HASH); - } - #[test] - fn test_partial_path_encoding() { - let check = |steps: &[u8], term| { - let (d, t) = PartialPath::decode(&PartialPath(steps.to_vec()).encode(term)); - assert_eq!(d.0, steps); - assert_eq!(t, term); - }; - for steps in [ - vec![0x1, 0x2, 0x3, 0x4], - vec![0x1, 0x2, 0x3], - vec![0x0, 0x1, 0x2], - vec![0x1, 0x2], - vec![0x1], - ] { - for term in [true, false] { - check(&steps, term) - } - } - } - #[test] - fn test_merkle_node_encoding() { - let check = |node: Node| { - let mut bytes = vec![0; node.dehydrated_len() as usize]; - node.dehydrate(&mut bytes).unwrap(); - - let mut mem = PlainMem::new(bytes.len() as u64, 0x0); - mem.write(0, &bytes); - println!("{bytes:?}"); - let node_ = Node::hydrate(0, &mem).unwrap(); - assert!(node == node_); - }; - let chd0 = [None; NBRANCH]; - let mut chd1 = chd0; - for node in chd1.iter_mut().take(NBRANCH / 2) { - *node = Some(DiskAddress::from(0xa)); - } - let mut chd_encoded: [Option>; NBRANCH] = Default::default(); - for encoded in chd_encoded.iter_mut().take(NBRANCH / 2) { - *encoded = Some(vec![0x1, 0x2, 0x3]); - } - for node in [ - Node::new_from_hash( - None, - None, - NodeType::Leaf(LeafNode( - PartialPath(vec![0x1, 0x2, 0x3]), - Data(vec![0x4, 0x5]), - )), - ), - Node::new_from_hash( - None, - None, - NodeType::Extension(ExtNode { - path: PartialPath(vec![0x1, 0x2, 0x3]), - child: DiskAddress::from(0x42), - child_encoded: None, - }), - ), - Node::new_from_hash( - None, - None, - NodeType::Extension(ExtNode { - path: PartialPath(vec![0x1, 0x2, 0x3]), - child: DiskAddress::null(), - child_encoded: Some(vec![0x1, 0x2, 0x3]), - }), - ), - Node::new_from_hash( - None, - None, - NodeType::Branch(BranchNode { - children: chd0, - value: Some(Data("hello, world!".as_bytes().to_vec())), - children_encoded: Default::default(), - }), - ), - Node::new_from_hash( - None, - None, - NodeType::Branch(BranchNode { - children: chd1, - value: None, - children_encoded: chd_encoded, - }), - ), - ] { - check(node); - } - } - #[test] - fn test_encode() { + fn create_test_merkle() -> Merkle> { const RESERVED: usize = 0x1000; let mut dm = shale::cached::DynamicMem::new(0x10000, 0); @@ -1465,58 +1341,81 @@ mod test { .expect("CompactSpace init fail"); let store = Box::new(space); - let merkle = Merkle::new(store); + Merkle::new(store) + } - { - let chd = Node::new(NodeType::Leaf(LeafNode( - PartialPath(vec![0x1, 0x2, 0x3]), - Data(vec![0x4, 0x5]), - ))); - let chd_ref = merkle.put_node(chd.clone()).unwrap(); - let chd_encoded = chd_ref.get_encoded(merkle.store.as_ref()); - let new_chd = Node::new(NodeType::decode(chd_encoded).unwrap()); - let new_chd_encoded = new_chd.get_encoded(merkle.store.as_ref()); - assert_eq!(chd_encoded, new_chd_encoded); - - let mut chd_encoded: [Option>; NBRANCH] = Default::default(); - chd_encoded[0] = Some(new_chd_encoded.to_vec()); - let node = Node::new(NodeType::Branch(BranchNode { - children: [None; NBRANCH], - value: Some(Data("value1".as_bytes().to_vec())), - children_encoded: chd_encoded, - })); - - let node_ref = merkle.put_node(node.clone()).unwrap(); - - let r = node_ref.get_encoded(merkle.store.as_ref()); - let new_node = Node::new(NodeType::decode(r).unwrap()); - let new_encoded = new_node.get_encoded(merkle.store.as_ref()); - assert_eq!(r, new_encoded); + fn branch(value: Vec, encoded_child: Option>) -> Node { + let children = Default::default(); + let value = Some(value).map(Data); + let mut children_encoded = <[Option>; NBRANCH]>::default(); + + if let Some(child) = encoded_child { + children_encoded[0] = Some(child); } - { - let chd = Node::new(NodeType::Branch(BranchNode { - children: [None; NBRANCH], - value: Some(Data("value1".as_bytes().to_vec())), - children_encoded: Default::default(), - })); - let chd_ref = merkle.put_node(chd.clone()).unwrap(); - let chd_encoded = chd_ref.get_encoded(merkle.store.as_ref()); - let new_chd = Node::new(NodeType::decode(chd_encoded).unwrap()); - let new_chd_encoded = new_chd.get_encoded(merkle.store.as_ref()); - assert_eq!(chd_encoded, new_chd_encoded); - - let node = Node::new(NodeType::Extension(ExtNode { - path: PartialPath(vec![0x1, 0x2, 0x3]), - child: DiskAddress::null(), - child_encoded: Some(chd_encoded.to_vec()), - })); - let node_ref = merkle.put_node(node.clone()).unwrap(); - - let r = node_ref.get_encoded(merkle.store.as_ref()); - let new_node = Node::new(NodeType::decode(r).unwrap()); - let new_encoded = new_node.get_encoded(merkle.store.as_ref()); - assert_eq!(r, new_encoded); + Node::branch(BranchNode { + children, + value, + children_encoded, + }) + } + + #[test_case(leaf(vec![1, 2, 3], vec![4, 5]) ; "leaf encoding")] + #[test_case(branch(b"value".to_vec(), vec![1, 2, 3].into()) ; "branch with value")] + #[test_case(branch(b"value".to_vec(), None); "branch without value")] + #[test_case(extension(vec![1, 2, 3], DiskAddress::null(), vec![4, 5].into()) ; "extension without child address")] + fn encode_(node: Node) { + let merkle = create_test_merkle(); + + let node_ref = merkle.put_node(node).unwrap(); + let encoded = node_ref.get_encoded(merkle.store.as_ref()); + let new_node = Node::from(NodeType::decode(encoded).unwrap()); + let new_node_encoded = new_node.get_encoded(merkle.store.as_ref()); + + assert_eq!(encoded, new_node_encoded); + } + + #[test] + fn insert_and_retrieve() { + let key = b"hello"; + let val = b"world"; + + let mut merkle = create_test_merkle(); + let root = merkle.init_root().unwrap(); + + merkle.insert(key, val.to_vec(), root).unwrap(); + + let fetched_val = merkle.get(key, root).unwrap(); + + assert_eq!(fetched_val.as_deref(), val.as_slice().into()); + } + + #[test] + fn insert_and_retrieve_multiple() { + let mut merkle = create_test_merkle(); + let root = merkle.init_root().unwrap(); + + // insert values + for key_val in u8::MIN..=u8::MAX { + let key = vec![key_val]; + let val = vec![key_val]; + + merkle.insert(&key, val.clone(), root).unwrap(); + + let fetched_val = merkle.get(&key, root).unwrap(); + + // make sure the value was inserted + assert_eq!(fetched_val.as_deref(), val.as_slice().into()); + } + + // make sure none of the previous values were forgotten after initial insert + for key_val in u8::MIN..=u8::MAX { + let key = vec![key_val]; + let val = vec![key_val]; + + let fetched_val = merkle.get(&key, root).unwrap(); + + assert_eq!(fetched_val.as_deref(), val.as_slice().into()); } } } diff --git a/firewood/src/merkle/node.rs b/firewood/src/merkle/node.rs index f4aacbcc4..f40e74cd5 100644 --- a/firewood/src/merkle/node.rs +++ b/firewood/src/merkle/node.rs @@ -421,6 +421,20 @@ impl NodeType { } } +impl From for Node { + fn from(inner: NodeType) -> Self { + let mut s = Self { + root_hash: OnceLock::new(), + is_encoded_longer_than_hash_len: OnceLock::new(), + encoded: OnceLock::new(), + inner, + lazy_dirty: AtomicBool::new(false), + }; + s.rehash(); + s + } +} + impl Node { const BRANCH_NODE: u8 = 0x0; const EXT_NODE: u8 = 0x1; @@ -468,24 +482,12 @@ impl Node { self.root_hash = OnceLock::new(); } - pub fn new(inner: NodeType) -> Self { - let mut s = Self { - root_hash: OnceLock::new(), - is_encoded_longer_than_hash_len: OnceLock::new(), - encoded: OnceLock::new(), - inner, - lazy_dirty: AtomicBool::new(false), - }; - s.rehash(); - s - } - pub fn branch(node: BranchNode) -> Self { - Self::new(NodeType::Branch(node)) + Self::from(NodeType::Branch(node)) } pub fn leaf(path: PartialPath, data: Data) -> Self { - Self::new(NodeType::Leaf(LeafNode(path, data))) + Self::from(NodeType::Leaf(LeafNode(path, data))) } pub fn inner(&self) -> &NodeType { @@ -842,3 +844,78 @@ impl Storable for Node { } } } + +#[cfg(test)] +pub(super) mod tests { + use std::array::from_fn; + + use super::*; + use shale::cached::PlainMem; + use test_case::test_case; + + pub fn leaf(path: Vec, data: Vec) -> Node { + Node::leaf(PartialPath(path), Data(data)) + } + + pub fn branch( + repeated_disk_address: usize, + value: Option>, + repeated_encoded_child: Option>, + ) -> Node { + let children: [Option; NBRANCH] = from_fn(|i| { + if i < NBRANCH / 2 { + DiskAddress::from(repeated_disk_address).into() + } else { + None + } + }); + + let children_encoded = repeated_encoded_child + .map(|child| { + from_fn(|i| { + if i < NBRANCH / 2 { + child.clone().into() + } else { + None + } + }) + }) + .unwrap_or_default(); + + Node::branch(BranchNode { + children, + value: value.map(Data), + children_encoded, + }) + } + + pub fn extension( + path: Vec, + child_address: DiskAddress, + child_encoded: Option>, + ) -> Node { + Node::from(NodeType::Extension(ExtNode { + path: PartialPath(path), + child: child_address, + child_encoded, + })) + } + + #[test_case(leaf(vec![0x01, 0x02, 0x03], vec![0x04, 0x05]); "leaf_node")] + #[test_case(extension(vec![0x01, 0x02, 0x03], DiskAddress::from(0x42), None); "extension with child address")] + #[test_case(extension(vec![0x01, 0x02, 0x03], DiskAddress::null(), vec![0x01, 0x02, 0x03].into()) ; "extension without child address")] + #[test_case(branch(0x0a, b"hello world".to_vec().into(), None); "branch with data")] + #[test_case(branch(0x0a, None, vec![0x01, 0x02, 0x03].into()); "branch without data")] + fn test_encoding(node: Node) { + let mut bytes = vec![0; node.dehydrated_len() as usize]; + + node.dehydrate(&mut bytes).unwrap(); + + let mut mem = PlainMem::new(node.dehydrated_len(), 0x00); + mem.write(0, &bytes); + + let hydrated_node = Node::hydrate(0, &mem).unwrap(); + + assert_eq!(node, hydrated_node); + } +} diff --git a/firewood/src/merkle/partial_path.rs b/firewood/src/merkle/partial_path.rs index b7aa70315..eb2bd94d8 100644 --- a/firewood/src/merkle/partial_path.rs +++ b/firewood/src/merkle/partial_path.rs @@ -71,3 +71,21 @@ impl PartialPath { } } } + +#[cfg(test)] +mod tests { + use super::*; + use test_case::test_case; + + #[test_case(&[1, 2, 3, 4], true)] + #[test_case(&[1, 2, 3], false)] + #[test_case(&[0, 1, 2], false)] + #[test_case(&[1, 2], true)] + #[test_case(&[1], true)] + fn test_encoding(steps: &[u8], term: bool) { + let path = PartialPath(steps.to_vec()).encode(term); + let (decoded, decoded_term) = PartialPath::decode(&path); + assert_eq!(&decoded.0, &steps); + assert_eq!(decoded_term, term); + } +} diff --git a/firewood/src/merkle/trie_hash.rs b/firewood/src/merkle/trie_hash.rs index 1f64788ae..c9777fb1d 100644 --- a/firewood/src/merkle/trie_hash.rs +++ b/firewood/src/merkle/trie_hash.rs @@ -4,18 +4,15 @@ use shale::{CachedStore, ShaleError, Storable}; use std::{ fmt::{self, Debug}, - io::{Cursor, Write}, + io::Write, }; pub const TRIE_HASH_LEN: usize = 32; +const U64_TRIE_HASH_LEN: u64 = TRIE_HASH_LEN as u64; #[derive(PartialEq, Eq, Clone)] pub struct TrieHash(pub [u8; TRIE_HASH_LEN]); -impl TrieHash { - const MSIZE: u64 = 32; -} - impl std::ops::Deref for TrieHash { type Target = [u8; TRIE_HASH_LEN]; fn deref(&self) -> &[u8; TRIE_HASH_LEN] { @@ -26,22 +23,21 @@ impl std::ops::Deref for TrieHash { impl Storable for TrieHash { fn hydrate(addr: usize, mem: &T) -> Result { let raw = mem - .get_view(addr, Self::MSIZE) + .get_view(addr, U64_TRIE_HASH_LEN) .ok_or(ShaleError::InvalidCacheView { offset: addr, - size: Self::MSIZE, + size: U64_TRIE_HASH_LEN, })?; - Ok(Self( - raw.as_deref()[..Self::MSIZE as usize].try_into().unwrap(), - )) + + Ok(Self(raw.as_deref()[..TRIE_HASH_LEN].try_into().unwrap())) } fn dehydrated_len(&self) -> u64 { - Self::MSIZE + U64_TRIE_HASH_LEN } - fn dehydrate(&self, to: &mut [u8]) -> Result<(), ShaleError> { - Cursor::new(to).write_all(&self.0).map_err(ShaleError::Io) + fn dehydrate(&self, mut to: &mut [u8]) -> Result<(), ShaleError> { + to.write_all(&self.0).map_err(ShaleError::Io) } } @@ -50,3 +46,18 @@ impl Debug for TrieHash { write!(f, "{}", hex::encode(self.0)) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_dehydrate() { + let zero_hash = TrieHash([0u8; TRIE_HASH_LEN]); + + let mut to = [1u8; TRIE_HASH_LEN]; + zero_hash.dehydrate(&mut to).unwrap(); + + assert_eq!(&to, &zero_hash.0); + } +} diff --git a/firewood/src/proof.rs b/firewood/src/proof.rs index 6f89f8377..c76765797 100644 --- a/firewood/src/proof.rs +++ b/firewood/src/proof.rs @@ -449,7 +449,7 @@ impl + Send> Proof { ) -> Result<(DiskAddress, Option, usize), ProofError> { let node = NodeType::decode(buf)?; let new_node = merkle - .put_node(Node::new(node)) + .put_node(Node::from(node)) .map_err(ProofError::InvalidNode)?; let addr = new_node.as_ptr(); match new_node.inner() {