Mpt trie panic refactor

mpt_trie/examples/hash_nodes.rs

@@ -54,9 +54,10 @@ use mpt_trie::partial_trie::PartialTrie;
 use mpt_trie::{
     nibbles::Nibbles,
     partial_trie::{HashedPartialTrie, Node},
+    trie_ops::TrieOpResult,
 };
 
-fn main() {
+fn main() -> TrieOpResult<()> {
     pretty_env_logger::try_init().unwrap();
 
     // Lets build the (binary) tries in the module-level docs. Since the example
@@ -66,10 +67,10 @@ fn main() {
 
     // Note the nibbles read the most significant nibble first (eg. `0x12` reads `1`
     // first).
-    full_trie.insert(Nibbles::from_str("0x00").unwrap(), large_val(1)); // 1st from left.
-    full_trie.insert(Nibbles::from_str("0x01").unwrap(), large_val(2)); // 2nd from left.
-    full_trie.insert(Nibbles::from(0x10_u64), large_val(3)); // 3rd from left.
-    full_trie.insert(Nibbles::from(0x11_u64), large_val(4)); // 4th from left.
+    full_trie.insert(Nibbles::from_str("0x00").unwrap(), large_val(1))?; // 1st from left.
+    full_trie.insert(Nibbles::from_str("0x01").unwrap(), large_val(2))?; // 2nd from left.
+    full_trie.insert(Nibbles::from(0x10_u64), large_val(3))?; // 3rd from left.
+    full_trie.insert(Nibbles::from(0x11_u64), large_val(4))?; // 4th from left.
 
     let full_trie_hash = full_trie.hash();
 
@@ -83,14 +84,16 @@ fn main() {
 
     // Hash version. `0` branch is replaced with a `Hash` node.
     let mut hash_trie = HashedPartialTrie::default();
-    hash_trie.insert(Nibbles::from_str("0x0").unwrap(), left_side_hash); // Hash node
-    hash_trie.insert(0x10_u64, large_val(3)); // 3rd from left.
-    hash_trie.insert(0x11_u64, large_val(4)); // 4th from left.
+    hash_trie.insert(Nibbles::from_str("0x0").unwrap(), left_side_hash)?; // Hash node
+    hash_trie.insert(0x10_u64, large_val(3))?; // 3rd from left.
+    hash_trie.insert(0x11_u64, large_val(4))?; // 4th from left.
 
     let hash_trie_hash = hash_trie.hash();
 
     // Hashes should be equal.
     assert_eq!(full_trie_hash, hash_trie_hash);
+
+    Ok(())
 }
 
 /// We want to ensure that all leafs are >= 32 bytes when RLP encoded in order

mpt_trie/examples/simple.rs

@@ -6,18 +6,18 @@ use mpt_trie::partial_trie::PartialTrie;
 use mpt_trie::{
     nibbles::{Nibbles, ToNibbles},
     partial_trie::{HashedPartialTrie, StandardTrie},
-    trie_ops::ValOrHash,
+    trie_ops::{TrieOpResult, ValOrHash},
 };
 
-fn main() {
+fn main() -> TrieOpResult<()> {
     // Construct an empty trie:
     let mut trie = StandardTrie::default();
 
     // Elements can be inserted into the trie by calling insert directly:
     trie.insert(
         Nibbles::from_bytes_be(b"hello").unwrap(),
         b"world!".to_vec(),
-    );
+    )?;
 
     // Or by initializing the trie with an iterator of key value pairs:
     let mut trie = StandardTrie::from_iter(vec![
@@ -44,7 +44,7 @@ fn main() {
 
     // Values can be deleted:
     let del_val = trie.delete(0x1234_u32);
-    assert_eq!(del_val, Some(b"some data".to_vec()));
+    assert_eq!(del_val.ok().unwrap(), Some(b"some data".to_vec()));
     assert_eq!(trie.get(0x1234_u32), None);
 
     // It's important to note how types are converted to `Nibbles`. This is
@@ -76,4 +76,6 @@ fn main() {
         format!("{:x}", Nibbles::from_bytes_le(&[69, 35, 1]).unwrap()),
         "0x012345"
     );
+
+    Ok(())
 }

mpt_trie/src/debug_tools/diff.rs

@@ -423,11 +423,11 @@ mod tests {
     fn depth_single_node_hash_diffs_work() {
         // TODO: Reduce duplication once we identify common structures across tests...
         let mut a = HashedPartialTrie::default();
-        a.insert(0x1234, vec![0]);
+        assert!(a.insert(0x1234, vec![0]).is_ok());
         let a_hash = a.hash();
 
         let mut b = a.clone();
-        b.insert(0x1234, vec![1]);
+        assert!(b.insert(0x1234, vec![1]).is_ok());
         let b_hash = b.hash();
 
         let diff = create_diff_between_tries(&a, &b);

mpt_trie/src/debug_tools/stats.rs

@@ -323,7 +323,7 @@ mod tests {
 
     #[test]
     fn hand_made_trie_has_correct_node_stats() {
-        let (trie, _) = handmade_trie_1();
+        let (trie, _) = handmade_trie_1().unwrap();
         let stats = get_trie_stats(&trie);
 
         assert_eq!(stats.counts.leaf, 4);
@@ -370,8 +370,8 @@ mod tests {
         let hash_entries = generate_n_random_fixed_trie_hash_entries(n_hash_nodes, seed + 1);
 
         let mut trie = HashedPartialTrie::default();
-        trie.extend(val_entries);
-        trie.extend(hash_entries);
+        assert!(trie.extend(val_entries).is_ok());
+        assert!(trie.extend(hash_entries).is_ok());
 
         let stats = get_trie_stats(&trie);
 

mpt_trie/src/partial_trie.rs

@@ -13,7 +13,7 @@ use serde::{Deserialize, Serialize};
 use crate::{
     nibbles::Nibbles,
     trie_hashing::{hash_trie, rlp_encode_and_hash_node, EncodedNode},
-    trie_ops::ValOrHash,
+    trie_ops::{TrieOpResult, ValOrHash},
     utils::bytes_to_h256,
 };
 
@@ -58,13 +58,13 @@ pub trait PartialTrie:
     fn new(n: Node<Self>) -> Self;
 
     /// Inserts a node into the trie.
-    fn insert<K, V>(&mut self, k: K, v: V)
+    fn insert<K, V>(&mut self, k: K, v: V) -> TrieOpResult<()>
     where
         K: Into<Nibbles>,
         V: Into<ValOrHash>;
 
     /// Add more nodes to the trie through an iterator
-    fn extend<K, V, I>(&mut self, nodes: I)
+    fn extend<K, V, I>(&mut self, nodes: I) -> TrieOpResult<()>
     where
         K: Into<Nibbles>,
         V: Into<ValOrHash>,
@@ -89,7 +89,7 @@ pub trait PartialTrie:
     /// are meant for parts of the trie that are not relevant, traversing one
     /// means that a `Hash` node was created that potentially should not have
     /// been.
-    fn delete<K>(&mut self, k: K) -> Option<Vec<u8>>
+    fn delete<K>(&mut self, k: K) -> TrieOpResult<Option<Vec<u8>>>
     where
         K: Into<Nibbles>;
 
@@ -214,15 +214,16 @@ impl PartialTrie for StandardTrie {
         Self(n)
     }
 
-    fn insert<K, V>(&mut self, k: K, v: V)
+    fn insert<K, V>(&mut self, k: K, v: V) -> TrieOpResult<()>
     where
         K: Into<Nibbles>,
         V: Into<ValOrHash>,
     {
-        self.0.trie_insert(k, v);
+        self.0.trie_insert(k, v)?;
+        Ok(())
     }
 
-    fn extend<K, V, I>(&mut self, nodes: I)
+    fn extend<K, V, I>(&mut self, nodes: I) -> TrieOpResult<()>
     where
         K: Into<Nibbles>,
         V: Into<ValOrHash>,
@@ -238,7 +239,7 @@ impl PartialTrie for StandardTrie {
         self.0.trie_get(k)
     }
 
-    fn delete<K>(&mut self, k: K) -> Option<Vec<u8>>
+    fn delete<K>(&mut self, k: K) -> TrieOpResult<Option<Vec<u8>>>
     where
         K: Into<Nibbles>,
     {
@@ -327,23 +328,25 @@ impl PartialTrie for HashedPartialTrie {
         }
     }
 
-    fn insert<K, V>(&mut self, k: K, v: V)
+    fn insert<K, V>(&mut self, k: K, v: V) -> TrieOpResult<()>
     where
         K: Into<crate::nibbles::Nibbles>,
         V: Into<crate::trie_ops::ValOrHash>,
     {
-        self.node.trie_insert(k, v);
+        self.node.trie_insert(k, v)?;
         self.set_hash(None);
+        Ok(())
     }
 
-    fn extend<K, V, I>(&mut self, nodes: I)
+    fn extend<K, V, I>(&mut self, nodes: I) -> TrieOpResult<()>
     where
         K: Into<crate::nibbles::Nibbles>,
         V: Into<crate::trie_ops::ValOrHash>,
         I: IntoIterator<Item = (K, V)>,
     {
-        self.node.trie_extend(nodes);
+        self.node.trie_extend(nodes)?;
         self.set_hash(None);
+        Ok(())
     }
 
     fn get<K>(&self, k: K) -> Option<&[u8]>
@@ -353,7 +356,7 @@ impl PartialTrie for HashedPartialTrie {
         self.node.trie_get(k)
     }
 
-    fn delete<K>(&mut self, k: K) -> Option<Vec<u8>>
+    fn delete<K>(&mut self, k: K) -> TrieOpResult<Option<Vec<u8>>>
     where
         K: Into<crate::nibbles::Nibbles>,
     {
@@ -434,7 +437,8 @@ where
     V: Into<ValOrHash>,
 {
     let mut root = N::new(Node::Empty);
-    root.extend(nodes);
-
-    root
+    match root.extend(nodes) {
+        Ok(_) => root,
+        Err(e) => panic!("Error extending trie: {}", e),
+    }
 }

mpt_trie/src/special_query.rs

@@ -131,7 +131,7 @@ mod test {
     #[test]
     fn query_iter_works_no_last_node() {
         common_setup();
-        let (trie, ks) = handmade_trie_1();
+        let (trie, ks) = handmade_trie_1().unwrap();
 
         // ks --> vec![0x1234, 0x1324, 0x132400005_u64, 0x2001, 0x2002];
         let res = vec![
@@ -175,7 +175,7 @@ mod test {
     #[test]
     fn query_iter_works_with_last_node() {
         common_setup();
-        let (trie, _) = handmade_trie_1();
+        let (trie, _) = handmade_trie_1().unwrap();
 
         let extension_expected = vec![
             TrieSegment::Branch(1),

mpt_trie/src/testing_utils.rs

@@ -10,7 +10,7 @@ use rand::{rngs::StdRng, seq::IteratorRandom, Rng, RngCore, SeedableRng};
 use crate::{
     nibbles::Nibbles,
     partial_trie::{HashedPartialTrie, Node, PartialTrie},
-    trie_ops::ValOrHash,
+    trie_ops::{TrieOpResult, ValOrHash},
     utils::is_even,
 };
 
@@ -192,19 +192,21 @@ fn gen_rand_u256_bytes(rng: &mut StdRng) -> Vec<u8> {
 
 /// Initializes a trie with keys large enough to force hashing (nodes less than
 /// 32 bytes are not hashed).
-pub(crate) fn create_trie_with_large_entry_nodes<T: Into<Nibbles> + Copy>(keys: &[T]) -> TrieType {
+pub(crate) fn create_trie_with_large_entry_nodes<T: Into<Nibbles> + Copy>(
+    keys: &[T],
+) -> TrieOpResult<TrieType> {
     let mut trie = TrieType::default();
     for (k, v) in keys.iter().map(|k| (*k).into()).map(large_entry) {
-        trie.insert(k, v.clone());
+        trie.insert(k, v.clone())?;
     }
 
-    trie
+    Ok(trie)
 }
 
-pub(crate) fn handmade_trie_1() -> (TrieType, Vec<Nibbles>) {
+pub(crate) fn handmade_trie_1() -> TrieOpResult<(TrieType, Vec<Nibbles>)> {
     let ks = vec![0x1234, 0x1324, 0x132400005_u64, 0x2001, 0x2002];
     let ks_nibbles: Vec<Nibbles> = ks.into_iter().map(|k| k.into()).collect();
-    let trie = create_trie_with_large_entry_nodes(&ks_nibbles);
+    let trie = create_trie_with_large_entry_nodes(&ks_nibbles)?;
 
     // Branch (0x)  --> 1, 2
     // Branch (0x1) --> 2, 3
@@ -219,5 +221,5 @@ pub(crate) fn handmade_trie_1() -> (TrieType, Vec<Nibbles>) {
     // Leaf  (0x2001) --> (n: 0x1, v: [3])
     // Leaf  (0x2002) --> (n: 0x2, v: [4])
 
-    (trie, ks_nibbles)
+    Ok((trie, ks_nibbles))
 }

mpt_trie/src/trie_hashing.rs

@@ -232,7 +232,7 @@ mod tests {
         let mut trie = HashedPartialTrie::new(Node::Empty);
 
         entries.map(move |(k, v)| {
-            trie.insert(k, v);
+            assert!(trie.insert(k, v).is_ok());
             trie.get_hash()
         })
     }
@@ -376,7 +376,8 @@ mod tests {
         let half_entries = entries.len() / 2;
         let entries_to_delete = entries.into_iter().take(half_entries);
         for (k, _) in entries_to_delete {
-            our_trie.delete(k);
+            let res = our_trie.delete(k);
+            assert_eq!(res.is_ok(), true, "Failed to delete key: {:?}", k);
             truth_trie.remove(&k.bytes_be()).unwrap();
 
             let truth_root_hash = H256(truth_trie.root_hash().unwrap().0);