Add ssz encoding, other pub fns

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<u8> {
+        SszContainer::from(self).as_ssz_bytes()
     }
 
-    fn as_bytes(&self) -> Vec<u8> {
-        unimplemented!()
-    }
-
-    fn from_bytes(bytes: &[u8]) -> Result<Self, String> {
-        unimplemented!()
+    pub fn from_bytes(bytes: &[u8]) -> Result<Self, String> {
+        SszContainer::from_ssz_bytes(bytes)
+            .map(Into::into)
+            .map_err(|e| format!("Failed to decode ProtoArrayForkChoice: {:?}", e))
     }
 }
 

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<usize>,

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<VoteTracker>,
+    balances: Vec<u64>,
+    prune_threshold: usize,
+    ffg_update_required: bool,
+    justified_epoch: Epoch,
+    finalized_epoch: Epoch,
+    finalized_root: Hash256,
+    nodes: Vec<ProtoNode>,
+    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<SszContainer> 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),
+        }
+    }
+}

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