From 8eccb3657dc1b7a1e0e91d8a726d0ce40f3ddb78 Mon Sep 17 00:00:00 2001
From: Thomas Coratger <thomas.coratger@gmail.com>
Date: Thu, 21 Nov 2024 21:19:51 +0100
Subject: [PATCH] trie: add unit tests for nibbles

---
 Cargo.lock                        |  1 +
 crates/trie/common/Cargo.toml     | 27 ++++-----
 crates/trie/common/src/nibbles.rs | 94 +++++++++++++++++++++++++++++++
 3 files changed, 109 insertions(+), 13 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index db7adde04793..a11d596660e8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -9370,6 +9370,7 @@ dependencies = [
  "reth-primitives-traits",
  "revm-primitives",
  "serde",
+ "serde_json",
 ]
 
 [[package]]
diff --git a/crates/trie/common/Cargo.toml b/crates/trie/common/Cargo.toml
index 0161fc7ff3d9..5b691d6e2033 100644
--- a/crates/trie/common/Cargo.toml
+++ b/crates/trie/common/Cargo.toml
@@ -40,22 +40,23 @@ proptest.workspace = true
 proptest-arbitrary-interop.workspace = true
 hash-db = "=0.15.2"
 plain_hasher = "0.2"
+serde_json.workspace = true
 
 [features]
 test-utils = [
-	"dep:plain_hasher",
-	"dep:hash-db",
-	"arbitrary",
-	"reth-primitives-traits/test-utils",
-	"reth-codecs/test-utils"
+    "dep:plain_hasher",
+    "dep:hash-db",
+    "arbitrary",
+    "reth-primitives-traits/test-utils",
+    "reth-codecs/test-utils",
 ]
 arbitrary = [
-	"alloy-trie/arbitrary",
-	"dep:arbitrary",
-	"reth-primitives-traits/arbitrary",
-	"alloy-consensus/arbitrary",
-	"alloy-primitives/arbitrary",
-	"nybbles/arbitrary",
-	"revm-primitives/arbitrary",
-	"reth-codecs/arbitrary"
+    "alloy-trie/arbitrary",
+    "dep:arbitrary",
+    "reth-primitives-traits/arbitrary",
+    "alloy-consensus/arbitrary",
+    "alloy-primitives/arbitrary",
+    "nybbles/arbitrary",
+    "revm-primitives/arbitrary",
+    "reth-codecs/arbitrary",
 ]
diff --git a/crates/trie/common/src/nibbles.rs b/crates/trie/common/src/nibbles.rs
index cf94f135f54b..797f194c1300 100644
--- a/crates/trie/common/src/nibbles.rs
+++ b/crates/trie/common/src/nibbles.rs
@@ -120,3 +120,97 @@ impl Compact for StoredNibblesSubKey {
         (Self(Nibbles::from_nibbles_unchecked(&buf[..len])), &buf[65..])
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use bytes::BytesMut;
+
+    #[test]
+    fn test_stored_nibbles_from_nibbles() {
+        let nibbles = Nibbles::from_nibbles_unchecked(vec![0x12, 0x34, 0x56]);
+        let stored = StoredNibbles::from(nibbles.clone());
+        assert_eq!(stored.0, nibbles);
+    }
+
+    #[test]
+    fn test_stored_nibbles_from_vec() {
+        let bytes = vec![0x12, 0x34, 0x56];
+        let stored = StoredNibbles::from(bytes.clone());
+        assert_eq!(stored.0.as_slice(), bytes.as_slice());
+    }
+
+    #[test]
+    fn test_stored_nibbles_equality() {
+        let bytes = vec![0x12, 0x34];
+        let stored = StoredNibbles::from(bytes.clone());
+        assert_eq!(stored, *bytes.as_slice());
+    }
+
+    #[test]
+    fn test_stored_nibbles_partial_cmp() {
+        let stored = StoredNibbles::from(vec![0x12, 0x34]);
+        let other = vec![0x12, 0x35];
+        assert!(stored < *other.as_slice());
+    }
+
+    #[test]
+    fn test_stored_nibbles_to_compact() {
+        let stored = StoredNibbles::from(vec![0x12, 0x34]);
+        let mut buf = BytesMut::with_capacity(10);
+        let len = stored.to_compact(&mut buf);
+        assert_eq!(len, 2);
+        assert_eq!(buf, &vec![0x12, 0x34][..]);
+    }
+
+    #[test]
+    fn test_stored_nibbles_from_compact() {
+        let buf = vec![0x12, 0x34, 0x56];
+        let (stored, remaining) = StoredNibbles::from_compact(&buf, 2);
+        assert_eq!(stored.0.as_slice(), &[0x12, 0x34]);
+        assert_eq!(remaining, &[0x56]);
+    }
+
+    #[test]
+    fn test_stored_nibbles_subkey_from_nibbles() {
+        let nibbles = Nibbles::from_nibbles_unchecked(vec![0x12, 0x34]);
+        let subkey = StoredNibblesSubKey::from(nibbles.clone());
+        assert_eq!(subkey.0, nibbles);
+    }
+
+    #[test]
+    fn test_stored_nibbles_subkey_to_compact() {
+        let subkey = StoredNibblesSubKey::from(vec![0x12, 0x34]);
+        let mut buf = BytesMut::with_capacity(65);
+        let len = subkey.to_compact(&mut buf);
+        assert_eq!(len, 65);
+        assert_eq!(buf[..2], [0x12, 0x34]);
+        assert_eq!(buf[64], 2); // Length byte
+    }
+
+    #[test]
+    fn test_stored_nibbles_subkey_from_compact() {
+        let mut buf = vec![0x12, 0x34];
+        buf.resize(65, 0);
+        buf[64] = 2;
+        let (subkey, remaining) = StoredNibblesSubKey::from_compact(&buf, 65);
+        assert_eq!(subkey.0.as_slice(), &[0x12, 0x34]);
+        assert_eq!(remaining, &[] as &[u8]);
+    }
+
+    #[test]
+    fn test_serialization_stored_nibbles() {
+        let stored = StoredNibbles::from(vec![0x12, 0x34]);
+        let serialized = serde_json::to_string(&stored).unwrap();
+        let deserialized: StoredNibbles = serde_json::from_str(&serialized).unwrap();
+        assert_eq!(stored, deserialized);
+    }
+
+    #[test]
+    fn test_serialization_stored_nibbles_subkey() {
+        let subkey = StoredNibblesSubKey::from(vec![0x12, 0x34]);
+        let serialized = serde_json::to_string(&subkey).unwrap();
+        let deserialized: StoredNibblesSubKey = serde_json::from_str(&serialized).unwrap();
+        assert_eq!(subkey, deserialized);
+    }
+}