feat(katana): state trie (#2607)

Cargo.toml

@@ -35,6 +35,7 @@ members = [
 	"crates/katana/storage/db",
 	"crates/katana/storage/provider",
 	"crates/katana/tasks",
+	"crates/katana/trie",
 	"crates/metrics",
 	"crates/saya/core",
 	"crates/saya/provider",
@@ -100,6 +101,7 @@ katana-rpc-types-builder = { path = "crates/katana/rpc/rpc-types-builder" }
 katana-runner = { path = "crates/katana/runner" }
 katana-slot-controller = { path = "crates/katana/controller" }
 katana-tasks = { path = "crates/katana/tasks" }
+katana-trie = { path = "crates/katana/trie" }
 
 # torii
 torii-client = { path = "crates/torii/client" }
@@ -254,4 +256,6 @@ alloy-transport = { version = "0.3", default-features = false }
 
 starknet = "0.12.0"
 starknet-crypto = "0.7.1"
-starknet-types-core = { version = "0.1.6", features = [ "arbitrary" ] }
+starknet-types-core = { version = "0.1.7", features = [ "arbitrary", "hash" ] }
+
+bitvec = "1.0.1"

crates/katana/storage/db/Cargo.toml

@@ -8,6 +8,7 @@ version.workspace = true
 
 [dependencies]
 katana-primitives = { workspace = true, features = [ "arbitrary" ] }
+katana-trie.workspace = true
 
 anyhow.workspace = true
 dojo-metrics.workspace = true
@@ -17,12 +18,16 @@ parking_lot.workspace = true
 roaring = { version = "0.10.3", features = [ "serde" ] }
 serde.workspace = true
 serde_json.workspace = true
+starknet.workspace = true
+starknet-types-core.workspace = true
 tempfile.workspace = true
 thiserror.workspace = true
 tracing.workspace = true
 
 # codecs
+bitvec.workspace = true
 postcard = { workspace = true, optional = true }
+smallvec = "1.13.2"
 
 [dependencies.libmdbx]
 git = "https://github.com/paradigmxyz/reth.git"
@@ -32,7 +37,6 @@ rev = "b34b0d3"
 [dev-dependencies]
 arbitrary.workspace = true
 criterion.workspace = true
-starknet.workspace = true
 
 [features]
 default = [ "postcard" ]

crates/katana/storage/db/src/codecs/postcard.rs

@@ -11,6 +11,7 @@ use crate::error::CodecError;
 use crate::models::block::StoredBlockBodyIndices;
 use crate::models::contract::ContractInfoChangeList;
 use crate::models::list::BlockList;
+use crate::models::trie::TrieDatabaseValue;
 
 macro_rules! impl_compress_and_decompress_for_table_values {
     ($($name:ty),*) => {
@@ -38,6 +39,7 @@ impl_compress_and_decompress_for_table_values!(
     Header,
     Receipt,
     Felt,
+    TrieDatabaseValue,
     ContractAddress,
     BlockList,
     GenericContractInfo,

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

@@ -13,6 +13,7 @@ pub mod error;
 pub mod mdbx;
 pub mod models;
 pub mod tables;
+pub mod trie;
 pub mod utils;
 pub mod version;
 

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

@@ -3,3 +3,4 @@ pub mod class;
 pub mod contract;
 pub mod list;
 pub mod storage;
+pub mod trie;

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

@@ -0,0 +1,81 @@
+use katana_trie::bonsai::ByteVec;
+use serde::{Deserialize, Serialize};
+
+use crate::codecs::{Decode, Encode};
+use crate::error::CodecError;
+
+#[repr(u8)]
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
+pub enum TrieDatabaseKeyType {
+    Trie = 0,
+    Flat,
+    TrieLog,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct TrieDatabaseKey {
+    pub r#type: TrieDatabaseKeyType,
+    pub key: Vec<u8>,
+}
+
+pub type TrieDatabaseValue = ByteVec;
+
+impl Encode for TrieDatabaseKey {
+    type Encoded = Vec<u8>;
+
+    fn encode(self) -> Self::Encoded {
+        let mut encoded = Vec::new();
+        encoded.push(self.r#type as u8);
+        encoded.extend(self.key);
+        encoded
+    }
+}
+
+impl Decode for TrieDatabaseKey {
+    fn decode<B: AsRef<[u8]>>(bytes: B) -> Result<Self, CodecError> {
+        let bytes = bytes.as_ref();
+        if bytes.is_empty() {
+            panic!("emptyy buffer")
+        }
+
+        let r#type = match bytes[0] {
+            0 => TrieDatabaseKeyType::Trie,
+            1 => TrieDatabaseKeyType::Flat,
+            2 => TrieDatabaseKeyType::TrieLog,
+            _ => panic!("Invalid trie database key type"),
+        };
+
+        let key = bytes[1..].to_vec();
+
+        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);
+    }
+}