Add set_all_blocks_to_optimistic

Author: paulhauner

consensus/proto_array/src/proto_array_fork_choice.rs

@@ -290,6 +290,77 @@ impl ProtoArrayForkChoice {
             .map_err(|e| format!("find_head failed: {:?}", e))
     }
 
+    /// For all nodes, regardless of their relationship to the finalized block, set their execution
+    /// status to be optimistic.
+    ///
+    /// In practice this means forgetting any `VALID` or `INVALID` statuses.
+    pub fn set_all_blocks_to_optimistic(&mut self) -> Result<(), String> {
+        // Iterate backwards through all nodes in the `proto_array`. Whilst it's not strictly
+        // required to do this process in reverse, it seems natural when we consider how LMD votes
+        // are counted.
+        //
+        // This function will touch all blocks, even those that do not descend from the finalized
+        // block. Since this function is expected to run at start-up during very rare
+        // circumstances we prefer simplicity over efficiency.
+        for node_index in (0..self.proto_array.nodes.len()).rev() {
+            let node = self
+                .proto_array
+                .nodes
+                .get_mut(node_index)
+                .ok_or("unreachable index out of bounds in proto_array nodes")?;
+
+            match node.execution_status {
+                ExecutionStatus::Invalid(block_hash) => {
+                    node.execution_status = ExecutionStatus::Optimistic(block_hash);
+
+                    // Restore the weight of the node, it would have been set to `0` in
+                    // `apply_score_changes` when it was invalidated.
+                    let restored_weight = self
+                        .votes
+                        .0
+                        .iter()
+                        .enumerate()
+                        .filter_map(|(validator_index, vote)| {
+                            if vote.current_root == node.root {
+                                Some(validator_index)
+                            } else {
+                                None
+                            }
+                        })
+                        // This `filter_map` asserts that any voting validator that does not have a
+                        // balance should be ignored. This is consistent with `compute_deltas`.
+                        .filter_map(|validator_index| self.balances.get(validator_index))
+                        .sum();
+
+                    // Add the restored weight to the node and all ancestors.
+                    if restored_weight > 0 {
+                        let mut node_or_ancestor = node;
+                        while let Some(parent_index) = node_or_ancestor.parent {
+                            node_or_ancestor
+                                .weight
+                                .checked_add(restored_weight)
+                                .ok_or("Overflow when adding weight to ancestor")?;
+                            node_or_ancestor = self
+                                .proto_array
+                                .nodes
+                                .get_mut(parent_index)
+                                .ok_or(format!("Missing parent index: {}", parent_index))?;
+                        }
+                    }
+                }
+                // There are no balance changes required if the node was either valid or
+                // optimistic.
+                ExecutionStatus::Valid(block_hash) | ExecutionStatus::Optimistic(block_hash) => {
+                    node.execution_status = ExecutionStatus::Optimistic(block_hash)
+                }
+                // An irrelevant node cannot become optimistic, this is a no-op.
+                ExecutionStatus::Irrelevant(_) => (),
+            }
+        }
+
+        Ok(())
+    }
+
     pub fn maybe_prune(&mut self, finalized_root: Hash256) -> Result<(), String> {
         self.proto_array
             .maybe_prune(finalized_root)