Add fork choice tree visualization

crates/blockchain/fork_choice/src/lib.rs

@@ -2,6 +2,30 @@ use std::collections::HashMap;
 
 use ethlambda_types::{attestation::AttestationData, primitives::H256};
 
+/// Compute per-block attestation weights for the fork choice tree.
+///
+/// For each validator attestation, walks backward from the attestation's head
+/// through the parent chain, incrementing weight for each block above start_slot.
+pub fn compute_block_weights(
+    start_slot: u64,
+    blocks: &HashMap<H256, (u64, H256)>,
+    attestations: &HashMap<u64, AttestationData>,
+) -> HashMap<H256, u64> {
+    let mut weights: HashMap<H256, u64> = HashMap::new();
+
+    for attestation_data in attestations.values() {
+        let mut current_root = attestation_data.head.root;
+        while let Some(&(slot, parent_root)) = blocks.get(&current_root)
+            && slot > start_slot
+        {
+            *weights.entry(current_root).or_default() += 1;
+            current_root = parent_root;
+        }
+    }
+
+    weights
+}
+
 /// Compute the LMD GHOST head of the chain, given a starting root, a set of blocks,
 /// a set of attestations, and a minimum score threshold.
 ///
@@ -24,17 +48,7 @@ pub fn compute_lmd_ghost_head(
             .expect("we already checked blocks is non-empty");
     }
     let start_slot = blocks[&start_root].0;
-    let mut weights: HashMap<H256, u64> = HashMap::new();
-
-    for attestation_data in attestations.values() {
-        let mut current_root = attestation_data.head.root;
-        while let Some(&(slot, parent_root)) = blocks.get(&current_root)
-            && slot > start_slot
-        {
-            *weights.entry(current_root).or_default() += 1;
-            current_root = parent_root;
-        }
-    }
+    let weights = compute_block_weights(start_slot, blocks, attestations);
 
     let mut children_map: HashMap<H256, Vec<H256>> = HashMap::new();
 
@@ -62,3 +76,57 @@ pub fn compute_lmd_ghost_head(
 
     head
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use ethlambda_types::state::Checkpoint;
+
+    fn make_attestation(head_root: H256, slot: u64) -> AttestationData {
+        AttestationData {
+            slot,
+            head: Checkpoint {
+                root: head_root,
+                slot,
+            },
+            target: Checkpoint::default(),
+            source: Checkpoint::default(),
+        }
+    }
+
+    #[test]
+    fn test_compute_block_weights() {
+        // Chain: root_a (slot 0) -> root_b (slot 1) -> root_c (slot 2)
+        let root_a = H256::from([1u8; 32]);
+        let root_b = H256::from([2u8; 32]);
+        let root_c = H256::from([3u8; 32]);
+
+        let mut blocks = HashMap::new();
+        blocks.insert(root_a, (0, H256::ZERO));
+        blocks.insert(root_b, (1, root_a));
+        blocks.insert(root_c, (2, root_b));
+
+        // Two validators: one attests to root_c, one attests to root_b
+        let mut attestations = HashMap::new();
+        attestations.insert(0, make_attestation(root_c, 2));
+        attestations.insert(1, make_attestation(root_b, 1));
+
+        let weights = compute_block_weights(0, &blocks, &attestations);
+
+        // root_c: 1 vote (validator 0)
+        assert_eq!(weights.get(&root_c).copied().unwrap_or(0), 1);
+        // root_b: 2 votes (validator 0 walks through it + validator 1 attests directly)
+        assert_eq!(weights.get(&root_b).copied().unwrap_or(0), 2);
+        // root_a: at slot 0 = start_slot, so not counted
+        assert_eq!(weights.get(&root_a).copied().unwrap_or(0), 0);
+    }
+
+    #[test]
+    fn test_compute_block_weights_empty() {
+        let blocks = HashMap::new();
+        let attestations = HashMap::new();
+
+        let weights = compute_block_weights(0, &blocks, &attestations);
+        assert!(weights.is_empty());
+    }
+}

crates/blockchain/src/store.rs

@@ -34,10 +34,7 @@ fn accept_new_attestations(store: &mut Store) {
 /// Update the head based on the fork choice rule.
 fn update_head(store: &mut Store) {
     let blocks = store.get_live_chain();
-    let attestations = extract_attestations_from_aggregated_payloads(
-        store,
-        store.iter_known_aggregated_payloads(),
-    );
+    let attestations = store.extract_latest_known_attestations();
     let old_head = store.head();
     let new_head = ethlambda_fork_choice::compute_lmd_ghost_head(
         store.latest_justified().root,
@@ -90,8 +87,7 @@ fn update_safe_target(store: &mut Store) {
     for (key, new_proofs) in store.iter_new_aggregated_payloads() {
         all_payloads.entry(key).or_default().extend(new_proofs);
     }
-    let attestations =
-        extract_attestations_from_aggregated_payloads(store, all_payloads.into_iter());
+    let attestations = store.extract_latest_attestations(all_payloads.into_iter());
     let safe_target = ethlambda_fork_choice::compute_lmd_ghost_head(
         store.latest_justified().root,
         &blocks,
@@ -101,34 +97,6 @@ fn update_safe_target(store: &mut Store) {
     store.set_safe_target(safe_target);
 }
 
-/// Reconstruct per-validator attestation data from aggregated payloads.
-///
-/// For each (validator_id, data_root) key in the payloads, looks up the
-/// attestation data by root. Returns the latest attestation per validator
-/// (by slot).
-fn extract_attestations_from_aggregated_payloads(
-    store: &Store,
-    payloads: impl Iterator<Item = (SignatureKey, Vec<StoredAggregatedPayload>)>,
-) -> HashMap<u64, AttestationData> {
-    let mut result: HashMap<u64, AttestationData> = HashMap::new();
-
-    for ((validator_id, data_root), _payload_list) in payloads {
-        let Some(data) = store.get_attestation_data_by_root(&data_root) else {
-            continue;
-        };
-
-        let should_update = result
-            .get(&validator_id)
-            .is_none_or(|existing| existing.slot < data.slot);
-
-        if should_update {
-            result.insert(validator_id, data);
-        }
-    }
-
-    result
-}
-
 /// Aggregate committee signatures at interval 2.
 ///
 /// Collects individual gossip signatures, aggregates them by attestation data,
@@ -765,10 +733,7 @@ pub fn produce_block_with_signatures(
     }
 
     // Convert known aggregated payloads to Attestation objects for build_block
-    let known_attestations = extract_attestations_from_aggregated_payloads(
-        store,
-        store.iter_known_aggregated_payloads(),
-    );
+    let known_attestations = store.extract_latest_known_attestations();
     let available_attestations: Vec<Attestation> = known_attestations
         .into_iter()
         .map(|(validator_id, data)| Attestation { validator_id, data })

crates/net/rpc/Cargo.toml

@@ -12,6 +12,7 @@ version.workspace = true
 [dependencies]
 axum = "0.8.1"
 tokio.workspace = true
+ethlambda-fork-choice.workspace = true
 ethlambda-metrics.workspace = true
 tracing.workspace = true
 ethlambda-storage.workspace = true