feat: implement StateTrieFetcher

portal-bridge/src/bin/fetch_state_from_portal.rs

@@ -0,0 +1,77 @@
+use anyhow::bail;
+use clap::Parser;
+use portal_bridge::{
+    beacon_block_fetcher::BeaconBlockFetcher, state_trie_fetcher::StateTrieFetcher,
+    types::genesis::GenesisConfig,
+};
+
+const LOCALHOST_BEACON_RPC_URL: &str = "http://localhost:9596/";
+const LOCALHOST_PORTAL_RPC_URL: &str = "http://localhost:8545/";
+
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+pub struct Args {
+    #[arg(long, num_args=1..)]
+    pub slots: Vec<u64>,
+    #[arg(long, default_value_t = String::from(LOCALHOST_BEACON_RPC_URL))]
+    pub beacon_rpc_url: String,
+    #[arg(long, default_value_t = String::from(LOCALHOST_PORTAL_RPC_URL))]
+    pub portal_rpc_url: String,
+}
+
+struct StateVerifier {
+    block_fetcher: BeaconBlockFetcher,
+    state_trie_fetcher: StateTrieFetcher,
+}
+
+impl StateVerifier {
+    fn new(args: &Args) -> anyhow::Result<Self> {
+        println!("Initializing...");
+        let block_fetcher =
+            BeaconBlockFetcher::new(&args.beacon_rpc_url, /* save_locally = */ false);
+        let state_trie_fetcher = StateTrieFetcher::new(&args.portal_rpc_url)?;
+        Ok(Self {
+            block_fetcher,
+            state_trie_fetcher,
+        })
+    }
+
+    async fn verify_state(&self, slot: u64) -> anyhow::Result<()> {
+        let root = if slot == 0 {
+            GenesisConfig::DEVNET6_STATE_ROOT
+        } else {
+            let Some(beacon_block) = self.block_fetcher.fetch_beacon_block(slot).await? else {
+                bail!("Beacon block for slot {slot} not found!")
+            };
+            beacon_block.message.body.execution_payload.state_root
+        };
+        println!("Veryfing slot {slot} with state root: {root}");
+        match self.state_trie_fetcher.fetch_state_trie(root).await {
+            Ok(synced_state_trie) => {
+                if synced_state_trie.root() == root {
+                    println!("SUCCESS");
+                } else {
+                    println!(
+                        "ERROR: State trie fetched but root is different! Expected {root} but received {}",
+                        synced_state_trie.root()
+                    )
+                }
+            }
+            Err(err) => {
+                println!("ERROR: Error while fetching state trie: {err}")
+            }
+        };
+        Ok(())
+    }
+}
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    let args = Args::parse();
+
+    let verifier = StateVerifier::new(&args)?;
+    for slot in args.slots {
+        verifier.verify_state(slot).await?;
+    }
+    Ok(())
+}

portal-bridge/src/bin/gossip_to_portal.rs

@@ -1,4 +1,7 @@
-use std::{collections::HashMap, time::Duration};
+use std::{
+    collections::{HashMap, HashSet},
+    time::Duration,
+};
 
 use alloy_primitives::B256;
 use clap::Parser;
@@ -50,7 +53,7 @@ impl Gossiper {
         let portal_client = HttpClientBuilder::new()
             .request_timeout(Duration::from_secs(60))
             .build(&args.portal_rpc_url)?;
-        let evm = VerkleEvm::new();
+        let evm = VerkleEvm::new(&read_genesis()?)?;
 
         Ok(Self {
             block_fetcher,
@@ -61,33 +64,44 @@ impl Gossiper {
 
     async fn gossip_genesis(&mut self) -> anyhow::Result<()> {
         let genesis_config = read_genesis()?;
-        let state_writes = self.evm.initialize_genesis(&genesis_config)?;
+        let state_writes = genesis_config.generate_state_diff().into();
         println!("Gossiping genesis...");
-        self.gossip_state_writes(GenesisConfig::DEVNET6_BLOCK_HASH, state_writes)
-            .await?;
+        self.gossip_state_writes(
+            GenesisConfig::DEVNET6_BLOCK_HASH,
+            state_writes,
+            HashSet::new(),
+        )
+        .await?;
         Ok(())
     }
 
     async fn gossip_slot(&mut self, slot: u64) -> anyhow::Result<()> {
-        let Some(beacon_block) = self.block_fetcher.fetch_beacon_block(slot).await? else {
+        let Ok(Some(beacon_block)) = self.block_fetcher.fetch_beacon_block(slot).await else {
             println!("Beacon block for slot {slot} not found!");
             return Ok(());
         };
         let execution_payload = &beacon_block.message.body.execution_payload;
-        let state_writes = self.evm.process_block(execution_payload)?;
+        let process_block_result = self.evm.process_block(execution_payload)?;
         println!(
-            "Gossiping slot {slot:04} (block - number={:04} hash={})",
-            execution_payload.block_number, execution_payload.block_hash
+            "Gossiping slot {slot:04} (block - number={:04} hash={} root={})",
+            execution_payload.block_number,
+            execution_payload.block_hash,
+            execution_payload.state_root
         );
-        self.gossip_state_writes(execution_payload.block_hash, state_writes)
-            .await?;
+        self.gossip_state_writes(
+            execution_payload.block_hash,
+            process_block_result.state_writes,
+            process_block_result.new_branch_nodes,
+        )
+        .await?;
         Ok(())
     }
 
     async fn gossip_state_writes(
         &self,
         block_hash: B256,
         state_writes: StateWrites,
+        new_branch_nodes: HashSet<TriePath>,
     ) -> anyhow::Result<()> {
         let mut content_to_gossip = HashMap::<VerkleContentKey, VerkleContentValue>::new();
         for stem_state_write in state_writes.iter() {
@@ -98,6 +112,13 @@ impl Gossiper {
 
                 // Branch bundle
                 let content_key = VerkleContentKey::Bundle(branch.commitment().clone());
+                if new_branch_nodes.contains(&trie_path)
+                    && content_to_gossip.contains_key(&content_key)
+                {
+                    // We already gossiped this bundle and all fragments
+                    continue;
+                }
+
                 content_to_gossip.entry(content_key).or_insert_with(|| {
                     let trie_proof = Self::create_bundle_trie_proof(
                         stem,
@@ -113,25 +134,36 @@ impl Gossiper {
                     })
                 });
 
+                let fragment_indices = if new_branch_nodes.contains(&trie_path) {
+                    0..PORTAL_NETWORK_NODE_WIDTH
+                } else {
+                    let fragment_index = stem[depth] as usize / PORTAL_NETWORK_NODE_WIDTH;
+                    fragment_index..fragment_index + 1
+                };
                 // Branch fragment
-                let (fragment_commitment, fragment) =
-                    branch.extract_fragment_node(stem[depth] as usize / PORTAL_NETWORK_NODE_WIDTH);
-                let content_key = VerkleContentKey::BranchFragment(fragment_commitment);
-                content_to_gossip.entry(content_key).or_insert_with(|| {
-                    let trie_proof = Self::create_branch_fragment_trie_proof(
-                        stem,
-                        depth,
-                        &path_to_leaf,
-                        branch.commitment(),
-                        &fragment,
-                    );
-                    VerkleContentValue::BranchFragmentWithProof(BranchFragmentNodeWithProof {
-                        node: fragment,
-                        block_hash,
-                        path: trie_path,
-                        proof: trie_proof,
-                    })
-                });
+                for fragment_index in fragment_indices {
+                    let (fragment_commitment, fragment) =
+                        branch.extract_fragment_node(fragment_index);
+                    if fragment_commitment.is_zero() {
+                        continue;
+                    }
+                    let content_key = VerkleContentKey::BranchFragment(fragment_commitment);
+                    content_to_gossip.entry(content_key).or_insert_with(|| {
+                        let trie_proof = Self::create_branch_fragment_trie_proof(
+                            stem,
+                            depth,
+                            &path_to_leaf,
+                            branch.commitment(),
+                            &fragment,
+                        );
+                        VerkleContentValue::BranchFragmentWithProof(BranchFragmentNodeWithProof {
+                            node: fragment,
+                            block_hash,
+                            path: trie_path.clone(),
+                            proof: trie_proof,
+                        })
+                    });
+                }
             }
 
             // Leaf bundle
@@ -245,7 +277,10 @@ impl Gossiper {
 async fn main() -> anyhow::Result<()> {
     let args = Args::parse();
 
+    println!("Initializing...");
     let mut gossiper = Gossiper::new(&args)?;
+
+    println!("Starting gossiping");
     gossiper.gossip_genesis().await?;
     for slot in 1..=args.slots {
         gossiper.gossip_slot(slot).await?;

portal-bridge/src/evm/verkle_evm.rs

@@ -1,4 +1,7 @@
+use std::collections::HashSet;
+
 use alloy_primitives::{address, keccak256, U8};
+use portal_verkle_trie::nodes::portal::ssz::TriePath;
 use verkle_core::{
     constants::{BALANCE_LEAF_KEY, CODE_KECCAK_LEAF_KEY, NONCE_LEAF_KEY, VERSION_LEAF_KEY},
     storage::AccountStorageLayout,
@@ -16,75 +19,71 @@ use crate::{
     verkle_trie::VerkleTrie,
 };
 
-#[derive(Default)]
 pub struct VerkleEvm {
-    next_block: u64,
-    state: VerkleTrie,
+    block: u64,
+    state_trie: VerkleTrie,
+}
+
+pub struct ProcessBlockResult {
+    pub state_writes: StateWrites,
+    pub new_branch_nodes: HashSet<TriePath>,
 }
 
 impl VerkleEvm {
-    pub fn new() -> Self {
-        Self::default()
+    pub fn new(genesis_config: &GenesisConfig) -> Result<Self, EvmError> {
+        let mut state_trie = VerkleTrie::new();
+        state_trie
+            .update(&genesis_config.generate_state_diff().into())
+            .map_err(EvmError::TrieError)?;
+        Ok(Self {
+            block: 0,
+            state_trie,
+        })
     }
 
     pub fn state_trie(&self) -> &VerkleTrie {
-        &self.state
+        &self.state_trie
     }
 
-    pub fn next_block(&self) -> u64 {
-        self.next_block
-    }
-
-    pub fn initialize_genesis(
-        &mut self,
-        genesis_config: &GenesisConfig,
-    ) -> Result<StateWrites, EvmError> {
-        if self.next_block != 0 {
-            return Err(EvmError::UnexpectedBlock {
-                expected: self.next_block,
-                actual: 0,
-            });
-        }
-        let state_writes = genesis_config.generate_state_diff().into();
-        self.state
-            .update(&state_writes)
-            .map_err(EvmError::TrieError)?;
-        self.next_block += 1;
-
-        Ok(state_writes)
+    pub fn block(&self) -> u64 {
+        self.block
     }
 
     pub fn process_block(
         &mut self,
         execution_payload: &ExecutionPayload,
-    ) -> Result<StateWrites, EvmError> {
-        if self.next_block != execution_payload.block_number.to::<u64>() {
+    ) -> Result<ProcessBlockResult, EvmError> {
+        if self.block + 1 != execution_payload.block_number.to::<u64>() {
             return Err(EvmError::UnexpectedBlock {
-                expected: self.next_block,
+                expected: self.block + 1,
                 actual: execution_payload.block_number.to(),
             });
         }
 
         let mut state_diff = execution_payload.execution_witness.state_diff.clone();
 
-        if self.next_block == 1 {
+        if self.block == 0 {
             update_state_diff_for_eip2935(&mut state_diff);
         }
 
         let state_writes = StateWrites::from(state_diff);
 
-        self.state
+        let new_branch_nodes = self
+            .state_trie
             .update(&state_writes)
             .map_err(EvmError::TrieError)?;
-        self.next_block += 1;
+        self.block += 1;
 
-        if self.state.root() != execution_payload.state_root {
+        if self.state_trie.root() != execution_payload.state_root {
             return Err(EvmError::WrongStateRoot {
                 expected: execution_payload.state_root,
-                actual: self.state.root(),
+                actual: self.state_trie.root(),
             });
         }
-        Ok(state_writes)
+        Ok(ProcessBlockResult {
+            state_writes,
+            new_branch_nodes,
+        })
     }
 }
 
@@ -134,30 +133,27 @@ mod tests {
         const STATE_ROOT: B256 =
             b256!("1fbf85345a3cbba9a6d44f991b721e55620a22397c2a93ee8d5011136ac300ee");
 
-        let mut evm = VerkleEvm::new();
-        evm.initialize_genesis(&read_genesis_for_test()?)?;
+        let evm = VerkleEvm::new(&read_genesis_for_test()?)?;
 
-        assert_eq!(evm.state.root(), STATE_ROOT);
+        assert_eq!(evm.state_trie.root(), STATE_ROOT);
         Ok(())
     }
 
     #[test]
     fn process_block_1() -> Result<()> {
-        let mut evm = VerkleEvm::new();
-        evm.initialize_genesis(&read_genesis_for_test()?)?;
+        let mut evm = VerkleEvm::new(&read_genesis_for_test()?)?;
 
         let reader = BufReader::new(File::open(test_path(beacon_slot_path(1)))?);
         let response: SuccessMessage = serde_json::from_reader(reader)?;
         let execution_payload = response.data.message.body.execution_payload;
         evm.process_block(&execution_payload)?;
-        assert_eq!(evm.state.root(), execution_payload.state_root);
+        assert_eq!(evm.state_trie.root(), execution_payload.state_root);
         Ok(())
     }
 
     #[test]
     fn process_block_1000() -> Result<()> {
-        let mut evm = VerkleEvm::new();
-        evm.initialize_genesis(&read_genesis_for_test()?)?;
+        let mut evm = VerkleEvm::new(&read_genesis_for_test()?)?;
 
         for block in 1..=1000 {
             let path = test_path(beacon_slot_path(block));

portal-bridge/src/lib.rs

@@ -1,6 +1,7 @@
 pub mod beacon_block_fetcher;
 pub mod evm;
 pub mod paths;
+pub mod state_trie_fetcher;
 pub mod types;
 pub mod utils;
 pub mod verkle_trie;