fix historical nodes not being encoded correctly

Author: kariy

crates/katana/core/src/backend/mod.rs

@@ -252,15 +252,16 @@ impl<'a, P: TrieWriter> UncommittedBlock<'a, P> {
 
     // state_commitment = hPos("STARKNET_STATE_V0", contract_trie_root, class_trie_root)
     fn compute_new_state_root(&self) -> Felt {
+        println!("ohayo im committing now");
         let class_trie_root = self
             .provider
             .trie_insert_declared_classes(self.header.number, &self.state_updates.declared_classes)
-            .unwrap();
+            .expect("failed to update class trie");
 
         let contract_trie_root = self
             .provider
             .trie_insert_contract_updates(self.header.number, self.state_updates)
-            .unwrap();
+            .expect("failed to update contract trie");
 
         hash::Poseidon::hash_array(&[
             short_string!("STARKNET_STATE_V0"),

crates/katana/rpc/rpc/src/starknet/mod.rs

@@ -1158,8 +1158,8 @@ where
                 let mut contract_leaves_data = Vec::new();
 
                 for address in addresses {
-                    let nonce = state.nonce(address)?.unwrap();
-                    let class_hash = state.class_hash_of_contract(address)?.unwrap();
+                    let nonce = state.nonce(address)?.unwrap_or_default();
+                    let class_hash = state.class_hash_of_contract(address)?.unwrap_or_default();
                     contract_leaves_data.push(ContractLeafData { class_hash, nonce });
                 }
 

crates/katana/storage/db/src/models/trie.rs

@@ -1,18 +1,48 @@
 use katana_trie::bonsai::ByteVec;
 use serde::{Deserialize, Serialize};
 
-use crate::codecs::{Decode, Encode};
+use crate::codecs::{Compress, Decode, Decompress, Encode};
 use crate::error::CodecError;
 
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub struct TrieHistoryEntry {
+    pub key: TrieDatabaseKey,
+    pub value: TrieDatabaseValue,
+}
+
+impl Compress for TrieHistoryEntry {
+    type Compressed = Vec<u8>;
+
+    fn compress(self) -> Self::Compressed {
+        let mut buf = Vec::new();
+        buf.extend(self.key.encode());
+        buf.extend(self.value.compress());
+        buf
+    }
+}
+
+impl Decompress for TrieHistoryEntry {
+    fn decompress<B: AsRef<[u8]>>(bytes: B) -> Result<Self, CodecError> {
+        let bytes = bytes.as_ref();
+
+        let key = TrieDatabaseKey::decode(bytes)?;
+        // first byte is the key type, second byte is the actual key length
+        let key_bytes_length = 1 + 1 + key.key.len();
+        let value = TrieDatabaseValue::decompress(&bytes[key_bytes_length..])?;
+
+        Ok(Self { key, value })
+    }
+}
+
 #[repr(u8)]
-#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
 pub enum TrieDatabaseKeyType {
     Trie = 0,
     Flat,
     TrieLog,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
 pub struct TrieDatabaseKey {
     pub r#type: TrieDatabaseKeyType,
     pub key: Vec<u8>,
@@ -26,6 +56,7 @@ impl Encode for TrieDatabaseKey {
     fn encode(self) -> Self::Encoded {
         let mut encoded = Vec::new();
         encoded.push(self.r#type as u8);
+        encoded.push(self.key.len() as u8); // Encode key length
         encoded.extend(self.key);
         encoded
     }
@@ -34,7 +65,9 @@ impl Encode for TrieDatabaseKey {
 impl Decode for TrieDatabaseKey {
     fn decode<B: AsRef<[u8]>>(bytes: B) -> Result<Self, CodecError> {
         let bytes = bytes.as_ref();
-        if bytes.is_empty() {
+
+        if bytes.len() < 2 {
+            // Need at least type and length bytes
             panic!("emptyy buffer")
         }
 
@@ -45,7 +78,13 @@ impl Decode for TrieDatabaseKey {
             _ => panic!("Invalid trie database key type"),
         };
 
-        let key = bytes[1..].to_vec();
+        let key_len = bytes[1] as usize;
+
+        if bytes.len() < 2 + key_len {
+            panic!("Buffer too short for key length");
+        }
+
+        let key = bytes[2..2 + key_len].to_vec();
 
         Ok(TrieDatabaseKey { r#type, key })
     }

crates/katana/storage/db/src/tables.rs

@@ -11,7 +11,7 @@ use crate::models::contract::{ContractClassChange, ContractInfoChangeList, Contr
 use crate::models::list::BlockList;
 use crate::models::stage::{StageCheckpoint, StageId};
 use crate::models::storage::{ContractStorageEntry, ContractStorageKey, StorageEntry};
-use crate::models::trie::{TrieDatabaseKey, TrieDatabaseValue};
+use crate::models::trie::{TrieDatabaseKey, TrieDatabaseValue, TrieHistoryEntry};
 
 pub trait Key: Encode + Decode + Clone + std::fmt::Debug {}
 pub trait Value: Compress + Decompress + std::fmt::Debug {}
@@ -37,7 +37,12 @@ pub trait DupSort: Table {
     type SubKey: Key;
 }
 
-pub trait Trie: Table<Key = TrieDatabaseKey, Value = TrieDatabaseValue> {}
+pub trait Trie: Table<Key = TrieDatabaseKey, Value = TrieDatabaseValue> {
+    /// Table for storing the trie entries according to the block its was committed.
+    type History: DupSort<Key = BlockNumber, SubKey = TrieDatabaseKey, Value = TrieHistoryEntry>;
+    /// Table for storing the trie change set.
+    type Changeset: Table<Key = TrieDatabaseKey, Value = BlockList>;
+}
 
 /// Enum for the types of tables present in libmdbx.
 #[derive(Debug, PartialEq, Copy, Clone)]
@@ -249,11 +254,11 @@ tables! {
     StoragesTrie: (TrieDatabaseKey) => TrieDatabaseValue,
 
     /// Class trie history
-    ClassesTrieHistory: (BlockNumber, TrieDatabaseKey) => TrieDatabaseValue,
+    ClassesTrieHistory: (BlockNumber, TrieDatabaseKey) => TrieHistoryEntry,
     /// Contract trie history
-    ContractsTrieHistory: (BlockNumber,TrieDatabaseKey) => TrieDatabaseValue,
+    ContractsTrieHistory: (BlockNumber, TrieDatabaseKey) => TrieHistoryEntry,
     /// Contract storage trie history
-    StoragesTrieHistory: (BlockNumber,TrieDatabaseKey) => TrieDatabaseValue,
+    StoragesTrieHistory: (BlockNumber, TrieDatabaseKey) => TrieHistoryEntry,
 
     /// Class trie change set
     ClassesTrieChangeSet: (TrieDatabaseKey) => BlockList,
@@ -263,10 +268,20 @@ tables! {
     StoragesTrieChangeSet: (TrieDatabaseKey) => BlockList
 }
 
-impl Trie for ClassesTrie {}
-impl Trie for ContractsTrie {}
-impl Trie for StoragesTrie {}
+impl Trie for ClassesTrie {
+    type History = ClassesTrieHistory;
+    type Changeset = ClassesTrieChangeSet;
+}
 
+impl Trie for ContractsTrie {
+    type History = ContractsTrieHistory;
+    type Changeset = ContractsTrieChangeSet;
+}
+
+impl Trie for StoragesTrie {
+    type History = StoragesTrieHistory;
+    type Changeset = StoragesTrieChangeSet;
+}
 #[cfg(test)]
 mod tests {