Rename lmd_ghost > proto_array_fork_choice

eth2/proto_array_fork_choice/src/lib.rs

@@ -0,0 +1,705 @@
+mod proto_array;
+mod ssz_container;
+
+use parking_lot::RwLock;
+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;
+
+#[derive(Clone, PartialEq, Debug)]
+pub enum Error {
+    FinalizedNodeUnknown(Hash256),
+    JustifiedNodeUnknown(Hash256),
+    InvalidFinalizedRootChange,
+    InvalidNodeIndex(usize),
+    InvalidParentIndex(usize),
+    InvalidBestChildIndex(usize),
+    InvalidJustifiedIndex(usize),
+    InvalidBestDescendant(usize),
+    InvalidParentDelta(usize),
+    InvalidNodeDelta(usize),
+    DeltaOverflow(usize),
+    IndexOverflow(&'static str),
+    InvalidDeltaLen { deltas: usize, indices: usize },
+    RevertedFinalizedEpoch,
+    InvalidFindHeadStartRoot,
+}
+
+#[derive(Default, PartialEq, Clone, Encode, Decode)]
+pub struct VoteTracker {
+    current_root: Hash256,
+    next_root: Hash256,
+    next_epoch: Epoch,
+}
+
+/// A Vec-wrapper which will grow to match any request.
+///
+/// E.g., a `get` or `insert` to an out-of-bounds element will cause the Vec to grow (using
+/// Default) to the smallest size required to fulfill the request.
+#[derive(Default, Clone, Debug, PartialEq)]
+pub struct ElasticList<T>(Vec<T>);
+
+impl<T> ElasticList<T>
+where
+    T: Default,
+{
+    fn ensure(&mut self, i: usize) {
+        if self.0.len() <= i {
+            self.0.resize_with(i + 1, Default::default);
+        }
+    }
+
+    pub fn get(&mut self, i: usize) -> &T {
+        self.ensure(i);
+        &self.0[i]
+    }
+
+    pub fn get_mut(&mut self, i: usize) -> &mut T {
+        self.ensure(i);
+        &mut self.0[i]
+    }
+
+    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> {
+        self.0.iter_mut()
+    }
+}
+
+pub struct ProtoArrayForkChoice {
+    proto_array: RwLock<ProtoArray>,
+    votes: RwLock<ElasticList<VoteTracker>>,
+    balances: RwLock<Vec<u64>>,
+}
+
+impl PartialEq for ProtoArrayForkChoice {
+    fn eq(&self, other: &Self) -> bool {
+        *self.proto_array.read() == *other.proto_array.read()
+            && *self.votes.read() == *other.votes.read()
+            && *self.balances.read() == *other.balances.read()
+    }
+}
+
+impl ProtoArrayForkChoice {
+    pub fn new(
+        justified_epoch: Epoch,
+        finalized_epoch: Epoch,
+        finalized_root: Hash256,
+    ) -> Result<Self, String> {
+        let mut proto_array = ProtoArray {
+            prune_threshold: DEFAULT_PRUNE_THRESHOLD,
+            ffg_update_required: false,
+            justified_epoch,
+            finalized_epoch,
+            finalized_root,
+            nodes: Vec::with_capacity(1),
+            indices: HashMap::with_capacity(1),
+        };
+
+        proto_array
+            .on_new_block(finalized_root, None, justified_epoch, finalized_epoch)
+            .map_err(|e| format!("Failed to add finalized block to proto_array: {:?}", e))?;
+
+        Ok(Self {
+            proto_array: RwLock::new(proto_array),
+            votes: RwLock::new(ElasticList::default()),
+            balances: RwLock::new(vec![]),
+        })
+    }
+
+    pub fn process_attestation(
+        &self,
+        validator_index: usize,
+        block_root: Hash256,
+        block_epoch: Epoch,
+    ) -> Result<(), String> {
+        let mut votes = self.votes.write();
+        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;
+        }
+
+        Ok(())
+    }
+
+    pub fn process_block(
+        &self,
+        block_root: Hash256,
+        parent_root: Hash256,
+        justified_epoch: Epoch,
+        finalized_epoch: Epoch,
+    ) -> Result<(), String> {
+        self.proto_array
+            .write()
+            .on_new_block(
+                block_root,
+                Some(parent_root),
+                justified_epoch,
+                finalized_epoch,
+            )
+            .map_err(|e| format!("process_block_error: {:?}", e))
+    }
+
+    pub fn find_head(
+        &self,
+        justified_epoch: Epoch,
+        justified_root: Hash256,
+        finalized_epoch: Epoch,
+        finalized_root: Hash256,
+        justified_state_balances: &[u64],
+    ) -> Result<Hash256, String> {
+        let mut proto_array = self.proto_array.write();
+        let mut votes = self.votes.write();
+        let mut old_balances = self.balances.write();
+
+        let new_balances = justified_state_balances;
+
+        proto_array
+            .maybe_prune(finalized_epoch, finalized_root)
+            .map_err(|e| format!("find_head maybe_prune failed: {:?}", e))?;
+
+        let deltas = compute_deltas(
+            &proto_array.indices,
+            &mut votes,
+            &old_balances,
+            &new_balances,
+        )
+        .map_err(|e| format!("find_head compute_deltas failed: {:?}", e))?;
+
+        proto_array
+            .apply_score_changes(deltas, justified_epoch)
+            .map_err(|e| format!("find_head apply_score_changes failed: {:?}", e))?;
+
+        *old_balances = new_balances.to_vec();
+
+        proto_array
+            .find_head(&justified_root)
+            .map_err(|e| format!("find_head failed: {:?}", e))
+    }
+
+    pub fn update_finalized_root(
+        &self,
+        finalized_epoch: Epoch,
+        finalized_root: Hash256,
+    ) -> Result<(), String> {
+        self.proto_array
+            .write()
+            .maybe_prune(finalized_epoch, finalized_root)
+            .map_err(|e| format!("find_head maybe_prune failed: {:?}", e))
+    }
+
+    pub fn set_prune_threshold(&self, prune_threshold: usize) {
+        self.proto_array.write().prune_threshold = prune_threshold;
+    }
+
+    pub fn len(&self) -> usize {
+        self.proto_array.read().nodes.len()
+    }
+
+    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
+        }
+    }
+
+    pub fn as_bytes(&self) -> Vec<u8> {
+        SszContainer::from(self).as_ssz_bytes()
+    }
+
+    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))
+    }
+}
+
+/// Returns a list of `deltas`, where there is one delta for each of the indices in
+/// `0..indices.len()`.
+///
+/// The deltas are formed by a change between `old_balances` and `new_balances`, and/or a change of vote in `votes`.
+///
+/// ## Errors
+///
+/// - If a value in `indices` is greater to or equal to `indices.len()`.
+/// - If some `Hash256` in `votes` is not a key in `indices` (except for `Hash256::zero()`, this is
+/// always valid).
+fn compute_deltas(
+    indices: &HashMap<Hash256, usize>,
+    votes: &mut ElasticList<VoteTracker>,
+    old_balances: &[u64],
+    new_balances: &[u64],
+) -> Result<Vec<i64>, Error> {
+    let mut deltas = vec![0_i64; indices.len()];
+
+    for (val_index, vote) in votes.iter_mut().enumerate() {
+        // There is no need to create a score change if the validator has never voted or both their
+        // votes are for the zero hash (alias to the genesis block).
+        if vote.current_root == Hash256::zero() && vote.next_root == Hash256::zero() {
+            continue;
+        }
+
+        // If the validator was not included in the _old_ balances (i.e., it did not exist yet)
+        // then say its balance was zero.
+        let old_balance = old_balances.get(val_index).copied().unwrap_or_else(|| 0);
+
+        // If the validators vote is not known in the _new_ balances, then use a balance of zero.
+        //
+        // It is possible that there is a vote for an unknown validator if we change our justified
+        // state to a new state with a higher epoch that is on a different fork because that fork may have
+        // on-boarded less validators than the prior fork.
+        let new_balance = new_balances.get(val_index).copied().unwrap_or_else(|| 0);
+
+        if vote.current_root != vote.next_root || old_balance != new_balance {
+            // We ignore the vote if it is not known in `indices`. We assume that it is outside
+            // of our tree (i.e., pre-finalization) and therefore not interesting.
+            if let Some(current_delta_index) = indices.get(&vote.current_root).copied() {
+                let delta = deltas
+                    .get_mut(current_delta_index)
+                    .ok_or_else(|| Error::InvalidNodeDelta(current_delta_index))?
+                    .checked_sub(old_balance as i64)
+                    .ok_or_else(|| Error::DeltaOverflow(current_delta_index))?;
+
+                // Array access safe due to check on previous line.
+                deltas[current_delta_index] = delta;
+            }
+
+            // We ignore the vote if it is not known in `indices`. We assume that it is outside
+            // of our tree (i.e., pre-finalization) and therefore not interesting.
+            if let Some(next_delta_index) = indices.get(&vote.next_root).copied() {
+                let delta = deltas
+                    .get(next_delta_index)
+                    .ok_or_else(|| Error::InvalidNodeDelta(next_delta_index))?
+                    .checked_add(new_balance as i64)
+                    .ok_or_else(|| Error::DeltaOverflow(next_delta_index))?;
+
+                // Array access safe due to check on previous line.
+                deltas[next_delta_index] = delta;
+            }
+
+            vote.current_root = vote.next_root;
+        }
+    }
+
+    Ok(deltas)
+}
+
+#[cfg(test)]
+mod test_compute_deltas {
+    use super::*;
+
+    /// Gives a hash that is not the zero hash (unless i is `usize::max_value)`.
+    fn hash_from_index(i: usize) -> Hash256 {
+        Hash256::from_low_u64_be(i as u64 + 1)
+    }
+
+    #[test]
+    fn zero_hash() {
+        let validator_count: usize = 16;
+
+        let mut indices = HashMap::new();
+        let mut votes = ElasticList::default();
+        let mut old_balances = vec![];
+        let mut new_balances = vec![];
+
+        for i in 0..validator_count {
+            indices.insert(hash_from_index(i), i);
+            votes.0.push(VoteTracker {
+                current_root: Hash256::zero(),
+                next_root: Hash256::zero(),
+                next_epoch: Epoch::new(0),
+            });
+            old_balances.push(0);
+            new_balances.push(0);
+        }
+
+        let deltas = compute_deltas(&indices, &mut votes, &old_balances, &new_balances)
+            .expect("should compute deltas");
+
+        assert_eq!(
+            deltas.len(),
+            validator_count,
+            "deltas should have expected length"
+        );
+        assert_eq!(
+            deltas,
+            vec![0; validator_count],
+            "deltas should all be zero"
+        );
+
+        for vote in votes.0 {
+            assert_eq!(
+                vote.current_root, vote.next_root,
+                "the vote shoulds should have been updated"
+            );
+        }
+    }
+
+    #[test]
+    fn all_voted_the_same() {
+        const BALANCE: u64 = 42;
+
+        let validator_count: usize = 16;
+
+        let mut indices = HashMap::new();
+        let mut votes = ElasticList::default();
+        let mut old_balances = vec![];
+        let mut new_balances = vec![];
+
+        for i in 0..validator_count {
+            indices.insert(hash_from_index(i), i);
+            votes.0.push(VoteTracker {
+                current_root: Hash256::zero(),
+                next_root: hash_from_index(0),
+                next_epoch: Epoch::new(0),
+            });
+            old_balances.push(BALANCE);
+            new_balances.push(BALANCE);
+        }
+
+        let deltas = compute_deltas(&indices, &mut votes, &old_balances, &new_balances)
+            .expect("should compute deltas");
+
+        assert_eq!(
+            deltas.len(),
+            validator_count,
+            "deltas should have expected length"
+        );
+
+        for (i, delta) in deltas.into_iter().enumerate() {
+            if i == 0 {
+                assert_eq!(
+                    delta,
+                    BALANCE as i64 * validator_count as i64,
+                    "zero'th root should have a delta"
+                );
+            } else {
+                assert_eq!(delta, 0, "all other deltas should be zero");
+            }
+        }
+
+        for vote in votes.0 {
+            assert_eq!(
+                vote.current_root, vote.next_root,
+                "the vote shoulds should have been updated"
+            );
+        }
+    }
+
+    #[test]
+    fn different_votes() {
+        const BALANCE: u64 = 42;
+
+        let validator_count: usize = 16;
+
+        let mut indices = HashMap::new();
+        let mut votes = ElasticList::default();
+        let mut old_balances = vec![];
+        let mut new_balances = vec![];
+
+        for i in 0..validator_count {
+            indices.insert(hash_from_index(i), i);
+            votes.0.push(VoteTracker {
+                current_root: Hash256::zero(),
+                next_root: hash_from_index(i),
+                next_epoch: Epoch::new(0),
+            });
+            old_balances.push(BALANCE);
+            new_balances.push(BALANCE);
+        }
+
+        let deltas = compute_deltas(&indices, &mut votes, &old_balances, &new_balances)
+            .expect("should compute deltas");
+
+        assert_eq!(
+            deltas.len(),
+            validator_count,
+            "deltas should have expected length"
+        );
+
+        for delta in deltas.into_iter() {
+            assert_eq!(
+                delta, BALANCE as i64,
+                "each root should have the same delta"
+            );
+        }
+
+        for vote in votes.0 {
+            assert_eq!(
+                vote.current_root, vote.next_root,
+                "the vote shoulds should have been updated"
+            );
+        }
+    }
+
+    #[test]
+    fn moving_votes() {
+        const BALANCE: u64 = 42;
+
+        let validator_count: usize = 16;
+
+        let mut indices = HashMap::new();
+        let mut votes = ElasticList::default();
+        let mut old_balances = vec![];
+        let mut new_balances = vec![];
+
+        for i in 0..validator_count {
+            indices.insert(hash_from_index(i), i);
+            votes.0.push(VoteTracker {
+                current_root: hash_from_index(0),
+                next_root: hash_from_index(1),
+                next_epoch: Epoch::new(0),
+            });
+            old_balances.push(BALANCE);
+            new_balances.push(BALANCE);
+        }
+
+        let deltas = compute_deltas(&indices, &mut votes, &old_balances, &new_balances)
+            .expect("should compute deltas");
+
+        assert_eq!(
+            deltas.len(),
+            validator_count,
+            "deltas should have expected length"
+        );
+
+        let total_delta = BALANCE as i64 * validator_count as i64;
+
+        for (i, delta) in deltas.into_iter().enumerate() {
+            if i == 0 {
+                assert_eq!(
+                    delta,
+                    0 - total_delta,
+                    "zero'th root should have a negative delta"
+                );
+            } else if i == 1 {
+                assert_eq!(delta, total_delta, "first root should have positive delta");
+            } else {
+                assert_eq!(delta, 0, "all other deltas should be zero");
+            }
+        }
+
+        for vote in votes.0 {
+            assert_eq!(
+                vote.current_root, vote.next_root,
+                "the vote shoulds should have been updated"
+            );
+        }
+    }
+
+    #[test]
+    fn move_out_of_tree() {
+        const BALANCE: u64 = 42;
+
+        let mut indices = HashMap::new();
+        let mut votes = ElasticList::default();
+
+        // There is only one block.
+        indices.insert(hash_from_index(1), 0);
+
+        // There are two validators.
+        let old_balances = vec![BALANCE; 2];
+        let new_balances = vec![BALANCE; 2];
+
+        // One validator moves their vote from the block to the zero hash.
+        votes.0.push(VoteTracker {
+            current_root: hash_from_index(1),
+            next_root: Hash256::zero(),
+            next_epoch: Epoch::new(0),
+        });
+
+        // One validator moves their vote from the block to something outside the tree.
+        votes.0.push(VoteTracker {
+            current_root: hash_from_index(1),
+            next_root: Hash256::from_low_u64_be(1337),
+            next_epoch: Epoch::new(0),
+        });
+
+        let deltas = compute_deltas(&indices, &mut votes, &old_balances, &new_balances)
+            .expect("should compute deltas");
+
+        assert_eq!(deltas.len(), 1, "deltas should have expected length");
+
+        assert_eq!(
+            deltas[0],
+            0 - BALANCE as i64 * 2,
+            "the block should have lost both balances"
+        );
+
+        for vote in votes.0 {
+            assert_eq!(
+                vote.current_root, vote.next_root,
+                "the vote shoulds should have been updated"
+            );
+        }
+    }
+
+    #[test]
+    fn changing_balances() {
+        const OLD_BALANCE: u64 = 42;
+        const NEW_BALANCE: u64 = OLD_BALANCE * 2;
+
+        let validator_count: usize = 16;
+
+        let mut indices = HashMap::new();
+        let mut votes = ElasticList::default();
+        let mut old_balances = vec![];
+        let mut new_balances = vec![];
+
+        for i in 0..validator_count {
+            indices.insert(hash_from_index(i), i);
+            votes.0.push(VoteTracker {
+                current_root: hash_from_index(0),
+                next_root: hash_from_index(1),
+                next_epoch: Epoch::new(0),
+            });
+            old_balances.push(OLD_BALANCE);
+            new_balances.push(NEW_BALANCE);
+        }
+
+        let deltas = compute_deltas(&indices, &mut votes, &old_balances, &new_balances)
+            .expect("should compute deltas");
+
+        assert_eq!(
+            deltas.len(),
+            validator_count,
+            "deltas should have expected length"
+        );
+
+        for (i, delta) in deltas.into_iter().enumerate() {
+            if i == 0 {
+                assert_eq!(
+                    delta,
+                    0 - OLD_BALANCE as i64 * validator_count as i64,
+                    "zero'th root should have a negative delta"
+                );
+            } else if i == 1 {
+                assert_eq!(
+                    delta,
+                    NEW_BALANCE as i64 * validator_count as i64,
+                    "first root should have positive delta"
+                );
+            } else {
+                assert_eq!(delta, 0, "all other deltas should be zero");
+            }
+        }
+
+        for vote in votes.0 {
+            assert_eq!(
+                vote.current_root, vote.next_root,
+                "the vote shoulds should have been updated"
+            );
+        }
+    }
+
+    #[test]
+    fn validator_appears() {
+        const BALANCE: u64 = 42;
+
+        let mut indices = HashMap::new();
+        let mut votes = ElasticList::default();
+
+        // There are two blocks.
+        indices.insert(hash_from_index(1), 0);
+        indices.insert(hash_from_index(2), 1);
+
+        // There is only one validator in the old balances.
+        let old_balances = vec![BALANCE; 1];
+        // There are two validators in the new balances.
+        let new_balances = vec![BALANCE; 2];
+
+        // Both validator move votes from block 1 to block 2.
+        for _ in 0..2 {
+            votes.0.push(VoteTracker {
+                current_root: hash_from_index(1),
+                next_root: hash_from_index(2),
+                next_epoch: Epoch::new(0),
+            });
+        }
+
+        let deltas = compute_deltas(&indices, &mut votes, &old_balances, &new_balances)
+            .expect("should compute deltas");
+
+        assert_eq!(deltas.len(), 2, "deltas should have expected length");
+
+        assert_eq!(
+            deltas[0],
+            0 - BALANCE as i64,
+            "block 1 should have only lost one balance"
+        );
+        assert_eq!(
+            deltas[1],
+            2 * BALANCE as i64,
+            "block 2 should have gained two balances"
+        );
+
+        for vote in votes.0 {
+            assert_eq!(
+                vote.current_root, vote.next_root,
+                "the vote shoulds should have been updated"
+            );
+        }
+    }
+
+    #[test]
+    fn validator_disappears() {
+        const BALANCE: u64 = 42;
+
+        let mut indices = HashMap::new();
+        let mut votes = ElasticList::default();
+
+        // There are two blocks.
+        indices.insert(hash_from_index(1), 0);
+        indices.insert(hash_from_index(2), 1);
+
+        // There are two validators in the old balances.
+        let old_balances = vec![BALANCE; 2];
+        // There is only one validator in the new balances.
+        let new_balances = vec![BALANCE; 1];
+
+        // Both validator move votes from block 1 to block 2.
+        for _ in 0..2 {
+            votes.0.push(VoteTracker {
+                current_root: hash_from_index(1),
+                next_root: hash_from_index(2),
+                next_epoch: Epoch::new(0),
+            });
+        }
+
+        let deltas = compute_deltas(&indices, &mut votes, &old_balances, &new_balances)
+            .expect("should compute deltas");
+
+        assert_eq!(deltas.len(), 2, "deltas should have expected length");
+
+        assert_eq!(
+            deltas[0],
+            0 - BALANCE as i64 * 2,
+            "block 1 should have lost both balances"
+        );
+        assert_eq!(
+            deltas[1], BALANCE as i64,
+            "block 2 should have only gained one balance"
+        );
+
+        for vote in votes.0 {
+            assert_eq!(
+                vote.current_root, vote.next_root,
+                "the vote shoulds should have been updated"
+            );
+        }
+    }
+}

eth2/proto_array_fork_choice/src/proto_array.rs

@@ -0,0 +1,427 @@
+use crate::Error;
+use ssz_derive::{Decode, Encode};
+use std::collections::HashMap;
+use types::{Epoch, Hash256};
+
+#[derive(Clone, PartialEq, Debug, Encode, Decode)]
+pub struct ProtoNode {
+    root: Hash256,
+    parent: Option<usize>,
+    justified_epoch: Epoch,
+    finalized_epoch: Epoch,
+    weight: u64,
+    best_child: Option<usize>,
+    best_descendant: Option<usize>,
+}
+
+impl ProtoNode {
+    /// Returns `true` if some node is "better" than the other, according to either weight or root.
+    ///
+    /// If `self == other`, then `true` is returned.
+    pub fn is_better_than(&self, other: &Self) -> bool {
+        if self.weight == other.weight {
+            self.root >= other.root
+        } else {
+            self.weight >= other.weight
+        }
+    }
+}
+
+#[derive(PartialEq)]
+pub struct ProtoArray {
+    /// Do not attempt to prune the tree unless it has at least this many nodes. Small prunes
+    /// simply waste time.
+    pub prune_threshold: usize,
+    /// Set to true when the Casper FFG justified/finalized epochs should be checked to ensure the
+    /// tree is filtered as per eth2 specs.
+    pub ffg_update_required: bool,
+    pub justified_epoch: Epoch,
+    pub finalized_epoch: Epoch,
+    pub finalized_root: Hash256,
+    pub nodes: Vec<ProtoNode>,
+    pub indices: HashMap<Hash256, usize>,
+}
+
+impl ProtoArray {
+    /// Iterate backwards through the array, touching all nodes and their parents and potentially
+    /// the best-child of each parent.
+    ///
+    /// The structure of the `self.nodes` array ensures that the child of each node is always
+    /// touched before it's parent.
+    ///
+    /// For each node, the following is done:
+    ///
+    /// - Update the nodes weight with the corresponding delta.
+    /// - Back-propgrate each nodes delta to its parents delta.
+    /// - Compare the current node with the parents best-child, updating it if the current node
+    /// should become the best child.
+    /// - Update the parents best-descendant with the current node or its best-descendant, if
+    /// required.
+    pub fn apply_score_changes(
+        &mut self,
+        mut deltas: Vec<i64>,
+        justified_epoch: Epoch,
+    ) -> Result<(), Error> {
+        if deltas.len() != self.indices.len() {
+            return Err(Error::InvalidDeltaLen {
+                deltas: deltas.len(),
+                indices: self.indices.len(),
+            });
+        }
+
+        // The `self.ffg_update_required` flag indicates if it is necessary to check the
+        // finalized/justified epoch of all nodes against the epochs in `self`.
+        //
+        // This behaviour is equivalent to the `filter_block_tree` function in the spec.
+        self.ffg_update_required = justified_epoch != self.justified_epoch;
+        if self.ffg_update_required {
+            self.justified_epoch = justified_epoch;
+        }
+
+        // Iterate backwards through all indices in `self.nodes`.
+        for node_index in (0..self.nodes.len()).rev() {
+            let node = &mut self
+                .nodes
+                .get_mut(node_index)
+                .ok_or_else(|| Error::InvalidNodeIndex(node_index))?;
+
+            // There is no need to adjust the balances or manage parent of the zero hash since it
+            // is an alias to the genesis block. The weight applied to the genesis block is
+            // irrelevant as we _always_ choose it and it's impossible for it to have a parent.
+            if node.root == Hash256::zero() {
+                continue;
+            }
+
+            let node_delta = deltas
+                .get(node_index)
+                .copied()
+                .ok_or_else(|| Error::InvalidNodeDelta(node_index))?;
+
+            // Apply the delta to the node.
+            if node_delta < 0 {
+                // Note: I am conflicted about whether to use `saturating_sub` or `checked_sub`
+                // here.
+                //
+                // I can't think of any valid reason why `node_delta.abs()` should be greater than
+                // `node.weight`, so I have chosen `checked_sub` to try and fail-fast if there is
+                // some error.
+                //
+                // However, I am not fully convinced that some valid case for `saturating_sub` does
+                // not exist.
+                node.weight = node
+                    .weight
+                    .checked_sub(node_delta.abs() as u64)
+                    .ok_or_else(|| Error::DeltaOverflow(node_index))?;
+            } else {
+                node.weight = node
+                    .weight
+                    .checked_add(node_delta as u64)
+                    .ok_or_else(|| Error::DeltaOverflow(node_index))?;
+            }
+
+            // If the node has a parent, try to update its best-child and best-descendant.
+            if let Some(parent_index) = node.parent {
+                let parent_delta = deltas
+                    .get_mut(parent_index)
+                    .ok_or_else(|| Error::InvalidParentDelta(parent_index))?;
+
+                // Back-propogate the nodes delta to its parent.
+                *parent_delta += node_delta;
+
+                let is_viable_for_head = self
+                    .nodes
+                    .get(node_index)
+                    .map(|node| self.node_is_viable_for_head(node))
+                    .ok_or_else(|| Error::InvalidNodeIndex(parent_index))?;
+
+                // If the given node is _not viable_ for the head and we are required to check
+                // for FFG changes, then remove the child if it is currently the parents
+                // best-child.
+                if !is_viable_for_head {
+                    if self.ffg_update_required {
+                        let parent_best_child = self
+                            .nodes
+                            .get(parent_index)
+                            .ok_or_else(|| Error::InvalidParentIndex(parent_index))?
+                            .best_child;
+
+                        if parent_best_child == Some(node_index) {
+                            let parent_node = self
+                                .nodes
+                                .get_mut(parent_index)
+                                .ok_or_else(|| Error::InvalidParentIndex(parent_index))?;
+
+                            parent_node.best_child = None;
+                            parent_node.best_descendant = None;
+                        }
+                    }
+
+                    continue;
+                }
+
+                // If the parent has a best-child, see if the current node is better. If it doesn't
+                // have a best child, set it to ours.
+                //
+                // Note: this code only runs if the node is viable for the head due to `continue`
+                // call in previous statement.
+                if let Some(parent_best_child_index) = self
+                    .nodes
+                    .get(parent_index)
+                    .ok_or_else(|| Error::InvalidParentIndex(parent_index))?
+                    .best_child
+                {
+                    // Here we set the best child to `node_index` when that is already the case.
+                    // This has the effect of ensuring the `best_descendant` is updated.
+                    if parent_best_child_index == node_index {
+                        self.set_best_child(parent_index, node_index)?;
+                        continue;
+                    }
+
+                    let parent_best_child = self
+                        .nodes
+                        .get(parent_best_child_index)
+                        .ok_or_else(|| Error::InvalidBestChildIndex(parent_best_child_index))?;
+
+                    let node_is_better_than_current_best_child = self
+                        .nodes
+                        .get(node_index)
+                        .ok_or_else(|| Error::InvalidNodeIndex(node_index))?
+                        .is_better_than(parent_best_child);
+
+                    let current_best_child_is_viable =
+                        self.node_is_viable_for_head(parent_best_child);
+
+                    // There are two conditions for replacing a best child:
+                    //
+                    // - The node is better than the present best-child.
+                    // - The present best-child in no longer viable (viz., it has been filtered out).
+                    if node_is_better_than_current_best_child || !current_best_child_is_viable {
+                        self.set_best_child(parent_index, node_index)?;
+                    }
+                } else {
+                    // If the best child is `None`, simply set it to the current node (noting that
+                    // this code only runs if the current node is viable for the head).
+                    self.set_best_child(parent_index, node_index)?;
+                };
+            }
+        }
+
+        self.ffg_update_required = false;
+
+        Ok(())
+    }
+
+    /// Register a new block with the fork choice.
+    ///
+    /// It is only sane to supply a `None` parent for the genesis block.
+    pub fn on_new_block(
+        &mut self,
+        root: Hash256,
+        parent: Option<Hash256>,
+        justified_epoch: Epoch,
+        finalized_epoch: Epoch,
+    ) -> Result<(), Error> {
+        let node_index = self.nodes.len();
+
+        let node = ProtoNode {
+            root,
+            parent: parent.and_then(|parent_root| self.indices.get(&parent_root).copied()),
+            justified_epoch,
+            finalized_epoch,
+            weight: 0,
+            best_child: None,
+            best_descendant: None,
+        };
+
+        self.indices.insert(node.root, node_index);
+        self.nodes.push(node.clone());
+
+        // If the blocks justified and finalized epochs match our values, then try and see if it
+        // becomes the best child.
+        if justified_epoch == self.justified_epoch && finalized_epoch == self.finalized_epoch {
+            if let Some(parent_index) = node.parent {
+                let parent = self
+                    .nodes
+                    .get(parent_index)
+                    .ok_or_else(|| Error::InvalidParentIndex(parent_index))?;
+
+                if let Some(parent_best_child_index) = parent.best_child {
+                    let parent_best_child = self
+                        .nodes
+                        .get(parent_best_child_index)
+                        .ok_or_else(|| Error::InvalidBestChildIndex(parent_best_child_index))?;
+
+                    if node.is_better_than(parent_best_child) {
+                        self.set_best_child(parent_index, node_index)?;
+                    }
+                } else {
+                    self.set_best_child(parent_index, node_index)?;
+                };
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Follows the best-descendant links to find the best-block (i.e., head-block).
+    ///
+    /// ## Notes
+    ///
+    /// The result of this function is not guaranteed to be accurate if `Self::on_new_block` has
+    /// been called without a subsequent `Self::apply_score_changes` call. This is because
+    /// `on_new_block` does not attempt to walk backwards through the tree and update the
+    /// best-child/best-descendant links.
+    pub fn find_head(&self, justified_root: &Hash256) -> Result<Hash256, Error> {
+        let justified_index = self
+            .indices
+            .get(justified_root)
+            .copied()
+            .ok_or_else(|| Error::JustifiedNodeUnknown(self.finalized_root))?;
+
+        let justified_node = self
+            .nodes
+            .get(justified_index)
+            .ok_or_else(|| Error::InvalidJustifiedIndex(justified_index))?;
+
+        // It is a logic error to try and find the head starting from a block that does not match
+        // the filter.
+        if justified_node.justified_epoch != self.justified_epoch
+            || justified_node.finalized_epoch != self.finalized_epoch
+        {
+            return Err(Error::InvalidFindHeadStartRoot);
+        }
+
+        let best_descendant_index = justified_node
+            .best_descendant
+            .unwrap_or_else(|| justified_index);
+
+        let best_node = self
+            .nodes
+            .get(best_descendant_index)
+            .ok_or_else(|| Error::InvalidBestDescendant(best_descendant_index))?;
+
+        Ok(best_node.root)
+    }
+
+    /// Update the tree with new finalization information. The tree is only actually pruned if both
+    /// of the two following criteria are met:
+    ///
+    /// - The supplied finalized epoch and root are different to the current values.
+    /// - The number of nodes in `self` is at least `self.prune_threshold`.
+    ///
+    /// # Errors
+    ///
+    /// Returns errors if:
+    ///
+    /// - The finalized epoch is less than the current one.
+    /// - The finalized epoch is equal to the current one, but the finalized root is different.
+    /// - There is some internal error relating to invalid indices inside `self`.
+    pub fn maybe_prune(
+        &mut self,
+        finalized_epoch: Epoch,
+        finalized_root: Hash256,
+    ) -> Result<(), Error> {
+        if finalized_epoch == self.finalized_epoch && self.finalized_root != finalized_root {
+            // It's illegal to swap finalized roots on the same epoch (this is reverting a
+            // finalized block).
+            return Err(Error::InvalidFinalizedRootChange);
+        } else if finalized_epoch < self.finalized_epoch {
+            // It's illegal to swap to an earlier finalized root (this is assumed to be reverting a
+            // finalized block).
+            return Err(Error::RevertedFinalizedEpoch);
+        } else if finalized_epoch != self.finalized_epoch {
+            self.finalized_epoch = finalized_epoch;
+            self.finalized_root = finalized_root;
+            self.ffg_update_required = true;
+        }
+
+        let finalized_index = *self
+            .indices
+            .get(&self.finalized_root)
+            .ok_or_else(|| Error::FinalizedNodeUnknown(self.finalized_root))?;
+
+        if finalized_index < self.prune_threshold {
+            // Pruning at small numbers incurs more cost than benefit.
+            return Ok(());
+        }
+
+        // Remove the `self.indices` key/values for all the to-be-deleted nodes.
+        for node_index in 0..finalized_index {
+            let root = &self
+                .nodes
+                .get(node_index)
+                .ok_or_else(|| Error::InvalidNodeIndex(node_index))?
+                .root;
+            self.indices.remove(root);
+        }
+
+        // Drop all the nodes prior to finalization.
+        self.nodes = self.nodes.split_off(finalized_index);
+
+        // Adjust the indices map.
+        for (_root, index) in self.indices.iter_mut() {
+            *index = index
+                .checked_sub(finalized_index)
+                .ok_or_else(|| Error::IndexOverflow("indices"))?;
+        }
+
+        // Iterate through all the existing nodes and adjust their indices to match the new layout
+        // of `self.nodes`.
+        for node in self.nodes.iter_mut() {
+            if let Some(parent) = node.parent {
+                // If `node.parent` is less than `finalized_index`, set it to `None`.
+                node.parent = parent.checked_sub(finalized_index);
+            }
+            if let Some(best_child) = node.best_child {
+                node.best_child = Some(
+                    best_child
+                        .checked_sub(finalized_index)
+                        .ok_or_else(|| Error::IndexOverflow("best_child"))?,
+                );
+            }
+            if let Some(best_descendant) = node.best_descendant {
+                node.best_descendant = Some(
+                    best_descendant
+                        .checked_sub(finalized_index)
+                        .ok_or_else(|| Error::IndexOverflow("best_descendant"))?,
+                );
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Sets the node at `parent_index` to have a best-child pointing to `child_index`. Also
+    /// updates the best-descendant.
+    fn set_best_child(&mut self, parent_index: usize, child_index: usize) -> Result<(), Error> {
+        let child_best_descendant = self
+            .nodes
+            .get(child_index)
+            .ok_or_else(|| Error::InvalidNodeIndex(child_index))?
+            .best_descendant;
+
+        let parent_node = self
+            .nodes
+            .get_mut(parent_index)
+            .ok_or_else(|| Error::InvalidParentIndex(parent_index))?;
+
+        parent_node.best_child = Some(child_index);
+        parent_node.best_descendant = if let Some(best_descendant) = child_best_descendant {
+            Some(best_descendant)
+        } else {
+            Some(child_index)
+        };
+
+        Ok(())
+    }
+
+    /// This is the equivalent to the `filter_block_tree` function in the eth2 spec:
+    ///
+    /// https://github.com/ethereum/eth2.0-specs/blob/v0.10.0/specs/phase0/fork-choice.md#filter_block_tree
+    ///
+    /// Any node that has a different finalized or justified epoch should not be viable for the
+    /// head.
+    fn node_is_viable_for_head(&self, node: &ProtoNode) -> bool {
+        node.justified_epoch == self.justified_epoch && node.finalized_epoch == self.finalized_epoch
+    }
+}

eth2/proto_array_fork_choice/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),
+        }
+    }
+}