Model testing of v1 and v2 tries against a HasMap model

Author: Kristoffer Ström

crates/eth-sparse-mpt/src/model_test.rs

@@ -1,19 +1,17 @@
-use crate::v2::trie::Trie;
+use crate::v1::sparse_mpt::{DeletionError as DeletionErrorV1, DiffTrie};
+use crate::v2::trie::{DeletionError as DeletionErrorV2, Trie};
+
+use alloy_primitives::{hex, keccak256, Bytes, FixedBytes};
 use quickcheck::{quickcheck, Arbitrary, Gen};
 use std::collections::HashMap;
 
-// The maximum key size. keeping it relatively
-// small increases the chance of multiple
-// operations being executed against the same
-// key, which will tease out more bugs.
-const KEY_SPACE: u8 = 16;
-
 #[derive(Clone, Debug)]
 enum Op {
-    Insert(Vec<u8>, Vec<u8>),
-    Get(Vec<u8>),
+    Insert(FixedKey, Vec<u8>),
+    Delete(FixedKey),
 }
 
+// helper trait to extend `choose` with exception handling
 trait ChooseNonempty {
     fn one_of<'a, T>(&'a mut self, entries: &'a [T]) -> &'a T;
 }
@@ -24,46 +22,129 @@ impl ChooseNonempty for Gen {
     }
 }
 
-// Arbitrary lets you create randomized instances
-// of types that you're interested in testing
-// properties with. QuickCheck will look for
-// this trait for things that are the arguments
-// to properties that it is testing.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+struct FixedKey(FixedBytes<32>);
+
+impl FixedKey {
+    fn as_slice(&self) -> &[u8; 32] {
+        &self.0
+    }
+
+    fn from_string(s: &str) -> Self {
+        Self(keccak256("a"))
+    }
+
+    fn from_bytes(bytes: [u8; 32]) -> Self {
+        Self(FixedBytes::new(bytes))
+    }
+
+    fn into_bytes(self) -> Bytes {
+        Bytes::from(self.0)
+    }
+}
+
+impl From<Bytes> for FixedKey {
+    fn from(bytes: Bytes) -> FixedKey {
+        let fbytes = FixedBytes::from_slice(bytes.as_ref());
+        FixedKey(fbytes)
+    }
+}
+
+// We chose a small number of keys, to make sure our error cases handle key collisions,
+// as well as shared prefixes, properties that would be very unlikely for random keys
+impl Arbitrary for FixedKey {
+    fn arbitrary(g: &mut Gen) -> Self {
+        let keys = [
+            FixedKey::from_bytes(hex!(
+                "0000000000000000000000000000000000000000000000000000000000000000"
+            )),
+            FixedKey::from_bytes(hex!(
+                "0000000000000000000000000000000000000000000000000000000000000001"
+            )),
+            FixedKey::from_bytes(hex!(
+                "0000000000000000000000000000001000000000000000000000000000000001"
+            )),
+            FixedKey::from_string("0"),
+            FixedKey::from_string("1"),
+            FixedKey::from_string("2"),
+            FixedKey::from_string("3"),
+            FixedKey::from_string("4"),
+            FixedKey::from_string("5"),
+            FixedKey::from_string("6"),
+            FixedKey::from_string("7"),
+        ];
+        *g.one_of(&keys)
+    }
+}
+
 impl Arbitrary for Op {
     fn arbitrary(g: &mut Gen) -> Self {
         // pick a random key to perform an operation on
-        let key = g
-            .one_of(&["key00", "key01", "odd", "key010"])
-            .as_bytes()
-            .to_owned();
+        let key = FixedKey::arbitrary(g);
 
         if *g.one_of(&[true, false]) {
             Op::Insert(key, "value".into())
         } else {
-            Op::Get(key)
+            Op::Delete(key)
         }
     }
 }
 
+/// This test fails, since the Trie is designed for fixed key sizes
+#[ignore]
+#[test]
+fn crash_example_v2() {
+    let mut trie = Trie::new_empty();
+    trie.insert(b"00aeee", b"ok").unwrap();
+    trie.insert(b"00ae", b"ok").unwrap();
+}
+
 quickcheck! {
-    fn model_test_v2(ops: Vec<Op>) -> bool {
+    fn model_test_v1_map(ops: Vec<Op>) -> bool {
+        let mut model = HashMap::new();
+        let mut implementation = DiffTrie::new_empty();
+
+        for op in ops {
+            match op {
+                Op::Insert(key, value) => {
+                    implementation.insert(key.into_bytes(), Bytes::from(value.clone())).unwrap();
+                    model.insert(key, value);
+                }
+                Op::Delete(key) => {
+                    match (implementation.delete(key.into_bytes()), model.remove(&key)) {
+                        (Err(DeletionErrorV1::KeyNotFound), None) => (),
+                        (Err(err), _) => panic!("Implementation error {err:?}"),
+                        (Ok(_), Some(_)) => (),
+                        (Ok(returned), None) => panic!("Implementation returned {returned:?} on delete"),
+                    }
+                }
+            }
+        }
+        true
+    }
+}
+
+quickcheck! {
+    fn model_test_v2_map(ops: Vec<Op>) -> bool {
         let mut model = HashMap::new();
         let mut implementation = Trie::new_empty();
 
         for op in ops {
             match op {
                 Op::Insert(k, v) => {
-                    implementation.insert(k.as_slice(), v.as_slice());
+                    implementation.insert(k.as_slice(), v.as_slice()).unwrap();
                     model.insert(k, v);
                 }
-                Op::Get(k) => {
-                    // if implementation.get(&k) != model.get(&k).map(AsRef::as_ref) {
-                    //     return false;
-                    // }
+                Op::Delete(k) => {
+                    match (implementation.delete(k.as_slice()), model.remove(&k)) {
+                        (Err(DeletionErrorV2::KeyNotFound), None) => (),
+                        (Err(e), _) => panic!("Implementation error {e:?}"),
+                        (Ok(_), Some(_)) => (),
+                        (Ok(a), None) => panic!("Implementation returned {a:?} on delete"),
+                    }
                 }
             }
         }
-
         true
     }
 }

crates/eth-sparse-mpt/src/v2/trie/mod.rs

@@ -361,66 +361,6 @@ impl Trie {
         self.delete_nibbles_key(&n)
     }
 
-    pub fn get(&mut self, key: &[u8]) -> Option<&[u8]> {
-        let n = Nibbles::unpack(key);
-        self.get_nibbles_key(&n)
-    }
-
-    pub fn get_nibbles_key(&mut self, nibbles_key: &Nibbles) -> Option<&[u8]> {
-        let get_key = nibbles_key.as_slice();
-        self.walk_path.clear();
-
-        let mut current_node = 0;
-        let mut path_walked = 0;
-
-        loop {
-            let node = self.nodes.get(current_node)?;
-
-            match node {
-                DiffTrieNode::Branch { children } => {
-                    // deleting from branch, key not found
-                    if get_key.len() == path_walked {
-                        return None;
-                    }
-
-                    let children = *children;
-
-                    let n = get_key[path_walked];
-                    self.walk_path.push((current_node, n));
-                    path_walked += 1;
-
-                    if let Some(child_ptr) = self.branch_node_children[children][n as usize] {
-                        current_node = child_ptr.local_ptr()?;
-                        continue;
-                    } else {
-                        return None;
-                    }
-                }
-                DiffTrieNode::Extension { key, next_node } => {
-                    let key = key.clone();
-                    let next_node = *next_node;
-
-                    if get_key[path_walked..].starts_with(&self.keys[key.clone()]) {
-                        self.walk_path.push((current_node, 0));
-                        path_walked += key.len();
-                        current_node = next_node.local_ptr()?;
-                        continue;
-                    }
-
-                    return None;
-                }
-                DiffTrieNode::Leaf { key, value } => {
-                    if self.keys[key.clone()] == get_key[path_walked..] {
-                        self.walk_path.push((current_node, 0));
-                        return Some(b"test");
-                    }
-                    return None;
-                }
-                DiffTrieNode::Null => return None,
-            }
-        }
-    }
-
     pub fn delete_nibbles_key(
         &mut self,
         nibbles_key: &Nibbles,