From 443547d67e7258e13bfd97eb43c4201f82848b49 Mon Sep 17 00:00:00 2001 From: Guillaume Ballet Date: Mon, 11 Nov 2019 16:42:53 +0100 Subject: [PATCH 01/10] Turn insert_leaf into self.insert + add tree trait --- src/node.rs | 556 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 556 insertions(+) diff --git a/src/node.rs b/src/node.rs index f1a7268..16f99c9 100644 --- a/src/node.rs +++ b/src/node.rs @@ -3,6 +3,11 @@ extern crate sha3; use super::utils::*; use sha3::{Digest, Keccak256}; +pub trait Tree { + fn is_leaf(&self) -> bool; + fn insert(&mut self, key: &K, value: V) -> Result<(), String>; +} + #[derive(Debug, Clone, PartialEq)] pub enum Node { Hash(Vec), @@ -12,6 +17,139 @@ pub enum Node { EmptySlot, } +impl Tree> for Node { + fn is_leaf(&self) -> bool { + match self { + Node::Leaf(_, _) => true, + _ => false, + } + } + + fn insert(&mut self, key: &NibbleKey, value: Vec) -> Result<(), String> { + use Node::*; + + if key.len() == 0 { + return Err("Attempted to insert a 0-byte key".to_string()); + } + + match self { + Leaf(leafkey, leafvalue) => { + // Find the common part of the current key with that of the + // leaf and create an intermediate full node. + let firstdiffindex = leafkey.factor_length(key); + + // Return an error if the leaf is already present. + if firstdiffindex == key.len() { + return Err(format!("Key is is already present!",)); + } + + // Create the new root, which is a full node. + let mut res = vec![EmptySlot; 16]; + // Add the initial leaf, with a key truncated by the common + // key part. + res[leafkey[firstdiffindex] as usize] = Leaf( + NibbleKey::from(leafkey[firstdiffindex + 1..].to_vec()), + leafvalue.to_vec(), + ); + // Add the node to be inserted + res[key[firstdiffindex] as usize] = + Leaf(NibbleKey::from(key[firstdiffindex + 1..].to_vec()), value); + // Put the common part into an extension node + *self = if firstdiffindex == 0 { + // Special case: no extension necessary + Branch(res) + } else { + Extension( + NibbleKey::from(key[..firstdiffindex].to_vec()), + Box::new(Branch(res)), + ) + }; + Ok(()) + } + Extension(extkey, box child) => { + // Find the common part of the current key with that of the + // extension and create an intermediate full node. + let firstdiffindex = extkey.factor_length(&NibbleKey::from(key.clone())); + + // Special case: key is longer than the extension key: + // recurse on the child node. + if firstdiffindex == extkey.len() { + child.insert(&NibbleKey::from(key[extkey.len()..].to_vec()), value)?; + return Ok(()); + } + + // Special case: key is completely unlike the extension key + if firstdiffindex == 0 { + let mut res = vec![EmptySlot; 16]; + + // Create the entry for the truncated extension key + // Was it an extension of 1 ? If so, place the node directly + // otherwise truncate the extension. + res[extkey[0] as usize] = if extkey.len() == 1 { + child.clone() + } else { + Extension( + NibbleKey::from(extkey[1..].to_vec()), + Box::new(child.clone()), + ) + }; + + // Create the entry for the node. If there was only a + // difference of one byte, that byte will be consumed by + // the fullnode and therefore the key in the leaf will be + // an empty slice `[]`. + res[key[0] as usize] = Leaf(NibbleKey::from(key[1..].to_vec()), value); + + *self = Branch(res); + return Ok(()); + } + + // Create the new root, which is a full node. + let mut res = vec![EmptySlot; 16]; + // Add the initial leaf, with a key truncated by the common + // key part. If the common part corresponds to the extension + // key length minus one, then there is no need for the creation + // of an extension node past the full node. + res[extkey[firstdiffindex] as usize] = if extkey.len() - firstdiffindex > 1 { + Extension( + NibbleKey::from(extkey[firstdiffindex + 1..].to_vec()), + Box::new(child.clone()), + ) + } else { + child.clone() + }; + // Add the node to be inserted + res[key[firstdiffindex] as usize] = + Leaf(NibbleKey::from(key[firstdiffindex + 1..].to_vec()), value); + // Put the common part into an extension node + *self = Extension( + NibbleKey::from(extkey[..firstdiffindex].to_vec()), + Box::new(Branch(res)), + ); + Ok(()) + } + Branch(ref mut vec) => { + let idx = key[0] as usize; + // If the slot isn't yet in use, fill it, and otherwise, + // recurse into the child node. + if vec[idx] == EmptySlot { + // XXX check that the value is at least 1 + vec[idx] = Leaf(NibbleKey::from(key[1..].to_vec()), value); + } else { + vec[idx].insert(&NibbleKey::from(key[1..].to_vec()), value)?; + } + // Return the root node with an updated entry + Ok(()) + } + EmptySlot => { + *self = Leaf(key.clone(), value); + Ok(()) + } + _ => panic!("Can not insert a node into a hashed node"), + } + } +} + impl Default for Node { fn default() -> Self { Node::EmptySlot @@ -268,6 +406,7 @@ impl Node { #[cfg(test)] mod tests { + use super::super::utils; use super::Node::*; use super::*; @@ -354,4 +493,421 @@ mod tests { ] ); } + + #[test] + fn insert_leaf_zero_length_key_after_fullnode() { + let mut root = Extension( + NibbleKey::from(vec![0u8; 31]), + Box::new(Branch(vec![ + EmptySlot, + Leaf(NibbleKey::from(vec![]), vec![0u8; 32]), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + ])), + ); + root.insert(&NibbleKey::from(vec![0u8; 32]), vec![1u8; 32]) + .unwrap(); + assert_eq!( + root, + Extension( + NibbleKey::from(vec![0u8; 31]), + Box::new(Branch(vec![ + Leaf(NibbleKey::from(vec![]), vec![1u8; 32]), + Leaf(NibbleKey::from(vec![]), vec![0u8; 32]), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot + ])) + ) + ); + } + + #[test] + fn insert_leaf_into_extension_root_all_bytes_in_key_common() { + let mut root = Extension( + NibbleKey::from(vec![0xd, 0xe, 0xa, 0xd]), + Box::new(Leaf(NibbleKey::from(vec![0u8; 28]), vec![1u8; 32])), + ); + let mut key = vec![1u8; 32]; + key[0] = 0xd; + key[1] = 0xe; + key[2] = 0xa; + key[3] = 0xd; + root.insert(&NibbleKey::from(key), vec![2u8; 32]).unwrap(); + assert_eq!( + root, + Extension( + NibbleKey::from(vec![0xd, 0xe, 0xa, 0xd]), + Box::new(Branch(vec![ + Leaf(NibbleKey::from(vec![0u8; 27]), vec![1u8; 32]), + Leaf(NibbleKey::from(vec![1u8; 27]), vec![2u8; 32]), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot + ])) + ) + ); + } + + #[test] + fn insert_leaf_into_extension_root_no_common_bytes_in_key() { + let mut root = Extension( + NibbleKey::from(vec![0xd, 0xe, 0xa, 0xd]), + Box::new(Leaf(NibbleKey::from(vec![0u8; 24]), vec![1u8; 32])), + ); + root.insert(&NibbleKey::from(vec![2u8; 32]), vec![1u8; 32]) + .unwrap(); + assert_eq!( + root, + Branch(vec![ + EmptySlot, + EmptySlot, + Leaf(NibbleKey::from(vec![2u8; 31]), vec![1u8; 32]), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + Extension( + NibbleKey::from(vec![14, 10, 13]), + Box::new(Leaf(NibbleKey::from(vec![0u8; 24]), vec![1u8; 32])) + ), + EmptySlot, + EmptySlot + ]) + ); + } + + #[test] + fn insert_leaf_into_extension_root_half_bytes_in_key_common() { + let mut root = Extension( + NibbleKey::from(vec![0xd, 0xe, 0xa, 0xd]), + Box::new(Leaf(NibbleKey::from(vec![0u8; 28]), vec![1u8; 32])), + ); + let mut key = vec![0u8; 32]; + key[0] = 0xd; + key[1] = 0xe; + root.insert(&NibbleKey::from(key), vec![1u8; 32]).unwrap(); + assert_eq!( + root, + Extension( + NibbleKey::from(vec![0xd, 0xe]), + Box::new(Branch(vec![ + Leaf(NibbleKey::from(vec![0u8; 29]), vec![1u8; 32]), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + Extension( + NibbleKey::from(vec![0xd]), + Box::new(Leaf(NibbleKey::from(vec![0u8; 28]), vec![1u8; 32])) + ), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot + ])) + ) + ); + } + + #[test] + fn insert_leaf_into_extension_root_almost_all_bytes_in_key_common() { + let mut root = Extension( + NibbleKey::from(vec![0xd, 0xe, 0xa, 0xd]), + Box::new(Leaf(NibbleKey::from(vec![0u8; 28]), vec![1u8; 32])), + ); + let mut key = vec![0u8; 32]; + key[0] = 0xd; + key[1] = 0xe; + key[2] = 0xa; + root.insert(&NibbleKey::from(key), vec![1u8; 32]).unwrap(); + assert_eq!( + root, + Extension( + NibbleKey::from(vec![0xd, 0xe, 0xa]), + Box::new(Branch(vec![ + Leaf(NibbleKey::from(vec![0u8; 28]), vec![1u8; 32]), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + Leaf(NibbleKey::from(vec![0u8; 28]), vec![1u8; 32]), + EmptySlot, + EmptySlot + ])) + ) + ); + } + + #[test] + fn insert_leaf_into_leaf_root_common_bytes_in_key() { + let mut key = vec![0u8; 32]; + for (i, v) in key.iter_mut().enumerate() { + if i >= 16 { + break; + } + *v = 2u8; + } + let mut root = Leaf(NibbleKey::from(key), vec![1u8; 32]); + root.insert(&NibbleKey::from(vec![2u8; 32]), vec![1u8; 32]) + .unwrap(); + assert_eq!( + root, + Extension( + NibbleKey::from(vec![2u8; 16]), + Box::new(Branch(vec![ + Leaf(NibbleKey::from(vec![0u8; 15]), vec![1u8; 32]), + EmptySlot, + Leaf(NibbleKey::from(vec![2u8; 15]), vec![1u8; 32]), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot + ])) + ) + ); + } + + #[test] + fn insert_leaf_into_leaf_root_no_common_bytes_in_key() { + let mut root = Leaf(NibbleKey::from(vec![1u8; 32]), vec![1u8; 32]); + root.insert(&NibbleKey::from(vec![2u8; 32]), vec![1u8; 32]) + .unwrap(); + assert_eq!( + root, + Branch(vec![ + EmptySlot, + Leaf(NibbleKey::from(vec![1u8; 31]), vec![1u8; 32]), + Leaf(NibbleKey::from(vec![2u8; 31]), vec![1u8; 32]), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot + ]) + ); + } + + #[test] + fn insert_leaf_into_emptyslot() { + let mut root = Node::default(); + assert_eq!(root, EmptySlot); + root.insert(&NibbleKey::from(vec![0u8; 32]), vec![1u8; 32]) + .unwrap(); + assert_eq!(root, Leaf(NibbleKey::from(vec![0u8; 32]), vec![1u8; 32])); + } + + #[test] + fn insert_two_leaves_into_emptyslot() { + let mut root = Node::default(); + assert_eq!(root, EmptySlot); + root.insert(&NibbleKey::from(vec![0u8; 32]), vec![1u8; 32]) + .unwrap(); + root.insert(&NibbleKey::from(vec![1u8; 32]), vec![1u8; 32]) + .unwrap(); + assert_eq!( + root, + Branch(vec![ + Leaf(NibbleKey::from(vec![0u8; 31]), vec![1u8; 32]), + Leaf(NibbleKey::from(vec![1u8; 31]), vec![1u8; 32]), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot + ]) + ); + } + + #[test] + fn insert_leaf_into_empty_root() { + let children = vec![EmptySlot; 16]; + let mut root = Branch(children); + root.insert(&NibbleKey::from(vec![0u8; 32]), vec![1u8; 32]) + .unwrap(); + assert_eq!( + root, + Branch(vec![ + Leaf(NibbleKey::from(vec![0u8; 31]), vec![1u8; 32]), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot + ]) + ); + } + + #[test] + fn insert_leaf_into_two_level_fullnodes() { + let mut root = Branch(vec![ + Branch(vec![EmptySlot; 16]), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + ]); + root.insert(&NibbleKey::from(vec![0u8; 32]), vec![1u8; 32]) + .unwrap(); + assert_eq!( + root, + Branch(vec![ + Branch(vec![ + Leaf(NibbleKey::from(vec![0u8; 30]), vec![1u8; 32]), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot + ]), + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot, + EmptySlot + ]) + ); + } + + #[test] + fn insert_leaf_two_embedded_nodes() { + let expected_root = + hex::decode(&"0x1e1f73cc50d595585797d261df5be9bf69037a57e6470c9e4ffc87b6221ab67a"[2..]) + .unwrap(); + let inputs = [ + ("0x1111111111111111111111111111111111111111", "0xffff"), + ("0x2222222222222222222222222222222222222222", "0xeeee"), + ]; + + let mut root = Branch(vec![EmptySlot; 16]); + for (ik, iv) in inputs.iter() { + let k = NibbleKey::from(utils::ByteKey(hex::decode(&ik[2..]).unwrap())); + let v = hex::decode(&iv[2..]).unwrap(); + root.insert(&k, v).unwrap(); + } + + let root_hash = root.hash(); + assert_eq!(expected_root, root_hash); + } } From 74340a1db0ebb76da7dcafa983ccdc0828a027b9 Mon Sep 17 00:00:00 2001 From: Guillaume Ballet Date: Mon, 11 Nov 2019 22:48:59 +0100 Subject: [PATCH 02/10] First attempt at implementing the iterator --- src/node.rs | 109 ++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 106 insertions(+), 3 deletions(-) diff --git a/src/node.rs b/src/node.rs index 16f99c9..a5ce3c5 100644 --- a/src/node.rs +++ b/src/node.rs @@ -3,9 +3,60 @@ extern crate sha3; use super::utils::*; use sha3::{Digest, Keccak256}; -pub trait Tree { +pub struct ChildIterator<'a, T: Tree + ?Sized> { + index: usize, + node: &'a T, +} + +impl<'a> std::iter::Iterator for ChildIterator<'a, Node> { + type Item = Node; + + #[inline] + fn next(&mut self) -> Option { + let node = &self.node; + if node.is_leaf() || node.is_empty() { + return None; + } + + // At this stage, the node is either a hash, a branch + // or an extension. A hash isn't supported and should + // not even have been created. + if let Node::Hash(_) = node { + panic!("Can not iterate over a hash"); + } + + match node { + Node::Extension(_, box child) => { + // The iterator will return a value the first time it's called, + // and nothing the second time since it only has one child. + if self.index == 0 { + self.index = 1; + Some(child.clone()) + } else { + None + } + } + Node::Branch(ref vec) => { + if self.index < vec.len() { + self.index += 1; + Some(vec[self.index - 1].clone()) + } else { + None + } + } + // Should not happen + _ => panic!("This should never be reached"), + } + } +} + +pub trait Tree { + type Key; + type Value; fn is_leaf(&self) -> bool; - fn insert(&mut self, key: &K, value: V) -> Result<(), String>; + fn is_empty(&self) -> bool; + fn children(&self) -> ChildIterator>; + fn insert(&mut self, key: &Self::Key, value: Self::Value) -> Result<(), String>; } #[derive(Debug, Clone, PartialEq)] @@ -17,7 +68,10 @@ pub enum Node { EmptySlot, } -impl Tree> for Node { +impl Tree for Node { + type Key = NibbleKey; + type Value = Vec; + fn is_leaf(&self) -> bool { match self { Node::Leaf(_, _) => true, @@ -25,6 +79,35 @@ impl Tree> for Node { } } + fn is_empty(&self) -> bool { + match self { + Node::EmptySlot => true, + _ => false, + } + } + + fn children(&self) -> ChildIterator>> { + match self { + Node::EmptySlot => ChildIterator { + index: 0, + node: self, + }, + Node::Leaf(_, _) => ChildIterator { + index: 0, + node: self, + }, + Node::Extension(_, _) => ChildIterator { + index: 0, + node: self, + }, + Node::Branch(_) => ChildIterator { + index: 0, + node: self, + }, + _ => panic!("not implemented"), + } + } + fn insert(&mut self, key: &NibbleKey, value: Vec) -> Result<(), String> { use Node::*; @@ -910,4 +993,24 @@ mod tests { let root_hash = root.hash(); assert_eq!(expected_root, root_hash); } + + #[test] + fn test_node_as_a_tree() { + let root = EmptySlot; + assert_eq!(root.is_empty(), true); + assert_eq!(root.is_leaf(), false); + for child in root.children() { + panic!("An empty slot has no child"); + } + + let root = Leaf(NibbleKey::from(vec![]), vec![]); + assert_eq!(root.is_empty(), false); + assert_eq!(root.is_leaf(), true); + + let mut root = EmptySlot; + root.insert(&NibbleKey::from(vec![1, 2, 3, 4]), vec![0, 0]); + root.insert(&NibbleKey::from(vec![1, 2, 1, 2]), vec![1, 1]); + assert_eq!(root.is_empty(), false); + assert_eq!(root.is_leaf(), false); + } } From 56b65d79e26d7f990fdb04d356bc35ae4d4018af Mon Sep 17 00:00:00 2001 From: Guillaume Ballet Date: Tue, 12 Nov 2019 08:24:40 +0100 Subject: [PATCH 03/10] Attempt at rewrite --- src/node.rs | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/src/node.rs b/src/node.rs index a5ce3c5..52f738b 100644 --- a/src/node.rs +++ b/src/node.rs @@ -3,12 +3,12 @@ extern crate sha3; use super::utils::*; use sha3::{Digest, Keccak256}; -pub struct ChildIterator<'a, T: Tree + ?Sized> { +pub struct NodeChildIterator<'a> { index: usize, - node: &'a T, + node: &'a Node, } -impl<'a> std::iter::Iterator for ChildIterator<'a, Node> { +impl<'a> std::iter::Iterator for NodeChildIterator<'a> { type Item = Node; #[inline] @@ -55,7 +55,8 @@ pub trait Tree { type Value; fn is_leaf(&self) -> bool; fn is_empty(&self) -> bool; - fn children(&self) -> ChildIterator>; + //fn children(&self) -> ChildIterator>; + fn children(&self) -> dyn Iterator; fn insert(&mut self, key: &Self::Key, value: Self::Value) -> Result<(), String>; } @@ -86,21 +87,21 @@ impl Tree for Node { } } - fn children(&self) -> ChildIterator>> { + fn children(&self) -> dyn Iterator { match self { - Node::EmptySlot => ChildIterator { + Node::EmptySlot => NodeChildIterator { index: 0, node: self, }, - Node::Leaf(_, _) => ChildIterator { + Node::Leaf(_, _) => NodeChildIterator { index: 0, node: self, }, - Node::Extension(_, _) => ChildIterator { + Node::Extension(_, _) => NodeChildIterator { index: 0, node: self, }, - Node::Branch(_) => ChildIterator { + Node::Branch(_) => NodeChildIterator { index: 0, node: self, }, From eaedd19d909f99ef3a0dba1567e274fdba00fe88 Mon Sep 17 00:00:00 2001 From: Guillaume Ballet Date: Tue, 12 Nov 2019 12:51:03 +0100 Subject: [PATCH 04/10] Rewrite iterator so that it works --- src/node.rs | 101 ++++++++++++++++++++++++++++++---------------------- 1 file changed, 58 insertions(+), 43 deletions(-) diff --git a/src/node.rs b/src/node.rs index 52f738b..fff6906 100644 --- a/src/node.rs +++ b/src/node.rs @@ -3,49 +3,21 @@ extern crate sha3; use super::utils::*; use sha3::{Digest, Keccak256}; -pub struct NodeChildIterator<'a> { +pub struct NodeChildIterator<'a, K, V> { index: usize, - node: &'a Node, + node: &'a dyn Tree, } -impl<'a> std::iter::Iterator for NodeChildIterator<'a> { - type Item = Node; +impl<'a> std::iter::Iterator for NodeChildIterator<'a, NibbleKey, Vec> { + type Item = &'a dyn Tree>; #[inline] fn next(&mut self) -> Option { - let node = &self.node; - if node.is_leaf() || node.is_empty() { - return None; - } - - // At this stage, the node is either a hash, a branch - // or an extension. A hash isn't supported and should - // not even have been created. - if let Node::Hash(_) = node { - panic!("Can not iterate over a hash"); - } - - match node { - Node::Extension(_, box child) => { - // The iterator will return a value the first time it's called, - // and nothing the second time since it only has one child. - if self.index == 0 { - self.index = 1; - Some(child.clone()) - } else { - None - } - } - Node::Branch(ref vec) => { - if self.index < vec.len() { - self.index += 1; - Some(vec[self.index - 1].clone()) - } else { - None - } - } - // Should not happen - _ => panic!("This should never be reached"), + if self.index < self.node.num_children() { + self.index += 1; + self.node.ith_child(self.index - 1) + } else { + None } } } @@ -53,10 +25,12 @@ impl<'a> std::iter::Iterator for NodeChildIterator<'a> { pub trait Tree { type Key; type Value; + fn is_leaf(&self) -> bool; fn is_empty(&self) -> bool; - //fn children(&self) -> ChildIterator>; - fn children(&self) -> dyn Iterator; + fn num_children(&self) -> usize; + fn ith_child(&self, index: usize) -> Option<&dyn Tree>; + fn children(&self) -> NodeChildIterator; fn insert(&mut self, key: &Self::Key, value: Self::Value) -> Result<(), String>; } @@ -87,7 +61,33 @@ impl Tree for Node { } } - fn children(&self) -> dyn Iterator { + fn num_children(&self) -> usize { + match self { + Node::EmptySlot | Node::Leaf(_, _) | Node::Hash(_) => 0usize, + Node::Extension(_, _) => 1usize, + Node::Branch(ref vec) => vec.iter().filter(|&x| *x != Node::EmptySlot).count(), + } + } + + fn ith_child(&self, index: usize) -> Option<&dyn Tree> { + if index >= self.num_children() { + panic!(format!( + "Requested child #{}, only have #{}", + index, + self.num_children() + )); + } + + match self { + Node::EmptySlot | Node::Leaf(_, _) | Node::Hash(_) => { + None as Option<&dyn Tree> + } + Node::Extension(_, box ext) => Some(&*ext), + Node::Branch(ref vec) => Some(&vec[index]), + } + } + + fn children(&self) -> NodeChildIterator { match self { Node::EmptySlot => NodeChildIterator { index: 0, @@ -1000,18 +1000,33 @@ mod tests { let root = EmptySlot; assert_eq!(root.is_empty(), true); assert_eq!(root.is_leaf(), false); - for child in root.children() { + for _ in root.children() { panic!("An empty slot has no child"); } let root = Leaf(NibbleKey::from(vec![]), vec![]); assert_eq!(root.is_empty(), false); assert_eq!(root.is_leaf(), true); + for _ in root.children() { + panic!("An leaf has no child"); + } let mut root = EmptySlot; - root.insert(&NibbleKey::from(vec![1, 2, 3, 4]), vec![0, 0]); - root.insert(&NibbleKey::from(vec![1, 2, 1, 2]), vec![1, 1]); + root.insert(&NibbleKey::from(vec![1, 2, 3, 4]), vec![0, 0]) + .unwrap(); + root.insert(&NibbleKey::from(vec![1, 2, 1, 2]), vec![1, 1]) + .unwrap(); + assert_eq!(root.is_empty(), false); + assert_eq!(root.is_leaf(), false); + assert_eq!(root.children().count(), 1); + + let mut root = EmptySlot; + root.insert(&NibbleKey::from(vec![1, 2, 3, 4]), vec![0, 0]) + .unwrap(); + root.insert(&NibbleKey::from(vec![2, 2, 2, 2]), vec![1, 1]) + .unwrap(); assert_eq!(root.is_empty(), false); assert_eq!(root.is_leaf(), false); + assert_eq!(root.children().count(), 2); } } From 39d8b2309f1fe5cd31ff306755d4df2ed3749ee4 Mon Sep 17 00:00:00 2001 From: Guillaume Ballet Date: Tue, 12 Nov 2019 14:56:18 +0100 Subject: [PATCH 05/10] Move `trait Tree` to its own module Also, set the number of children reported by Node::Branch to 16 instead of the number of non-empty nodes. --- src/lib.rs | 1 + src/node.rs | 42 +++++++++--------------------------------- src/tree.rs | 40 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 33 deletions(-) create mode 100644 src/tree.rs diff --git a/src/lib.rs b/src/lib.rs index d145f9e..3dd7619 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6,6 +6,7 @@ extern crate sha3; pub mod instruction; pub mod multiproof; pub mod node; +pub mod tree; pub mod utils; use instruction::*; diff --git a/src/node.rs b/src/node.rs index fff6906..eaeb28b 100644 --- a/src/node.rs +++ b/src/node.rs @@ -1,39 +1,9 @@ extern crate sha3; +use super::tree::*; use super::utils::*; use sha3::{Digest, Keccak256}; -pub struct NodeChildIterator<'a, K, V> { - index: usize, - node: &'a dyn Tree, -} - -impl<'a> std::iter::Iterator for NodeChildIterator<'a, NibbleKey, Vec> { - type Item = &'a dyn Tree>; - - #[inline] - fn next(&mut self) -> Option { - if self.index < self.node.num_children() { - self.index += 1; - self.node.ith_child(self.index - 1) - } else { - None - } - } -} - -pub trait Tree { - type Key; - type Value; - - fn is_leaf(&self) -> bool; - fn is_empty(&self) -> bool; - fn num_children(&self) -> usize; - fn ith_child(&self, index: usize) -> Option<&dyn Tree>; - fn children(&self) -> NodeChildIterator; - fn insert(&mut self, key: &Self::Key, value: Self::Value) -> Result<(), String>; -} - #[derive(Debug, Clone, PartialEq)] pub enum Node { Hash(Vec), @@ -61,14 +31,20 @@ impl Tree for Node { } } + /// Returns the *maximum* child count of the tree's root. + /// + /// Branch nodes will always report 16, as empty slots are counted as children. fn num_children(&self) -> usize { match self { Node::EmptySlot | Node::Leaf(_, _) | Node::Hash(_) => 0usize, Node::Extension(_, _) => 1usize, - Node::Branch(ref vec) => vec.iter().filter(|&x| *x != Node::EmptySlot).count(), + Node::Branch(_) => 16usize, } } + /// Return the tree root's *ith* child. + /// + /// The function will return `EmptySlot` instead of `None` if a branch node has no *ith* child. fn ith_child(&self, index: usize) -> Option<&dyn Tree> { if index >= self.num_children() { panic!(format!( @@ -1027,6 +1003,6 @@ mod tests { .unwrap(); assert_eq!(root.is_empty(), false); assert_eq!(root.is_leaf(), false); - assert_eq!(root.children().count(), 2); + assert_eq!(root.children().count(), 16); } } diff --git a/src/tree.rs b/src/tree.rs new file mode 100644 index 0000000..2bdbac1 --- /dev/null +++ b/src/tree.rs @@ -0,0 +1,40 @@ +use super::utils::*; + +pub struct NodeChildIterator<'a, K, V> { + pub index: usize, + pub node: &'a dyn Tree, +} + +/// A trait representing the an underlying tree structure. +pub trait Tree { + /// The tree's key type. Must be specified when implementing this trait. + type Key; + /// The tree's value type. Must be specified when implementing this trait. + type Value; + + /// Specifies if the current tree is a simple leaf. + fn is_leaf(&self) -> bool; + /// Specifies if the current tree is empty. + fn is_empty(&self) -> bool; + /// Returns the tree root's *maximal* number of children. + fn num_children(&self) -> usize; + /// Returns a pointer to child #i, or `None` if no such child exists. + fn ith_child(&self, index: usize) -> Option<&dyn Tree>; + /// Returns an iterator to the node's children. Some of these nodes can be empty. + fn children(&self) -> NodeChildIterator; + fn insert(&mut self, key: &Self::Key, value: Self::Value) -> Result<(), String>; +} + +impl<'a> std::iter::Iterator for NodeChildIterator<'a, NibbleKey, Vec> { + type Item = &'a dyn Tree>; + + #[inline] + fn next(&mut self) -> Option { + if self.index < self.node.num_children() { + self.index += 1; + self.node.ith_child(self.index - 1) + } else { + None + } + } +} From 94657491863a4038dd68404c5c55084a493bdf28 Mon Sep 17 00:00:00 2001 From: Guillaume Ballet Date: Wed, 13 Nov 2019 14:16:22 +0100 Subject: [PATCH 06/10] Remove insert_leaf and associated code duplication --- src/lib.rs | 652 ++++++---------------------------------------------- src/tree.rs | 1 + 2 files changed, 67 insertions(+), 586 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 3dd7619..ad32df2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -111,130 +111,6 @@ pub fn rebuild(proof: &Multiproof) -> Result { .ok_or(String::from("Stack underflow, expected root node")) } -// Insert a `(key,value)` pair into a (sub-)tree represented by `root`. -// It returns the root of the updated (sub-)tree. -pub fn insert_leaf(root: &mut Node, key: &NibbleKey, value: Vec) -> Result { - use Node::*; - - if key.len() == 0 { - return Err("Attempted to insert a 0-byte key".to_string()); - } - - match root { - Leaf(leafkey, leafvalue) => { - // Find the common part of the current key with that of the - // leaf and create an intermediate full node. - let firstdiffindex = leafkey.factor_length(key); - - // Return an error if the leaf is already present. - if firstdiffindex == key.len() { - return Err(format!("Key is is already present!",)); - } - - // Create the new root, which is a full node. - let mut res = vec![EmptySlot; 16]; - // Add the initial leaf, with a key truncated by the common - // key part. - res[leafkey[firstdiffindex] as usize] = Leaf( - NibbleKey::from(leafkey[firstdiffindex + 1..].to_vec()), - leafvalue.to_vec(), - ); - // Add the node to be inserted - res[key[firstdiffindex] as usize] = - Leaf(NibbleKey::from(key[firstdiffindex + 1..].to_vec()), value); - // Put the common part into an extension node - if firstdiffindex == 0 { - // Special case: no extension necessary - Ok(Branch(res)) - } else { - Ok(Extension( - NibbleKey::from(key[..firstdiffindex].to_vec()), - Box::new(Branch(res)), - )) - } - } - Extension(extkey, box child) => { - // Find the common part of the current key with that of the - // extension and create an intermediate full node. - let firstdiffindex = extkey.factor_length(&NibbleKey::from(key.clone())); - - // Special case: key is longer than the extension key: - // recurse on the child node. - if firstdiffindex == extkey.len() { - let childroot = insert_leaf( - &mut child.clone(), - &NibbleKey::from(key[extkey.len()..].to_vec()), - value, - )?; - return Ok(Extension(extkey.clone(), Box::new(childroot))); - } - - // Special case: key is completely unlike the extension key - if firstdiffindex == 0 { - let mut res = vec![EmptySlot; 16]; - - // Create the entry for the truncated extension key - // Was it an extension of 1 ? If so, place the node directly - // otherwise truncate the extension. - res[extkey[0] as usize] = if extkey.len() == 1 { - child.clone() - } else { - Extension( - NibbleKey::from(extkey[1..].to_vec()), - Box::new(child.clone()), - ) - }; - - // Create the entry for the node. If there was only a - // difference of one byte, that byte will be consumed by - // the fullnode and therefore the key in the leaf will be - // an empty slice `[]`. - res[key[0] as usize] = Leaf(NibbleKey::from(key[1..].to_vec()), value); - - return Ok(Branch(res)); - } - - // Create the new root, which is a full node. - let mut res = vec![EmptySlot; 16]; - // Add the initial leaf, with a key truncated by the common - // key part. If the common part corresponds to the extension - // key length minus one, then there is no need for the creation - // of an extension node past the full node. - res[extkey[firstdiffindex] as usize] = if extkey.len() - firstdiffindex > 1 { - Extension( - NibbleKey::from(extkey[firstdiffindex + 1..].to_vec()), - Box::new(child.clone()), - ) - } else { - child.clone() - }; - // Add the node to be inserted - res[key[firstdiffindex] as usize] = - Leaf(NibbleKey::from(key[firstdiffindex + 1..].to_vec()), value); - // Put the common part into an extension node - Ok(Extension( - NibbleKey::from(extkey[..firstdiffindex].to_vec()), - Box::new(Branch(res)), - )) - } - Branch(ref mut vec) => { - let idx = key[0] as usize; - // If the slot isn't yet in use, fill it, and otherwise, - // recurse into the child node. - vec[idx] = if vec[idx] == EmptySlot { - // XXX check that the value is at least 1 - Leaf(NibbleKey::from(key[1..].to_vec()), value) - } else { - insert_leaf(&mut vec[idx], &NibbleKey::from(key[1..].to_vec()), value)? - }; - // Return the root node with an updated entry - Ok(Branch(vec.to_vec())) - } - EmptySlot => Ok(Leaf(key.clone(), value)), - _ => panic!("Can not insert a node into a hashed node"), - } -} - // Helper function that generates a multiproof based on one `(key.value)` // pair. pub fn make_multiproof(root: &Node, keys: Vec) -> Result { @@ -383,6 +259,7 @@ pub fn make_multiproof(root: &Node, keys: Vec) -> Result= 16 { - break; - } - *v = 2u8; - } - let mut root = Leaf(NibbleKey::from(key), vec![1u8; 32]); - let out = insert_leaf(&mut root, &NibbleKey::from(vec![2u8; 32]), vec![1u8; 32]).unwrap(); - assert_eq!( - out, - Extension( - NibbleKey::from(vec![2u8; 16]), - Box::new(Branch(vec![ - Leaf(NibbleKey::from(vec![0u8; 15]), vec![1u8; 32]), - EmptySlot, - Leaf(NibbleKey::from(vec![2u8; 15]), vec![1u8; 32]), - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot - ])) - ) - ); - } - - #[test] - fn insert_leaf_into_leaf_root_no_common_bytes_in_key() { - let mut root = Leaf(NibbleKey::from(vec![1u8; 32]), vec![1u8; 32]); - let out = insert_leaf(&mut root, &NibbleKey::from(vec![2u8; 32]), vec![1u8; 32]).unwrap(); - assert_eq!( - out, - Branch(vec![ - EmptySlot, - Leaf(NibbleKey::from(vec![1u8; 31]), vec![1u8; 32]), - Leaf(NibbleKey::from(vec![2u8; 31]), vec![1u8; 32]), - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot - ]) - ); - } - - #[test] - fn insert_leaf_into_emptyslot() { - let mut root = Node::default(); - assert_eq!(root, EmptySlot); - let out = insert_leaf(&mut root, &NibbleKey::from(vec![0u8; 32]), vec![1u8; 32]).unwrap(); - assert_eq!(out, Leaf(NibbleKey::from(vec![0u8; 32]), vec![1u8; 32])); - } - - #[test] - fn insert_two_leaves_into_emptyslot() { - let mut root = Node::default(); - assert_eq!(root, EmptySlot); - root = insert_leaf(&mut root, &NibbleKey::from(vec![0u8; 32]), vec![1u8; 32]).unwrap(); - root = insert_leaf(&mut root, &NibbleKey::from(vec![1u8; 32]), vec![1u8; 32]).unwrap(); - assert_eq!( - root, - Branch(vec![ - Leaf(NibbleKey::from(vec![0u8; 31]), vec![1u8; 32]), - Leaf(NibbleKey::from(vec![1u8; 31]), vec![1u8; 32]), - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot - ]) - ); - } - - #[test] - fn insert_leaf_into_empty_root() { - let children = vec![EmptySlot; 16]; - let mut root = Branch(children); - let out = insert_leaf(&mut root, &NibbleKey::from(vec![0u8; 32]), vec![1u8; 32]); - assert_eq!( - out.unwrap(), - Branch(vec![ - Leaf(NibbleKey::from(vec![0u8; 31]), vec![1u8; 32]), - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot - ]) - ); - } - - #[test] - fn insert_leaf_into_two_level_fullnodes() { - let mut root = Branch(vec![ - Branch(vec![EmptySlot; 16]), - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - ]); - let out = insert_leaf(&mut root, &NibbleKey::from(vec![0u8; 32]), vec![1u8; 32]); - assert_eq!( - out.unwrap(), - Branch(vec![ - Branch(vec![ - Leaf(NibbleKey::from(vec![0u8; 30]), vec![1u8; 32]), - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot - ]), - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot, - EmptySlot - ]) - ); - } - - #[test] - fn insert_leaf_two_embedded_nodes() { - let expected_root = - hex::decode(&"0x1e1f73cc50d595585797d261df5be9bf69037a57e6470c9e4ffc87b6221ab67a"[2..]) - .unwrap(); - let inputs = [ - ("0x1111111111111111111111111111111111111111", "0xffff"), - ("0x2222222222222222222222222222222222222222", "0xeeee"), - ]; - - let mut root = Branch(vec![EmptySlot; 16]); - for (ik, iv) in inputs.iter() { - let k = NibbleKey::from(utils::ByteKey(hex::decode(&ik[2..]).unwrap())); - let v = hex::decode(&iv[2..]).unwrap(); - insert_leaf(&mut root, &k, v).unwrap(); - } - - let root_hash = root.hash(); - assert_eq!(expected_root, root_hash); - } - #[test] fn tree_with_just_one_leaf() { let proof = Multiproof { @@ -1164,38 +645,35 @@ mod tests { #[test] fn roundtrip() { let mut tree_root = Node::Branch(vec![Node::EmptySlot; 16]); - let new_root = insert_leaf( - &mut tree_root, - &NibbleKey::from(vec![1u8; 32]), - vec![2u8; 32], - ) - .unwrap(); + tree_root + .insert(&NibbleKey::from(vec![1u8; 32]), vec![2u8; 32]) + .unwrap(); assert_eq!( - new_root.hash(), + tree_root.hash(), vec![ 86, 102, 96, 191, 106, 199, 70, 178, 131, 236, 157, 14, 50, 168, 100, 69, 123, 66, 223, 122, 0, 97, 18, 144, 20, 79, 250, 219, 73, 190, 134, 108 ] ); - let proof = make_multiproof(&new_root, vec![NibbleKey::from(vec![1u8; 32])]).unwrap(); + let proof = make_multiproof(&tree_root, vec![NibbleKey::from(vec![1u8; 32])]).unwrap(); // RLP roundtrip let proof_rlp = rlp::encode(&proof); let proof = rlp::decode(&proof_rlp).unwrap(); let rebuilt_root = rebuild(&proof).unwrap(); - assert_eq!(new_root, rebuilt_root); + assert_eq!(tree_root, rebuilt_root); } #[test] fn test_nullkey_leaf() { let missing_key = NibbleKey::from(vec![1u8; 32]); let mut root = Node::default(); - root = insert_leaf(&mut root, &NibbleKey::from(vec![0u8; 32]), vec![0u8; 32]).unwrap(); - root = insert_leaf( - &mut root, + root.insert(&NibbleKey::from(vec![0u8; 32]), vec![0u8; 32]) + .unwrap(); + root.insert( &NibbleKey::from(ByteKey::from( hex::decode("11111111111111110000000000000000").unwrap(), )), @@ -1259,16 +737,14 @@ mod tests { fn test_nullkey_ext() { let missing_key = NibbleKey::from(vec![1u8; 32]); let mut root = Node::default(); - root = insert_leaf( - &mut root, + root.insert( &NibbleKey::from(ByteKey::from( hex::decode("11111111111111182222222222222222").unwrap(), )), vec![0u8; 32], ) .unwrap(); - root = insert_leaf( - &mut root, + root.insert( &NibbleKey::from(ByteKey::from( hex::decode("11111111111111180000000000000000").unwrap(), )), @@ -1305,16 +781,14 @@ mod tests { fn test_nullkey_branch() { let missing_key = NibbleKey::from(vec![2u8; 32]); let mut root = Node::default(); - root = insert_leaf( - &mut root, + root.insert( &NibbleKey::from(ByteKey::from( hex::decode("11111111111111182222222222222222").unwrap(), )), vec![0u8; 32], ) .unwrap(); - root = insert_leaf( - &mut root, + root.insert( &NibbleKey::from(ByteKey::from( hex::decode("01111111111111180000000000000000").unwrap(), )), @@ -1388,9 +862,9 @@ mod tests { let k3 = &NibbleKey::from(ByteKey::from( hex::decode("00111111111111110000000000000000").unwrap(), )); - root = insert_leaf(&mut root, k1, vec![0u8; 32]).unwrap(); - root = insert_leaf(&mut root, k2, vec![0u8; 32]).unwrap(); - root = insert_leaf(&mut root, k3, vec![0u8; 32]).unwrap(); + root.insert(k1, vec![0u8; 32]).unwrap(); + root.insert(k2, vec![0u8; 32]).unwrap(); + root.insert(k3, vec![0u8; 32]).unwrap(); assert_eq!( root[&NibbleKey::from(&k2[..2])], @@ -1410,9 +884,9 @@ mod tests { let k3 = &NibbleKey::from(ByteKey::from( hex::decode("11111111111111120000000000000000").unwrap(), )); - root = insert_leaf(&mut root, k1, vec![0u8; 32]).unwrap(); - root = insert_leaf(&mut root, k2, vec![0u8; 32]).unwrap(); - root = insert_leaf(&mut root, k3, vec![0u8; 32]).unwrap(); + root.insert(k1, vec![0u8; 32]).unwrap(); + root.insert(k2, vec![0u8; 32]).unwrap(); + root.insert(k3, vec![0u8; 32]).unwrap(); // check that a partial key returns the whole key assert_eq!(root[&NibbleKey::from(&k2[..2])], root); @@ -1435,9 +909,9 @@ mod tests { let k3 = &NibbleKey::from(ByteKey::from( hex::decode("00111111111111110000000000000000").unwrap(), )); - root = insert_leaf(&mut root, k1, vec![0u8; 32]).unwrap(); - root = insert_leaf(&mut root, k2, vec![0u8; 32]).unwrap(); - root = insert_leaf(&mut root, k3, vec![0u8; 32]).unwrap(); + root.insert(k1, vec![0u8; 32]).unwrap(); + root.insert(k2, vec![0u8; 32]).unwrap(); + root.insert(k3, vec![0u8; 32]).unwrap(); assert_eq!(root[&NibbleKey::from(vec![])], root); } @@ -1455,9 +929,9 @@ mod tests { hex::decode("11111111111111111111111111111111").unwrap(), )); - root = insert_leaf(&mut root, k1, vec![0u8; 32]).unwrap(); - root = insert_leaf(&mut root, k2, vec![0u8; 32]).unwrap(); - root = insert_leaf(&mut root, k3, vec![0u8; 32]).unwrap(); + root.insert(k1, vec![0u8; 32]).unwrap(); + root.insert(k2, vec![0u8; 32]).unwrap(); + root.insert(k3, vec![0u8; 32]).unwrap(); assert!(root.is_key_present(k1)); } @@ -1475,8 +949,8 @@ mod tests { hex::decode("11111111111111111111111111111111").unwrap(), )); - root = insert_leaf(&mut root, k2, vec![0u8; 32]).unwrap(); - root = insert_leaf(&mut root, k3, vec![0u8; 32]).unwrap(); + root.insert(k2, vec![0u8; 32]).unwrap(); + root.insert(k3, vec![0u8; 32]).unwrap(); assert!(!root.is_key_present(k1)); } @@ -1485,7 +959,8 @@ mod tests { fn check_payload_length_exactly_32_bytes() { let mut root = Node::default(); - root = insert_leaf(&mut root, &NibbleKey::from(vec![1u8; 16]), vec![1u8; 20]).unwrap(); + root.insert(&NibbleKey::from(vec![1u8; 16]), vec![1u8; 20]) + .unwrap(); assert_eq!(root.hash().len(), 32); assert_eq!( root.hash(), @@ -1500,7 +975,8 @@ mod tests { fn check_leaf_length_less_than_32_bytes() { let mut root = Node::default(); - root = insert_leaf(&mut root, &NibbleKey::from(vec![1u8; 2]), vec![1u8; 20]).unwrap(); + root.insert(&NibbleKey::from(vec![1u8; 2]), vec![1u8; 20]) + .unwrap(); assert_eq!( root.composition(), vec![216, 130, 32, 17, 148, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] @@ -1517,8 +993,10 @@ mod tests { #[test] fn check_branch_less_than_32_bytes() { let mut root = Node::default(); - root = insert_leaf(&mut root, &NibbleKey::from(vec![1u8; 4]), vec![1u8; 2]).unwrap(); - root = insert_leaf(&mut root, &NibbleKey::from(vec![2u8; 4]), vec![1u8; 2]).unwrap(); + root.insert(&NibbleKey::from(vec![1u8; 4]), vec![1u8; 2]) + .unwrap(); + root.insert(&NibbleKey::from(vec![2u8; 4]), vec![1u8; 2]) + .unwrap(); assert_eq!( root.composition(), @@ -1542,8 +1020,10 @@ mod tests { second_key.extend(vec![2u8; 2]); let mut root = Node::default(); - root = insert_leaf(&mut root, &NibbleKey::from(vec![1u8; 4]), vec![1u8; 2]).unwrap(); - root = insert_leaf(&mut root, &NibbleKey::from(second_key), vec![1u8; 2]).unwrap(); + root.insert(&NibbleKey::from(vec![1u8; 4]), vec![1u8; 2]) + .unwrap(); + root.insert(&NibbleKey::from(second_key), vec![1u8; 2]) + .unwrap(); assert_eq!( root.composition(), diff --git a/src/tree.rs b/src/tree.rs index 2bdbac1..6d5d814 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -22,6 +22,7 @@ pub trait Tree { fn ith_child(&self, index: usize) -> Option<&dyn Tree>; /// Returns an iterator to the node's children. Some of these nodes can be empty. fn children(&self) -> NodeChildIterator; + /// Insert a `(key,value)` pair into a (sub-)tree represented by `root`. fn insert(&mut self, key: &Self::Key, value: Self::Value) -> Result<(), String>; } From 7beb0e0d964916d212ad9960e71c4592e8740c1d Mon Sep 17 00:00:00 2001 From: Guillaume Ballet Date: Thu, 21 Nov 2019 19:02:51 +0100 Subject: [PATCH 07/10] multiproof generation independent from tree type --- src/lib.rs | 169 +++++++++++++++++++++++----------------------- src/multiproof.rs | 6 ++ src/node.rs | 89 +++++++++++++++++------- src/tree.rs | 37 +++++++--- 4 files changed, 180 insertions(+), 121 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index ad32df2..d99b661 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -12,103 +12,100 @@ pub mod utils; use instruction::*; use multiproof::*; use node::*; +use tree::{KeyValueStore, Tree}; use utils::*; -// Rebuilds the tree based on the multiproof components -pub fn rebuild(proof: &Multiproof) -> Result { - use Instruction::*; - use Node::*; +impl + rlp::Decodable> ProofToTree for Multiproof { + fn rebuild(&self) -> Result { + use Instruction::*; - let mut hiter = proof.hashes.iter(); - let iiter = proof.instructions.iter(); - let mut kviter = proof.keyvals.iter().map(|encoded| { - // Deserialize the keys as they are read - rlp::decode::(encoded).unwrap() - }); + let mut hiter = self.hashes.iter(); + let iiter = self.instructions.iter(); + let mut kviter = self.keyvals.iter().map(|encoded| { + // Deserialize the keys as they are read + rlp::decode::(encoded).unwrap() + }); - let mut stack = Vec::::new(); + let mut stack = Vec::::new(); - for instr in iiter { - match instr { - HASHER => { - if let Some(h) = hiter.next() { - stack.push(Hash(h.to_vec())); - } else { - return Err(format!("Proof requires one more hash in HASHER")); + for instr in iiter { + match instr { + HASHER => { + if let Some(h) = hiter.next() { + stack.push(T::from_hash(h.to_vec())); + } else { + return Err(format!("Proof requires one more hash in HASHER")); + } } - } - LEAF(keylength) => { - if let Some(Leaf(key, value)) = kviter.next() { - // If the value is empty, we have a NULL key and - // therefore an EmptySlot should be returned. - if value.len() != 0 { - stack.push(Leaf( - NibbleKey::from(key[key.len() - *keylength..].to_vec()), - value.to_vec(), - )); + LEAF(keylength) => { + if let Some(leaf) = kviter.next() { + // If the value is empty, we have a NULL key and + // therefore an EmptySlot should be returned. + match leaf.value() { + None => stack.push(T::new_empty()), + Some(_) if leaf.value_length().unwrap() == 0usize => { + stack.push(T::new_empty()) + } + Some(_) => stack.push(leaf), + } } else { - stack.push(EmptySlot); + return Err(format!( + "Proof requires one more (key,value) pair in LEAF({})", + keylength + )); } - } else { - return Err(format!( - "Proof requires one more (key,value) pair in LEAF({})", - keylength - )); } - } - BRANCH(digit) => { - if let Some(node) = stack.pop() { - let mut children = vec![Node::EmptySlot; 16]; - children[*digit] = node; - stack.push(Branch(children)) - } else { - return Err(format!( - "Could not pop a value from the stack, that is required for a BRANCH({})", - digit - )); + BRANCH(digit) => { + if let Some(ref node) = stack.pop() { + let mut b = T::new_branch(); + b.set_ith_child(*digit, node); + stack.push(b) + } else { + return Err(format!( + "Could not pop a value from the stack, that is required for a BRANCH({})", + digit + )); + } } - } - EXTENSION(key) => { - if let Some(node) = stack.pop() { - stack.push(Extension(NibbleKey::from(key.to_vec()), Box::new(node))); - } else { - return Err(format!( - "Could not find a node on the stack, that is required for an EXTENSION({:?})", - key - )); + EXTENSION(key) => { + if let Some(node) = stack.pop() { + stack.push(T::new_extension(key.to_vec(), /*Box::new(node)*/ node)); + } else { + return Err(format!( + "Could not find a node on the stack, that is required for an EXTENSION({:?})", + key + )); + } } - } - ADD(digit) => { - if let (Some(el1), Some(el2)) = (stack.pop(), stack.last_mut()) { - match el2 { - Branch(ref mut n2) => { - if *digit >= n2.len() { + ADD(digit) => { + if let (Some(el1), Some(el2)) = (stack.pop(), stack.last_mut()) { + // Only true if this is a branch node + if el2.num_children() > 1 { + if *digit >= el2.num_children() { return Err(format!( "Incorrect full node index: {} > {}", digit, - n2.len() - 1 + el2.num_children() )); } - // A hash needs to be fed into the hash sponge, any other node is simply - // a child (el1) of the parent node (el2). this is done during resolve. - n2[*digit] = el1; - } - Hash(_) => { - return Err(String::from("Hash node no longer supported in this case")); + // Any node is simply a child (el1) of the parent node (el2). This is done + // during resolve. + el2.set_ith_child(*digit, &el1); + } else { + return Err(String::from("Could not find enough parameters to ADD")); } - _ => return Err(String::from("Unexpected node type")), + } else { + return Err(String::from("Could not find enough parameters to ADD")); } - } else { - return Err(String::from("Could not find enough parameters to ADD")); } } } - } - stack - .pop() - .ok_or(String::from("Stack underflow, expected root node")) + stack + .pop() + .ok_or(String::from("Stack underflow, expected root node")) + } } // Helper function that generates a multiproof based on one `(key.value)` @@ -435,7 +432,7 @@ mod tests { let proof = make_multiproof(&root, vec![NibbleKey::from(vec![2u8; 32])]).unwrap(); - let res = rebuild(&proof); + let res = proof.rebuild(); assert_eq!(res.unwrap().hash(), pre_root_hash); } @@ -464,7 +461,7 @@ mod tests { ]; let proof = make_multiproof(&root, keys).unwrap(); - let res = rebuild(&proof); + let res = proof.rebuild(); assert_eq!(res.unwrap().hash(), pre_root_hash); } @@ -527,7 +524,7 @@ mod tests { ])], instructions: vec![LEAF(0)], }; - let out = rebuild(&proof).unwrap(); + let out = proof.rebuild().unwrap(); assert_eq!(out, Leaf(NibbleKey::from(vec![]), vec![4, 5, 6])) } @@ -541,7 +538,7 @@ mod tests { ])], instructions: vec![LEAF(0), BRANCH(0)], }; - let out = rebuild(&proof).unwrap(); + let out = proof.rebuild().unwrap(); assert_eq!( out, Branch(vec![ @@ -575,7 +572,7 @@ mod tests { ], instructions: vec![LEAF(0), BRANCH(0), LEAF(1), ADD(2)], }; - let out = rebuild(&proof).unwrap(); + let out = proof.rebuild().unwrap(); assert_eq!( out, Branch(vec![ @@ -615,7 +612,7 @@ mod tests { rlp::encode_list::, Vec>(&vec![vec![7, 8, 9], vec![10, 11, 12]]), ], }; - let out = rebuild(&proof).unwrap(); + let out = proof.rebuild().unwrap(); assert_eq!( out, Extension( @@ -661,9 +658,9 @@ mod tests { // RLP roundtrip let proof_rlp = rlp::encode(&proof); - let proof = rlp::decode(&proof_rlp).unwrap(); + let proof: Multiproof = rlp::decode(&proof_rlp).unwrap(); - let rebuilt_root = rebuild(&proof).unwrap(); + let rebuilt_root = proof.rebuild().unwrap(); assert_eq!(tree_root, rebuilt_root); } @@ -698,7 +695,7 @@ mod tests { ); assert_eq!(proof.instructions.len(), 4); - let rebuilt = rebuild(&proof).unwrap(); + let rebuilt = proof.rebuild().unwrap(); assert_eq!( rebuilt, Branch(vec![ @@ -763,7 +760,7 @@ mod tests { assert_eq!(proof.keyvals.len(), 0); assert_eq!(proof.instructions.len(), 2); - let rebuilt = rebuild(&proof).unwrap(); + let rebuilt = proof.rebuild().unwrap(); assert_eq!( rebuilt, Extension( @@ -809,7 +806,7 @@ mod tests { assert_eq!(proof.keyvals.len(), 1); assert_eq!(proof.instructions.len(), 6); - let rebuilt = rebuild(&proof).unwrap(); + let rebuilt = proof.rebuild().unwrap(); assert_eq!( rebuilt, Branch(vec![ @@ -846,7 +843,7 @@ mod tests { let missing_key = NibbleKey::from(vec![2u8; 32]); let proof = make_multiproof(&root, vec![missing_key.clone()]).unwrap(); - let rebuilt = rebuild(&proof).unwrap(); + let rebuilt = proof.rebuild().unwrap(); assert!(!rebuilt.is_key_present(&missing_key)); } diff --git a/src/multiproof.rs b/src/multiproof.rs index 4d36545..5703917 100644 --- a/src/multiproof.rs +++ b/src/multiproof.rs @@ -1,4 +1,10 @@ use super::instruction::*; +use super::tree::{KeyValueStore, Tree}; + +pub trait ProofToTree> { + /// Rebuilds a tree of type `T` based on the proof's components. + fn rebuild(&self) -> Result; +} #[derive(Debug, PartialEq)] pub struct Multiproof { diff --git a/src/node.rs b/src/node.rs index eaeb28b..9cd3128 100644 --- a/src/node.rs +++ b/src/node.rs @@ -13,10 +13,12 @@ pub enum Node { EmptySlot, } -impl Tree for Node { +impl KeyValueStore for Node { type Key = NibbleKey; type Value = Vec; +} +impl Tree for Node { fn is_leaf(&self) -> bool { match self { Node::Leaf(_, _) => true, @@ -45,7 +47,7 @@ impl Tree for Node { /// Return the tree root's *ith* child. /// /// The function will return `EmptySlot` instead of `None` if a branch node has no *ith* child. - fn ith_child(&self, index: usize) -> Option<&dyn Tree> { + fn ith_child(&self, index: usize) -> Option<&Self> { if index >= self.num_children() { panic!(format!( "Requested child #{}, only have #{}", @@ -55,33 +57,38 @@ impl Tree for Node { } match self { - Node::EmptySlot | Node::Leaf(_, _) | Node::Hash(_) => { - None as Option<&dyn Tree> - } + Node::EmptySlot | Node::Leaf(_, _) | Node::Hash(_) => None, Node::Extension(_, box ext) => Some(&*ext), Node::Branch(ref vec) => Some(&vec[index]), } } - fn children(&self) -> NodeChildIterator { - match self { - Node::EmptySlot => NodeChildIterator { - index: 0, - node: self, - }, - Node::Leaf(_, _) => NodeChildIterator { - index: 0, - node: self, - }, - Node::Extension(_, _) => NodeChildIterator { - index: 0, - node: self, - }, - Node::Branch(_) => NodeChildIterator { - index: 0, - node: self, - }, - _ => panic!("not implemented"), + fn set_ith_child(&mut self, index: usize, child: &Self) { + if index >= self.num_children() { + panic!(format!( + "Requested child #{}, only have #{}", + index, + self.num_children() + )); + } + + if let Node::Branch(ref mut vec) = self { + if let Node::EmptySlot = &vec[index] { + // NOTE attempt to use pointers instead of cloning + vec[index] = child.clone(); + } else { + panic!("Refusing to overwrite child node"); + } + } else { + panic!("Only branch nodes can be set in this implementation."); + } + } + + fn children(&self) -> NodeChildIterator { + NodeChildIterator { + index: 0, + key: None, + node: &self, } } @@ -208,6 +215,40 @@ impl Tree for Node { _ => panic!("Can not insert a node into a hashed node"), } } + + fn value(&self) -> Option<&Vec> { + match self { + Node::Leaf(_, ref v) => Some(v), + _ => None, + } + } + + fn value_length(&self) -> Option { + match self { + Node::Leaf(_, ref v) => Some(v.len()), + _ => None, + } + } + + fn from_hash(h: Vec) -> Self { + Node::Hash(h.to_vec()) + } + + fn new_extension(ext: Vec, child: Self) -> Self { + Node::Extension(NibbleKey::from(ext), Box::new(child)) + } + + fn new_branch() -> Self { + Node::Branch(vec![Node::EmptySlot; 16]) + } + + fn new_leaf(key: Vec, value: Vec) -> Self { + Node::Leaf(NibbleKey::from(key), value) + } + + fn new_empty() -> Self { + Node::EmptySlot + } } impl Default for Node { diff --git a/src/tree.rs b/src/tree.rs index 6d5d814..33c753e 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -1,17 +1,21 @@ -use super::utils::*; - -pub struct NodeChildIterator<'a, K, V> { +pub struct NodeChildIterator<'a, S: KeyValueStore, T: Tree> { + /// The iterator's current child index. pub index: usize, - pub node: &'a dyn Tree, + /// The node's current key, if available. + pub key: Option<&'a S::Key>, + /// The node whose children are being iterated on. + pub node: &'a T, } -/// A trait representing the an underlying tree structure. -pub trait Tree { +pub trait KeyValueStore { /// The tree's key type. Must be specified when implementing this trait. type Key; /// The tree's value type. Must be specified when implementing this trait. type Value; +} +/// A trait representing the an underlying tree structure. +pub trait Tree: Sized { /// Specifies if the current tree is a simple leaf. fn is_leaf(&self) -> bool; /// Specifies if the current tree is empty. @@ -19,15 +23,26 @@ pub trait Tree { /// Returns the tree root's *maximal* number of children. fn num_children(&self) -> usize; /// Returns a pointer to child #i, or `None` if no such child exists. - fn ith_child(&self, index: usize) -> Option<&dyn Tree>; + fn ith_child(&self, index: usize) -> Option<&Self>; + /// Set child node #i of the current node. + fn set_ith_child(&mut self, index: usize, child: &Self); /// Returns an iterator to the node's children. Some of these nodes can be empty. - fn children(&self) -> NodeChildIterator; + fn children(&self) -> NodeChildIterator; /// Insert a `(key,value)` pair into a (sub-)tree represented by `root`. - fn insert(&mut self, key: &Self::Key, value: Self::Value) -> Result<(), String>; + fn insert(&mut self, key: &S::Key, value: S::Value) -> Result; + + fn value(&self) -> Option<&S::Value>; + fn value_length(&self) -> Option; + + fn from_hash(h: Vec) -> Self; + fn new_empty() -> Self; + fn new_extension(ext: Vec, child: Self) -> Self; + fn new_branch() -> Self; + fn new_leaf(key: Vec, value: Vec) -> Self; } -impl<'a> std::iter::Iterator for NodeChildIterator<'a, NibbleKey, Vec> { - type Item = &'a dyn Tree>; +impl<'a, S: KeyValueStore, T: Tree> std::iter::Iterator for NodeChildIterator<'a, S, T> { + type Item = &'a T; #[inline] fn next(&mut self) -> Option { From e7ff767293dc9281d9aa2d3a83f5a9f9cb26ee70 Mon Sep 17 00:00:00 2001 From: Guillaume Ballet Date: Thu, 21 Nov 2019 19:37:01 +0100 Subject: [PATCH 08/10] Fix test build --- src/lib.rs | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index d99b661..65a0c95 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -284,7 +284,7 @@ mod tests { keyvals: proof.keyvals, instructions: proof.instructions, }; - let new_root = rebuild(&proof).unwrap(); + let new_root: Node = proof.rebuild().unwrap(); assert_eq!( new_root, @@ -432,9 +432,9 @@ mod tests { let proof = make_multiproof(&root, vec![NibbleKey::from(vec![2u8; 32])]).unwrap(); - let res = proof.rebuild(); + let res: Node = proof.rebuild().unwrap(); - assert_eq!(res.unwrap().hash(), pre_root_hash); + assert_eq!(res.hash(), pre_root_hash); } #[test] @@ -461,9 +461,9 @@ mod tests { ]; let proof = make_multiproof(&root, keys).unwrap(); - let res = proof.rebuild(); + let res: Node = proof.rebuild().unwrap(); - assert_eq!(res.unwrap().hash(), pre_root_hash); + assert_eq!(res.hash(), pre_root_hash); } #[test] @@ -524,7 +524,7 @@ mod tests { ])], instructions: vec![LEAF(0)], }; - let out = proof.rebuild().unwrap(); + let out: Node = proof.rebuild().unwrap(); assert_eq!(out, Leaf(NibbleKey::from(vec![]), vec![4, 5, 6])) } @@ -538,7 +538,7 @@ mod tests { ])], instructions: vec![LEAF(0), BRANCH(0)], }; - let out = proof.rebuild().unwrap(); + let out: Node = proof.rebuild().unwrap(); assert_eq!( out, Branch(vec![ @@ -572,7 +572,7 @@ mod tests { ], instructions: vec![LEAF(0), BRANCH(0), LEAF(1), ADD(2)], }; - let out = proof.rebuild().unwrap(); + let out: Node = proof.rebuild().unwrap(); assert_eq!( out, Branch(vec![ @@ -612,7 +612,7 @@ mod tests { rlp::encode_list::, Vec>(&vec![vec![7, 8, 9], vec![10, 11, 12]]), ], }; - let out = proof.rebuild().unwrap(); + let out: Node = proof.rebuild().unwrap(); assert_eq!( out, Extension( @@ -695,7 +695,7 @@ mod tests { ); assert_eq!(proof.instructions.len(), 4); - let rebuilt = proof.rebuild().unwrap(); + let rebuilt: Node = proof.rebuild().unwrap(); assert_eq!( rebuilt, Branch(vec![ @@ -760,7 +760,7 @@ mod tests { assert_eq!(proof.keyvals.len(), 0); assert_eq!(proof.instructions.len(), 2); - let rebuilt = proof.rebuild().unwrap(); + let rebuilt: Node = proof.rebuild().unwrap(); assert_eq!( rebuilt, Extension( @@ -806,7 +806,7 @@ mod tests { assert_eq!(proof.keyvals.len(), 1); assert_eq!(proof.instructions.len(), 6); - let rebuilt = proof.rebuild().unwrap(); + let rebuilt: Node = proof.rebuild().unwrap(); assert_eq!( rebuilt, Branch(vec![ @@ -843,7 +843,7 @@ mod tests { let missing_key = NibbleKey::from(vec![2u8; 32]); let proof = make_multiproof(&root, vec![missing_key.clone()]).unwrap(); - let rebuilt = proof.rebuild().unwrap(); + let rebuilt: Node = proof.rebuild().unwrap(); assert!(!rebuilt.is_key_present(&missing_key)); } From 500b3ac26942c0d8c738d48f936eba75bfdb7ccc Mon Sep 17 00:00:00 2001 From: Guillaume Ballet Date: Thu, 21 Nov 2019 19:37:24 +0100 Subject: [PATCH 09/10] Fix cargo build --- Cargo.toml | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 64d07f2..5e65af9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,8 +9,5 @@ edition = "2018" [dependencies] sha3 = "0.8.2" rlp = "0.4.2" -hex = "0.3.2" +hex = "0.4.0" serde_json = "1.0.40" - -[dev-dependencies] -hex = "0.3.2" From 7a6af4a06ab27ca2fabadcd68890ddcc1964e9ac Mon Sep 17 00:00:00 2001 From: Guillaume Ballet Date: Fri, 22 Nov 2019 11:04:55 +0100 Subject: [PATCH 10/10] Fix unit tests This commit effectively makes the parameter in LEAF() obsolete, because it is assumed that the leaf only contains the label that it needs. --- src/lib.rs | 23 +++++++++++++---------- src/node.rs | 4 +--- src/utils.rs | 27 +++++++++++++++++++++++++++ 3 files changed, 41 insertions(+), 13 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 65a0c95..4117ee9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -43,10 +43,13 @@ impl + rlp::Decodable> ProofToTree for Multip // therefore an EmptySlot should be returned. match leaf.value() { None => stack.push(T::new_empty()), - Some(_) if leaf.value_length().unwrap() == 0usize => { - stack.push(T::new_empty()) + Some(_) => { + if leaf.value_length().unwrap() == 0usize { + stack.push(T::new_empty()) + } else { + stack.push(leaf) + } } - Some(_) => stack.push(leaf), } } else { return Err(format!( @@ -69,7 +72,7 @@ impl + rlp::Decodable> ProofToTree for Multip } EXTENSION(key) => { if let Some(node) = stack.pop() { - stack.push(T::new_extension(key.to_vec(), /*Box::new(node)*/ node)); + stack.push(T::new_extension(key.to_vec(), node)); } else { return Err(format!( "Could not find a node on the stack, that is required for an EXTENSION({:?})", @@ -519,7 +522,7 @@ mod tests { let proof = Multiproof { hashes: vec![], keyvals: vec![rlp::encode_list::, Vec>(&vec![ - vec![1, 2, 3], + vec![], vec![4, 5, 6], ])], instructions: vec![LEAF(0)], @@ -533,7 +536,7 @@ mod tests { let proof = Multiproof { hashes: vec![], keyvals: vec![rlp::encode_list::, Vec>(&vec![ - vec![1, 2, 3], + vec![], vec![4, 5, 6], ])], instructions: vec![LEAF(0), BRANCH(0)], @@ -567,8 +570,8 @@ mod tests { let proof = Multiproof { hashes: vec![], keyvals: vec![ - rlp::encode_list::, Vec>(&vec![vec![1, 2, 3], vec![4, 5, 6]]), - rlp::encode_list::, Vec>(&vec![vec![7, 8, 9], vec![10, 11, 12]]), + rlp::encode_list::, Vec>(&vec![vec![], vec![4, 5, 6]]), + rlp::encode_list::, Vec>(&vec![vec![25], vec![10, 11, 12]]), ], instructions: vec![LEAF(0), BRANCH(0), LEAF(1), ADD(2)], }; @@ -608,8 +611,8 @@ mod tests { EXTENSION(vec![13, 14, 15]), ], keyvals: vec![ - rlp::encode_list::, Vec>(&vec![vec![1, 2, 3], vec![4, 5, 6]]), - rlp::encode_list::, Vec>(&vec![vec![7, 8, 9], vec![10, 11, 12]]), + rlp::encode_list::, Vec>(&vec![vec![], vec![4, 5, 6]]), + rlp::encode_list::, Vec>(&vec![vec![25], vec![10, 11, 12]]), ], }; let out: Node = proof.rebuild().unwrap(); diff --git a/src/node.rs b/src/node.rs index 9cd3128..51f11c2 100644 --- a/src/node.rs +++ b/src/node.rs @@ -299,9 +299,7 @@ impl rlp::Decodable for Node { return Err(rlp::DecoderError::RlpExpectedToBeList); } let keyval = rlp.as_list::>()?; - let key_bytes = ByteKey(keyval[0].clone()); - let key_nibbles = NibbleKey::from(key_bytes); - // TODO: remove indicator prefix if node is a leaf or extension + let key_nibbles = NibbleKey::remove_hex_prefix(&keyval[0]); Ok(Node::Leaf(key_nibbles, keyval[1].clone())) } } diff --git a/src/utils.rs b/src/utils.rs index 370c575..dc37d50 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -86,6 +86,33 @@ impl NibbleKey { output } + + pub fn remove_hex_prefix(payload: &Vec) -> NibbleKey { + if payload.len() == 0 { + return NibbleKey(payload.clone()); + } + match payload[0] { + x if x & 16 == 16 => { + // Odd payload.len() + let mut out = vec![0u8; (payload.len() - 1) * 2 + 1]; + out[0] = x & 0xF; + for i in 1..payload.len() { + out[2 * i - 1] = payload[i] >> 4; + out[2 * i] = payload[i] & 0xF; + } + NibbleKey(out) + } + _ => { + // Even payload.len() + let mut out = vec![0u8; (payload.len() - 1) * 2]; + for i in 1..payload.len() { + out[2 * (i - 1)] = payload[i - 1] & 0xF; + out[2 * (i - 1) + 1] = payload[i] >> 4; + } + NibbleKey(out) + } + } + } } impl rlp::Encodable for NibbleKey {