feat: compute world state

Cargo.toml

@@ -32,3 +32,5 @@ tokio = { version = "1.38.0", features = ["full"] }
 thiserror = "1.0.61"
 hex = "0.4.3"
 lazy_static = "1.5.0"
+patricia-merkle-tree = { git = "https://github.com/lambdaclass/merkle_patricia_tree.git" }
+sha3 = "0.10.8"

crates/core/Cargo.toml

@@ -12,13 +12,13 @@ serde.workspace = true
 serde_json.workspace = true
 thiserror.workspace = true
 keccak-hash = "0.10.0"
-sha3 = "0.10.8"
+sha3.workspace = true
 secp256k1 = { version = "0.29", default-features = false, features = [
     "global-context",
     "recovery",
 ] }
 revm = { version = "10.0.0", features = ["serde", "std", "serde-json"] }
-patricia-merkle-tree = { git = "https://github.com/lambdaclass/merkle_patricia_tree.git" }
+patricia-merkle-tree.workspace = true
 bytes.workspace = true
 hex.workspace = true
 lazy_static.workspace = true

crates/core/src/rlp/encode.rs

@@ -13,6 +13,12 @@ pub trait RLPEncode {
         self.encode(&mut buf);
         buf.len()
     }
+
+    fn encode_to_vec(&self) -> Vec<u8> {
+        let mut buf = Vec::new();
+        self.encode(&mut buf);
+        buf
+    }
 }
 
 impl RLPEncode for bool {

crates/core/src/types/account.rs

@@ -2,8 +2,15 @@ use std::collections::HashMap;
 
 use bytes::Bytes;
 use ethereum_types::{H256, U256};
+use patricia_merkle_tree::PatriciaMerkleTree;
+use sha3::Keccak256;
 
-use crate::rlp::{encode::RLPEncode, structs::Encoder};
+use crate::rlp::{
+    decode::RLPDecode,
+    encode::RLPEncode,
+    error::RLPDecodeError,
+    structs::{Decoder, Encoder},
+};
 
 use super::GenesisAccount;
 
@@ -22,6 +29,14 @@ pub struct AccountInfo {
     pub nonce: u64,
 }
 
+#[derive(Debug, PartialEq)]
+pub struct AccountState {
+    pub nonce: u64,
+    pub balance: U256,
+    pub storage_root: H256,
+    pub code_hash: H256,
+}
+
 impl From<GenesisAccount> for Account {
     fn from(genesis: GenesisAccount) -> Self {
         Self {
@@ -50,6 +65,80 @@ impl RLPEncode for AccountInfo {
     }
 }
 
+impl RLPDecode for AccountInfo {
+    fn decode_unfinished(rlp: &[u8]) -> Result<(AccountInfo, &[u8]), RLPDecodeError> {
+        let decoder = Decoder::new(rlp)?;
+        let (code_hash, decoder) = decoder.decode_field("code_hash")?;
+        let (balance, decoder) = decoder.decode_field("balance")?;
+        let (nonce, decoder) = decoder.decode_field("nonce")?;
+
+        let account_info = AccountInfo {
+            code_hash,
+            balance,
+            nonce,
+        };
+        Ok((account_info, decoder.finish()?))
+    }
+}
+
+impl RLPEncode for AccountState {
+    fn encode(&self, buf: &mut dyn bytes::BufMut) {
+        Encoder::new(buf)
+            .encode_field(&self.nonce)
+            .encode_field(&self.balance)
+            .encode_field(&self.storage_root)
+            .encode_field(&self.code_hash)
+            .finish();
+    }
+}
+
+impl RLPDecode for AccountState {
+    fn decode_unfinished(rlp: &[u8]) -> Result<(AccountState, &[u8]), RLPDecodeError> {
+        let decoder = Decoder::new(rlp)?;
+        let (nonce, decoder) = decoder.decode_field("nonce")?;
+        let (balance, decoder) = decoder.decode_field("balance")?;
+        let (storage_root, decoder) = decoder.decode_field("storage_root")?;
+        let (code_hash, decoder) = decoder.decode_field("code_hash")?;
+        let state = AccountState {
+            nonce,
+            balance,
+            storage_root,
+            code_hash,
+        };
+        Ok((state, decoder.finish()?))
+    }
+}
+
+pub fn compute_storage_root(storage: &HashMap<H256, H256>) -> H256 {
+    let rlp_storage = storage
+        .iter()
+        .map(|(k, v)| {
+            let mut k_buf = vec![];
+            let mut v_buf = vec![];
+            k.encode(&mut k_buf);
+            v.encode(&mut v_buf);
+            (k_buf, v_buf)
+        })
+        .collect::<Vec<_>>();
+    let root =
+        PatriciaMerkleTree::<_, _, Keccak256>::compute_hash_from_sorted_iter(rlp_storage.iter());
+    H256(root.into())
+}
+
+impl AccountState {
+    pub fn from_info_and_storage(
+        info: &AccountInfo,
+        storage: &HashMap<H256, H256>,
+    ) -> AccountState {
+        AccountState {
+            nonce: info.nonce,
+            balance: info.balance,
+            storage_root: compute_storage_root(storage),
+            code_hash: info.code_hash,
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::str::FromStr;

crates/storage/Cargo.toml

@@ -10,3 +10,8 @@ ethereum_rust-core.workspace = true
 
 libmdbx.workspace = true
 anyhow = "1.0.86"
+patricia-merkle-tree.workspace = true
+sha3.workspace = true
+
+[dev-dependencies]
+hex.workspace = true

crates/storage/src/account.rs

@@ -1,12 +1,12 @@
 use libmdbx::orm::{Decodable, Encodable};
 
-pub struct AddressRLP(Vec<u8>);
+pub struct AddressRLP(pub Vec<u8>);
 
-pub struct AccountInfoRLP(Vec<u8>);
+pub struct AccountInfoRLP(pub Vec<u8>);
 
-pub struct AccountStorageKeyRLP(Vec<u8>);
+pub struct AccountStorageKeyRLP(pub [u8; 32]);
 
-pub struct AccountStorageValueRLP(Vec<u8>);
+pub struct AccountStorageValueRLP(pub [u8; 32]);
 
 pub struct AccountCodeHashRLP(Vec<u8>);
 
@@ -41,7 +41,7 @@ impl Decodable for AccountInfoRLP {
 }
 
 impl Encodable for AccountStorageKeyRLP {
-    type Encoded = Vec<u8>;
+    type Encoded = [u8; 32];
 
     fn encode(self) -> Self::Encoded {
         self.0
@@ -50,12 +50,12 @@ impl Encodable for AccountStorageKeyRLP {
 
 impl Decodable for AccountStorageKeyRLP {
     fn decode(b: &[u8]) -> anyhow::Result<Self> {
-        Ok(AccountStorageKeyRLP(b.to_vec()))
+        Ok(AccountStorageKeyRLP(b.try_into()?))
     }
 }
 
 impl Encodable for AccountStorageValueRLP {
-    type Encoded = Vec<u8>;
+    type Encoded = [u8; 32];
 
     fn encode(self) -> Self::Encoded {
         self.0
@@ -64,7 +64,7 @@ impl Encodable for AccountStorageValueRLP {
 
 impl Decodable for AccountStorageValueRLP {
     fn decode(b: &[u8]) -> anyhow::Result<Self> {
-        Ok(AccountStorageValueRLP(b.to_vec()))
+        Ok(AccountStorageValueRLP(b.try_into()?))
     }
 }
 

crates/storage/src/lib.rs

@@ -1,6 +1,7 @@
 mod account;
 mod block;
 mod receipt;
+mod world_state;
 
 use account::{
     AccountCodeHashRLP, AccountCodeRLP, AccountInfoRLP, AccountStorageKeyRLP,
@@ -31,7 +32,7 @@ table!(
 );
 dupsort!(
     /// Account storages table.
-    ( AccountStorages ) AddressRLP[AccountStorageKeyRLP] => AccountStorageValueRLP
+    ( AccountStorages ) AddressRLP => (AccountStorageKeyRLP, AccountStorageValueRLP) [AccountStorageKeyRLP]
 );
 table!(
     /// Account codes table.

crates/storage/src/world_state.rs

@@ -0,0 +1,160 @@
+use std::collections::HashMap;
+
+use anyhow::anyhow;
+use ethereum_rust_core::{
+    rlp::{decode::RLPDecode, encode::RLPEncode},
+    types::{compute_storage_root, AccountInfo, AccountState},
+    Address, H256,
+};
+use libmdbx::orm::Database;
+use patricia_merkle_tree::PatriciaMerkleTree;
+use sha3::{digest::core_api::CoreWrapper, Keccak256Core};
+
+use crate::{AccountInfos, AccountStorages};
+
+/// A Merkle Tree containing a mapping from account addresses to account states
+#[allow(unused)]
+#[derive(Default)]
+pub struct WorldState(pub PatriciaMerkleTree<Vec<u8>, Vec<u8>, CoreWrapper<Keccak256Core>>);
+
+#[allow(unused)]
+impl WorldState {
+    pub fn get_account_state(&self, addr: &Address) -> Option<AccountState> {
+        self.0
+            .get(&addr.encode_to_vec())
+            .and_then(|encoded| AccountState::decode(encoded).ok())
+    }
+
+    pub fn build_from_db(db: &Database) -> Result<WorldState, anyhow::Error> {
+        let db = db.begin_read()?;
+        // Fetch & Decode AccountInfos
+        let mut account_infos = HashMap::<_, _>::new();
+        let mut account_infos_db = db.cursor::<AccountInfos>()?;
+        while let Some((rlp_address, rlp_info)) = account_infos_db.next()? {
+            account_infos.insert(
+                Address::decode(&rlp_address.0)?,
+                AccountInfo::decode(&rlp_info.0)?,
+            );
+        }
+        // Fetch & Decode Account Storages
+        let mut account_storages = HashMap::<Address, HashMap<H256, H256>>::new();
+        let mut account_storages_db = db.cursor::<AccountStorages>()?;
+        while let Some((rlp_address, (storage_key_bytes, storage_value_bytes))) =
+            account_storages_db.next()?
+        {
+            let entry = account_storages
+                .entry(Address::decode(&rlp_address.0)?)
+                .or_default();
+            entry.insert(
+                H256::from(&storage_key_bytes.0),
+                H256::from(&storage_value_bytes.0),
+            );
+        }
+        // Fill World State merkle tree
+        let mut world_state = WorldState::default();
+        for (addr, account_info) in account_infos {
+            let storage = account_storages
+                .get(&addr)
+                .ok_or(anyhow!("No storage found in db for account address {addr}"))?;
+            let account_state = AccountState {
+                nonce: account_info.nonce,
+                balance: account_info.balance,
+                storage_root: compute_storage_root(storage),
+                code_hash: account_info.code_hash,
+            };
+
+            world_state
+                .0
+                .insert(addr.encode_to_vec(), account_state.encode_to_vec());
+        }
+
+        Ok(world_state)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use ethereum_rust_core::rlp::encode::RLPEncode;
+    use std::str::FromStr;
+
+    use crate::{
+        account::{AccountInfoRLP, AccountStorageKeyRLP, AccountStorageValueRLP, AddressRLP},
+        init_db,
+    };
+
+    use super::*;
+
+    #[test]
+    fn test_build_world_state() {
+        let db = init_db(None::<String>);
+        let write = db.begin_readwrite().unwrap();
+        let account_info = AccountInfo {
+            code_hash: H256::from_slice(
+                &hex::decode("3d9209c0aa535c1b05fb6000abf2e8239ac31d12ff08b2c6db0c6ef68cf7795f")
+                    .unwrap(),
+            ),
+            balance: 12.into(),
+            nonce: 0,
+        };
+        let address =
+            Address::from_slice(&hex::decode("a94f5374fce5edbc8e2a8697c15331677e6ebf0b").unwrap());
+
+        let storage = [
+            (
+                H256::from_str(
+                    "0x1000000000000000000000000000000000000000000000000000000000000022",
+                )
+                .unwrap(),
+                H256::from_str(
+                    "0xf5a5fd42d16a20302798ef6ed309979b43003d2320d9f0e8ea9831a92759fb4b",
+                )
+                .unwrap(),
+            ),
+            (
+                H256::from_str(
+                    "0x1000000000000000000000000000000000000000000000000000000000000038",
+                )
+                .unwrap(),
+                H256::from_str(
+                    "0xe71f0aa83cc32edfbefa9f4d3e0174ca85182eec9f3a09f6a6c0df6377a510d7",
+                )
+                .unwrap(),
+            ),
+        ];
+        write
+            .upsert::<AccountInfos>(
+                AddressRLP(address.encode_to_vec()),
+                AccountInfoRLP(account_info.encode_to_vec()),
+            )
+            .unwrap();
+        write
+            .upsert::<AccountStorages>(
+                AddressRLP(address.encode_to_vec()),
+                (
+                    AccountStorageKeyRLP(storage[0].0 .0),
+                    AccountStorageValueRLP(storage[0].1 .0),
+                ),
+            )
+            .unwrap();
+        write
+            .upsert::<AccountStorages>(
+                AddressRLP(address.encode_to_vec()),
+                (
+                    AccountStorageKeyRLP(storage[1].0 .0),
+                    AccountStorageValueRLP(storage[1].1 .0),
+                ),
+            )
+            .unwrap();
+        write.commit().unwrap();
+
+        let world_state = WorldState::build_from_db(&db).unwrap();
+        let account_state = world_state.get_account_state(&address).unwrap();
+        let expected_account_state = AccountState {
+            nonce: account_info.nonce,
+            balance: account_info.balance,
+            storage_root: compute_storage_root(&HashMap::from(storage)),
+            code_hash: account_state.code_hash,
+        };
+        assert_eq!(account_state, expected_account_state);
+    }
+}