From 070711cc9493a4bcb1dce822314a2a8de602a860 Mon Sep 17 00:00:00 2001
From: Ammar Arif <evergreenkary@gmail.com>
Date: Thu, 31 Oct 2024 15:31:34 -0400
Subject: [PATCH] skip the first 4 bits

---
 crates/katana/storage/db/src/models/trie.rs   | 35 +++++++++++++++++--
 crates/katana/storage/db/src/trie.rs          |  4 +--
 .../storage/provider/src/providers/db/trie.rs |  6 ++--
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/crates/katana/storage/db/src/models/trie.rs b/crates/katana/storage/db/src/models/trie.rs
index 6d573048e9..b10456af57 100644
--- a/crates/katana/storage/db/src/models/trie.rs
+++ b/crates/katana/storage/db/src/models/trie.rs
@@ -5,14 +5,14 @@ use crate::codecs::{Decode, Encode};
 use crate::error::CodecError;
 
 #[repr(u8)]
-#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
 pub enum TrieDatabaseKeyType {
-    Trie,
+    Trie = 0,
     Flat,
     TrieLog,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct TrieDatabaseKey {
     pub r#type: TrieDatabaseKeyType,
     pub key: Vec<u8>,
@@ -50,3 +50,32 @@ impl Decode for TrieDatabaseKey {
         Ok(TrieDatabaseKey { r#type, key })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_trie_key_roundtrip() {
+        let key = TrieDatabaseKey { r#type: TrieDatabaseKeyType::Trie, key: vec![1, 2, 3] };
+        let encoded = key.clone().encode();
+        let decoded = TrieDatabaseKey::decode(encoded).unwrap();
+        assert_eq!(key, decoded);
+    }
+
+    #[test]
+    fn test_flat_key_roundtrip() {
+        let key = TrieDatabaseKey { r#type: TrieDatabaseKeyType::Flat, key: vec![4, 5, 6] };
+        let encoded = key.clone().encode();
+        let decoded = TrieDatabaseKey::decode(encoded).unwrap();
+        assert_eq!(key, decoded);
+    }
+
+    #[test]
+    fn test_trielog_key_roundtrip() {
+        let key = TrieDatabaseKey { r#type: TrieDatabaseKeyType::TrieLog, key: vec![7, 8, 9] };
+        let encoded = key.clone().encode();
+        let decoded = TrieDatabaseKey::decode(encoded).unwrap();
+        assert_eq!(key, decoded);
+    }
+}
diff --git a/crates/katana/storage/db/src/trie.rs b/crates/katana/storage/db/src/trie.rs
index 4045418397..7edefbe521 100644
--- a/crates/katana/storage/db/src/trie.rs
+++ b/crates/katana/storage/db/src/trie.rs
@@ -122,11 +122,11 @@ where
     fn snapshot(&mut self, _: BasicId) {}
 
     fn merge(&mut self, _: Self::Transaction) -> Result<(), Self::DatabaseError> {
-        Ok(())
+        todo!();
     }
 
     fn transaction(&self, _: BasicId) -> Option<Self::Transaction> {
-        None
+        todo!();
     }
 }
 
diff --git a/crates/katana/storage/provider/src/providers/db/trie.rs b/crates/katana/storage/provider/src/providers/db/trie.rs
index de4f91418a..0a9e4e29e7 100644
--- a/crates/katana/storage/provider/src/providers/db/trie.rs
+++ b/crates/katana/storage/provider/src/providers/db/trie.rs
@@ -62,9 +62,9 @@ impl<Db: Database> ContractTrieWriter for DbProvider<Db> {
             // First we insert the contract storage changes
             for (address, storage_entries) in &state_updates.storage_updates {
                 for (key, value) in storage_entries {
-                    let keys = key.to_bytes_be();
-                    let keys: BitVec<u8, Msb0> = keys.as_bits().to_owned();
-                    storage_trie_db.insert(&address.to_bytes_be(), &keys, value).unwrap();
+                    let bytes = key.to_bytes_be();
+                    let bv: BitVec<u8, Msb0> = bytes.as_bits()[5..].to_owned();
+                    storage_trie_db.insert(&address.to_bytes_be(), &bv, value).unwrap();
                 }
                 // insert the contract address in the contract_leafs to put the storage root later
                 contract_leafs.insert(*address, Default::default());