feat: add Checkpoint container and match with spec (ReamLabs#701)

Author: syjn99

crates/common/chain/lean/src/lean_chain.rs

@@ -3,8 +3,8 @@ use std::collections::HashMap;
 use alloy_primitives::B256;
 use anyhow::anyhow;
 use ream_consensus_lean::{
-    block::Block, get_fork_choice_head, get_latest_justified_hash, is_justifiable_slot,
-    process_block, state::LeanState, vote::Vote,
+    block::Block, checkpoint::Checkpoint, get_fork_choice_head, get_latest_justified_hash,
+    is_justifiable_slot, process_block, state::LeanState, vote::Vote,
 };
 use ssz_types::VariableList;
 use tree_hash::TreeHash;
@@ -58,7 +58,7 @@ impl LeanChain {
     pub fn latest_finalized_hash(&self) -> Option<B256> {
         self.post_states
             .get(&self.head)
-            .map(|state| state.latest_finalized_hash)
+            .map(|state| state.latest_finalized.root)
     }
 
     /// Compute the latest block that the staker is allowed to choose as the target
@@ -119,7 +119,7 @@ impl LeanChain {
                 .known_votes
                 .clone()
                 .into_iter()
-                .filter(|vote| vote.source == state.latest_justified_hash)
+                .filter(|vote| vote.source.root == state.latest_justified.root)
                 .filter(|vote| !new_block.votes.contains(vote))
                 .collect::<Vec<_>>();
 
@@ -173,7 +173,7 @@ impl LeanChain {
 
         // If the latest finalized slot is very far back, then only some slots are
         // valid to justify, make sure the target is one of those
-        while !is_justifiable_slot(&state.latest_finalized_slot, &target_block.slot) {
+        while !is_justifiable_slot(&state.latest_finalized.slot, &target_block.slot) {
             target_block = self.chain.get(&target_block.parent).ok_or_else(|| {
                 anyhow!(
                     "Block not found for target block's parent hash: {}",
@@ -193,12 +193,18 @@ impl LeanChain {
             // IDs.
             validator_id: 0,
             slot: get_current_slot(),
-            head: self.head,
-            head_slot: head_block.slot,
-            target: target_block.tree_hash_root(),
-            target_slot: target_block.slot,
-            source: state.latest_justified_hash,
-            source_slot: state.latest_justified_slot,
+            head: Checkpoint {
+                root: self.head,
+                slot: head_block.slot,
+            },
+            target: Checkpoint {
+                root: target_block.tree_hash_root(),
+                slot: target_block.slot,
+            },
+            source: Checkpoint {
+                root: state.latest_justified.root,
+                slot: state.latest_justified.slot,
+            },
         })
     }
 

crates/common/chain/lean/src/service.rs

@@ -104,13 +104,13 @@ impl LeanChainService {
                         let vote = &signed_vote.data;
                         info!(
                             "Received signed vote from validator {} for head {:?} / source_slot {:?} at slot {}",
-                            vote.validator_id, vote.head, vote.source_slot, vote.slot
+                            vote.validator_id, vote.head, vote.source.slot, vote.slot
                         );
                     }
                     VoteItem::Unsigned(vote) => {
                         info!(
                             "Received unsigned vote from validator {} for head {:?} / source_slot {:?} at slot {}",
-                            vote.validator_id, vote.head, vote.source_slot, vote.slot
+                            vote.validator_id, vote.head, vote.source.slot, vote.slot
                         );
                     }
                 }
@@ -183,15 +183,15 @@ impl LeanChainService {
 
         if is_known_vote || is_new_vote {
             // Do nothing
-        } else if lean_chain.chain.contains_key(&vote.head) {
+        } else if lean_chain.chain.contains_key(&vote.head.root) {
             drop(lean_chain);
 
             // We should acquire another write lock
             let mut lean_chain = self.lean_chain.write().await;
             lean_chain.new_votes.push(vote);
         } else {
             self.dependencies
-                .entry(vote.head)
+                .entry(vote.head.root)
                 .or_default()
                 .push(QueueItem::VoteItem(VoteItem::Unsigned(vote)));
         }

crates/common/consensus/lean/src/block.rs

@@ -9,7 +9,7 @@ use crate::vote::Vote;
 
 #[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize, Encode, Decode, TreeHash)]
 pub struct SignedBlock {
-    pub message: Block,
+    pub data: Block,
     pub signature: PQSignature,
 }
 

crates/common/consensus/lean/src/checkpoint.rs

@@ -0,0 +1,12 @@
+use alloy_primitives::B256;
+use serde::{Deserialize, Serialize};
+use ssz_derive::{Decode, Encode};
+use tree_hash_derive::TreeHash;
+
+#[derive(
+    Debug, Default, PartialEq, Eq, Clone, Serialize, Deserialize, Encode, Decode, TreeHash,
+)]
+pub struct Checkpoint {
+    pub root: B256,
+    pub slot: u64,
+}

crates/common/consensus/lean/src/lib.rs

@@ -1,4 +1,5 @@
 pub mod block;
+pub mod checkpoint;
 pub mod config;
 pub mod state;
 pub mod vote;
@@ -75,37 +76,37 @@ pub fn process_block(pre_state: &LeanState, block: &Block) -> anyhow::Result<Lea
         // Ignore votes whose source is not already justified,
         // or whose target is not in the history, or whose target is not a
         // valid justifiable slot
-        if !state.justified_slots[vote.source_slot as usize]
-            || vote.source != state.historical_block_hashes[vote.source_slot as usize]
-            || vote.target != state.historical_block_hashes[vote.target_slot as usize]
-            || vote.target_slot <= vote.source_slot
-            || !is_justifiable_slot(&state.latest_finalized_slot, &vote.target_slot)
+        if !state.justified_slots[vote.source.slot as usize]
+            || vote.source.root != state.historical_block_hashes[vote.source.slot as usize]
+            || vote.target.root != state.historical_block_hashes[vote.target.slot as usize]
+            || vote.target.slot <= vote.source.slot
+            || !is_justifiable_slot(&state.latest_finalized.slot, &vote.target.slot)
         {
             continue;
         }
 
         // Track attempts to justify new hashes
-        state.initialize_justifications_for_root(&vote.target)?;
-        state.set_justification(&vote.target, &vote.validator_id, true)?;
+        state.initialize_justifications_for_root(&vote.target.root)?;
+        state.set_justification(&vote.target.root, &vote.validator_id, true)?;
 
-        let count = state.count_justifications(&vote.target)?;
+        let count = state.count_justifications(&vote.target.root)?;
 
         // If 2/3 voted for the same new valid hash to justify
         if count == (2 * state.config.num_validators) / 3 {
-            state.latest_justified_hash = vote.target;
-            state.latest_justified_slot = vote.target_slot;
-            state.justified_slots[vote.target_slot as usize] = true;
+            state.latest_justified.root = vote.target.root;
+            state.latest_justified.slot = vote.target.slot;
+            state.justified_slots[vote.target.slot as usize] = true;
 
-            state.remove_justifications(&vote.target)?;
+            state.remove_justifications(&vote.target.root)?;
 
             // Finalization: if the target is the next valid justifiable
             // hash after the source
-            let is_target_next_valid_justifiable_slot = !((vote.source_slot + 1)..vote.target_slot)
-                .any(|slot| is_justifiable_slot(&state.latest_finalized_slot, &slot));
+            let is_target_next_valid_justifiable_slot = !((vote.source.slot + 1)..vote.target.slot)
+                .any(|slot| is_justifiable_slot(&state.latest_finalized.slot, &slot));
 
             if is_target_next_valid_justifiable_slot {
-                state.latest_finalized_hash = vote.source;
-                state.latest_finalized_slot = vote.source_slot;
+                state.latest_finalized.root = vote.source.root;
+                state.latest_finalized.slot = vote.source.slot;
             }
         }
     }
@@ -117,8 +118,8 @@ pub fn process_block(pre_state: &LeanState, block: &Block) -> anyhow::Result<Lea
 pub fn get_latest_justified_hash(post_states: &HashMap<B256, LeanState>) -> Option<B256> {
     post_states
         .values()
-        .max_by_key(|state| state.latest_justified_slot)
-        .map(|state| state.latest_justified_hash)
+        .max_by_key(|state| state.latest_justified.slot)
+        .map(|state| state.latest_justified.root)
 }
 
 /// Use LMD GHOST to get the head, given a particular root (usually the
@@ -159,8 +160,8 @@ pub fn get_fork_choice_head(
     let mut vote_weights = HashMap::<B256, u64>::new();
 
     for vote in latest_votes.values() {
-        if blocks.contains_key(&vote.head) {
-            let mut block_hash = vote.head;
+        if blocks.contains_key(&vote.head.root) {
+            let mut block_hash = vote.head.root;
             while {
                 let current_block = blocks
                     .get(&block_hash)
@@ -195,21 +196,16 @@ pub fn get_fork_choice_head(
     // choose the child with the most latest votes, tiebreaking by slot then hash
     let mut current_root = root;
 
-    loop {
-        match children_map.get(&current_root) {
-            None => {
-                break Ok(current_root);
-            }
-            Some(children) => {
-                current_root = *children
-                    .iter()
-                    .max_by_key(|child_hash| {
-                        let vote_weight = vote_weights.get(*child_hash).unwrap_or(&0);
-                        let slot = blocks.get(*child_hash).map(|block| block.slot).unwrap_or(0);
-                        (*vote_weight, slot, *(*child_hash))
-                    })
-                    .ok_or_else(|| anyhow!("No children found for current root: {current_root}"))?;
-            }
-        }
+    while let Some(children) = children_map.get(&current_root) {
+        current_root = *children
+            .iter()
+            .max_by_key(|child_hash| {
+                let vote_weight = vote_weights.get(*child_hash).unwrap_or(&0);
+                let slot = blocks.get(*child_hash).map(|block| block.slot).unwrap_or(0);
+                (*vote_weight, slot, *(*child_hash))
+            })
+            .ok_or_else(|| anyhow!("No children found for current root: {current_root}"))?;
     }
+
+    Ok(current_root)
 }

crates/common/consensus/lean/src/state.rs

@@ -9,16 +9,14 @@ use ssz_types::{
 };
 use tree_hash_derive::TreeHash;
 
-use crate::config::Config;
+use crate::{checkpoint::Checkpoint, config::Config};
 
 #[derive(Debug, Eq, PartialEq, Clone, Serialize, Deserialize, Encode, Decode, TreeHash)]
 pub struct LeanState {
     pub config: Config,
 
-    pub latest_justified_hash: B256,
-    pub latest_justified_slot: u64,
-    pub latest_finalized_hash: B256,
-    pub latest_finalized_slot: u64,
+    pub latest_justified: Checkpoint,
+    pub latest_finalized: Checkpoint,
 
     pub historical_block_hashes: VariableList<B256, U262144>,
     pub justified_slots: VariableList<bool, U262144>,
@@ -36,10 +34,8 @@ impl LeanState {
         LeanState {
             config: Config { num_validators },
 
-            latest_justified_hash: B256::ZERO,
-            latest_justified_slot: 0,
-            latest_finalized_hash: B256::ZERO,
-            latest_finalized_slot: 0,
+            latest_justified: Checkpoint::default(),
+            latest_finalized: Checkpoint::default(),
 
             historical_block_hashes: VariableList::empty(),
             justified_slots: VariableList::empty(),

crates/common/consensus/lean/src/vote.rs

@@ -1,9 +1,10 @@
-use alloy_primitives::B256;
 use ream_pqc::PQSignature;
 use serde::{Deserialize, Serialize};
 use ssz_derive::{Decode, Encode};
 use tree_hash_derive::TreeHash;
 
+use crate::checkpoint::Checkpoint;
+
 #[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize, Encode, Decode, TreeHash)]
 pub struct SignedVote {
     pub data: Vote,
@@ -14,10 +15,7 @@ pub struct SignedVote {
 pub struct Vote {
     pub validator_id: u64,
     pub slot: u64,
-    pub head: B256,
-    pub head_slot: u64,
-    pub target: B256,
-    pub target_slot: u64,
-    pub source: B256,
-    pub source_slot: u64,
+    pub head: Checkpoint,
+    pub target: Checkpoint,
+    pub source: Checkpoint,
 }

crates/common/validator/lean/src/service.rs

@@ -110,7 +110,7 @@ impl ValidatorService {
                             // Build the vote from LeanChain, and modify its validator ID
                             let vote_template = self.lean_chain.read().await.build_vote().expect("Failed to build vote");
 
-                            info!("Built vote template for head {:?} at slot {} with target {:?}", vote_template.head, vote_template.slot, vote_template.target);
+                            info!("Built vote template for head {:?} at slot {} with target {:?}", vote_template.head, vote_template.slot, vote_template.target.slot);
 
                             let votes = self.keystores.iter().map(|keystore| {
                                 let mut vote = vote_template.clone();