diff --git a/eth2/lmd_ghost/src/lib.rs b/eth2/lmd_ghost/src/lib.rs index 98b706fdff..46102c7615 100644 --- a/eth2/lmd_ghost/src/lib.rs +++ b/eth2/lmd_ghost/src/lib.rs @@ -1,10 +1,13 @@ mod proto_array; +mod ssz_container; use parking_lot::RwLock; -use std::collections::HashMap; -use types::{Epoch, Hash256, Slot}; - use proto_array::ProtoArray; +use ssz::{Decode, Encode}; +use ssz_container::SszContainer; +use ssz_derive::{Decode, Encode}; +use std::collections::HashMap; +use types::{Epoch, Hash256}; pub const DEFAULT_PRUNE_THRESHOLD: usize = 256; @@ -27,7 +30,7 @@ pub enum Error { InvalidFindHeadStartRoot, } -#[derive(Default, PartialEq, Clone)] +#[derive(Default, PartialEq, Clone, Encode, Decode)] pub struct VoteTracker { current_root: Hash256, next_root: Hash256, @@ -114,9 +117,9 @@ impl ProtoArrayForkChoice { block_epoch: Epoch, ) -> Result<(), String> { let mut votes = self.votes.write(); + let vote = votes.get_mut(validator_index); - if block_epoch > votes.get(validator_index).next_epoch { - let vote = votes.get_mut(validator_index); + if block_epoch > vote.next_epoch || *vote == VoteTracker::default() { vote.next_root = block_root; vote.next_epoch = block_epoch; } @@ -198,20 +201,30 @@ impl ProtoArrayForkChoice { self.proto_array.read().nodes.len() } - fn latest_message(&self, validator_index: usize) -> Option<(Hash256, Slot)> { - unimplemented!() + pub fn latest_message(&self, validator_index: usize) -> Option<(Hash256, Epoch)> { + let votes = self.votes.read(); + + if validator_index < votes.0.len() { + let vote = &votes.0[validator_index]; + + if *vote == VoteTracker::default() { + None + } else { + Some((vote.next_root, vote.next_epoch)) + } + } else { + None + } } - fn verify_integrity(&self) -> Result<(), String> { - unimplemented!() + pub fn as_bytes(&self) -> Vec { + SszContainer::from(self).as_ssz_bytes() } - fn as_bytes(&self) -> Vec { - unimplemented!() - } - - fn from_bytes(bytes: &[u8]) -> Result { - unimplemented!() + pub fn from_bytes(bytes: &[u8]) -> Result { + SszContainer::from_ssz_bytes(bytes) + .map(Into::into) + .map_err(|e| format!("Failed to decode ProtoArrayForkChoice: {:?}", e)) } } diff --git a/eth2/lmd_ghost/src/proto_array.rs b/eth2/lmd_ghost/src/proto_array.rs index bc427516d0..c0cdfd6c1f 100644 --- a/eth2/lmd_ghost/src/proto_array.rs +++ b/eth2/lmd_ghost/src/proto_array.rs @@ -1,8 +1,9 @@ use crate::Error; +use ssz_derive::{Decode, Encode}; use std::collections::HashMap; use types::{Epoch, Hash256}; -#[derive(Clone, PartialEq, Debug)] +#[derive(Clone, PartialEq, Debug, Encode, Decode)] pub struct ProtoNode { root: Hash256, parent: Option, diff --git a/eth2/lmd_ghost/src/ssz_container.rs b/eth2/lmd_ghost/src/ssz_container.rs new file mode 100644 index 0000000000..4931bde42a --- /dev/null +++ b/eth2/lmd_ghost/src/ssz_container.rs @@ -0,0 +1,60 @@ +use crate::{ + proto_array::{ProtoArray, ProtoNode}, + ElasticList, ProtoArrayForkChoice, VoteTracker, +}; +use parking_lot::RwLock; +use ssz_derive::{Decode, Encode}; +use std::collections::HashMap; +use std::iter::FromIterator; +use types::{Epoch, Hash256}; + +#[derive(Encode, Decode)] +pub struct SszContainer { + votes: Vec, + balances: Vec, + prune_threshold: usize, + ffg_update_required: bool, + justified_epoch: Epoch, + finalized_epoch: Epoch, + finalized_root: Hash256, + nodes: Vec, + indices: Vec<(Hash256, usize)>, +} + +impl From<&ProtoArrayForkChoice> for SszContainer { + fn from(from: &ProtoArrayForkChoice) -> Self { + let proto_array = from.proto_array.read(); + + Self { + votes: from.votes.read().0.clone(), + balances: from.balances.read().clone(), + prune_threshold: proto_array.prune_threshold, + ffg_update_required: proto_array.ffg_update_required, + justified_epoch: proto_array.justified_epoch, + finalized_epoch: proto_array.finalized_epoch, + finalized_root: proto_array.finalized_root, + nodes: proto_array.nodes.clone(), + indices: proto_array.indices.iter().map(|(k, v)| (*k, *v)).collect(), + } + } +} + +impl From for ProtoArrayForkChoice { + fn from(from: SszContainer) -> Self { + let proto_array = ProtoArray { + prune_threshold: from.prune_threshold, + ffg_update_required: from.ffg_update_required, + justified_epoch: from.justified_epoch, + finalized_epoch: from.finalized_epoch, + finalized_root: from.finalized_root, + nodes: from.nodes, + indices: HashMap::from_iter(from.indices.into_iter()), + }; + + Self { + proto_array: RwLock::new(proto_array), + votes: RwLock::new(ElasticList(from.votes)), + balances: RwLock::new(from.balances), + } + } +} diff --git a/eth2/lmd_ghost/tests/test.rs b/eth2/lmd_ghost/tests/test.rs index d99760b4fc..ae9dae51aa 100644 --- a/eth2/lmd_ghost/tests/test.rs +++ b/eth2/lmd_ghost/tests/test.rs @@ -6,6 +6,16 @@ fn get_hash(i: u64) -> Hash256 { Hash256::from_low_u64_be(i) } +fn check_bytes_round_trip(original: &ProtoArrayForkChoice) { + let bytes = original.as_bytes(); + let decoded = + ProtoArrayForkChoice::from_bytes(&bytes).expect("fork choice should decode from bytes"); + assert!( + *original == decoded, + "fork choice should encode and decode without change" + ); +} + /// This tests does not use any validator votes, it just relies on hash-sorting to find the /// head. #[test] @@ -15,6 +25,8 @@ fn no_votes() { let fork_choice = ProtoArrayForkChoice::new(Epoch::new(0), Epoch::new(0), get_hash(0)) .expect("should create fork choice"); + check_bytes_round_trip(&fork_choice); + assert_eq!( fork_choice .find_head( @@ -128,6 +140,8 @@ fn no_votes() { .process_block(get_hash(4), get_hash(2), Epoch::new(0), Epoch::new(0)) .expect("should process block"); + check_bytes_round_trip(&fork_choice); + // Ensure the head is 4. // // 0 @@ -555,6 +569,8 @@ fn votes() { .process_block(get_hash(5), get_hash(4), Epoch::new(1), Epoch::new(1)) .expect("should process block"); + check_bytes_round_trip(&fork_choice); + // Ensure that 5 is filtered out and the head stays at 4. // // 0 @@ -910,6 +926,8 @@ fn votes() { // Set pruning to an unreachable value. fork_choice.set_prune_threshold(usize::max_value()); + check_bytes_round_trip(&fork_choice); + // Run find-head to trigger a prune. assert_eq!( fork_choice @@ -966,6 +984,8 @@ fn votes() { // Ensure that pruning happened. assert_eq!(fork_choice.len(), 6, "there should be 6 blocks"); + check_bytes_round_trip(&fork_choice); + // Add block 11 // // 5 6