From 2a938715eb8ed16b180ec33ae3ac94b5c14d73e1 Mon Sep 17 00:00:00 2001 From: dapplion <35266934+dapplion@users.noreply.github.com> Date: Mon, 8 Jun 2026 15:34:16 +0200 Subject: [PATCH] Reuse fork-choice child traversal for compliance leaves --- consensus/proto_array/src/proto_array.rs | 74 +++++++++++++ .../src/proto_array_fork_choice.rs | 104 ++---------------- 2 files changed, 81 insertions(+), 97 deletions(-) diff --git a/consensus/proto_array/src/proto_array.rs b/consensus/proto_array/src/proto_array.rs index 08bf9f87f1..6931bdc34a 100644 --- a/consensus/proto_array/src/proto_array.rs +++ b/consensus/proto_array/src/proto_array.rs @@ -1298,6 +1298,80 @@ impl ProtoArray { } } + /// Returns every leaf node in the filtered block tree, along with its fork-choice weight. + /// + /// This is similar to `find_head_walk`, except it walks every viable branch instead of taking + /// the maximum child at each step. + #[allow(clippy::too_many_arguments)] + pub(crate) fn filtered_block_tree_leaves_and_weights( + &self, + justified_root: &Hash256, + current_slot: Slot, + justified_checkpoint: Checkpoint, + finalized_checkpoint: Checkpoint, + proposer_boost_root: Hash256, + justified_balances: &JustifiedBalances, + spec: &ChainSpec, + ) -> Result, Error> { + let start_index = self + .indices + .get(justified_root) + .copied() + .ok_or(Error::NodeUnknown(*justified_root))?; + + let viable_nodes = self.get_filtered_block_tree::( + start_index, + current_slot, + justified_checkpoint, + finalized_checkpoint, + )?; + + let apply_proposer_boost = + self.should_apply_proposer_boost::(proposer_boost_root, justified_balances, spec)?; + + let mut leaves = Vec::new(); + let mut stack = vec![IndexedForkChoiceNode { + root: *justified_root, + proto_node_index: start_index, + payload_status: PayloadStatus::Pending, + }]; + + while let Some(fc_node) = stack.pop() { + let proto_node = self + .nodes + .get(fc_node.proto_node_index) + .ok_or(Error::InvalidNodeIndex(fc_node.proto_node_index))?; + + let children: Vec<_> = self + .get_node_children(&fc_node)? + .into_iter() + .filter(|(child, _)| viable_nodes.contains(&child.proto_node_index)) + .collect(); + + if children.is_empty() { + let leaf_node = if proto_node.payload_received().is_err() { + fc_node.with_status(PayloadStatus::Pending) + } else { + fc_node + }; + let weight = self.get_weight::( + &leaf_node, + proto_node, + apply_proposer_boost, + proposer_boost_root, + current_slot, + justified_balances, + spec, + )?; + leaves.push((leaf_node.root, leaf_node.payload_status, weight)); + } else { + stack.extend(children.into_iter().map(|(child, _)| child)); + } + } + + Ok(leaves) + } + /// Returns the canonical payload status of a block, matching the decision /// `get_head` would make between `(root, FULL)` and `(root, EMPTY)`. pub(crate) fn get_canonical_payload_status( diff --git a/consensus/proto_array/src/proto_array_fork_choice.rs b/consensus/proto_array/src/proto_array_fork_choice.rs index 2620ff08a5..c75a5d023f 100644 --- a/consensus/proto_array/src/proto_array_fork_choice.rs +++ b/consensus/proto_array/src/proto_array_fork_choice.rs @@ -1133,107 +1133,17 @@ impl ProtoArrayForkChoice { justified_balances: &JustifiedBalances, spec: &ChainSpec, ) -> Result, String> { - let start_index = self - .proto_array - .indices - .get(justified_root) - .copied() - .ok_or_else(|| { - format!( - "filtered_block_tree_leaves_and_weights: justified node \ - {justified_root:?} unknown" - ) - })?; - let viable_indices = self - .proto_array - .get_filtered_block_tree::( - start_index, + self.proto_array + .filtered_block_tree_leaves_and_weights::( + justified_root, current_slot, justified_checkpoint, finalized_checkpoint, + proposer_boost_root, + justified_balances, + spec, ) - .map_err(|e| format!("get_filtered_block_tree failed: {e:?}"))?; - - let apply_proposer_boost = self - .proto_array - .should_apply_proposer_boost::(proposer_boost_root, justified_balances, spec) - .map_err(|e| format!("should_apply_proposer_boost failed: {e:?}"))?; - - let mut leaves = Vec::new(); - let mut stack = vec![IndexedForkChoiceNode { - root: *justified_root, - proto_node_index: start_index, - payload_status: PayloadStatus::Pending, - }]; - - while let Some(fc_node) = stack.pop() { - let proto_node = self - .proto_array - .nodes - .get(fc_node.proto_node_index) - .ok_or_else(|| format!("invalid viable node index {}", fc_node.proto_node_index))?; - - let children = if proto_node.payload_received().is_ok() { - if fc_node.payload_status == PayloadStatus::Pending { - let mut children = vec![fc_node.with_status(PayloadStatus::Empty)]; - if proto_node.payload_received().is_ok_and(|received| received) { - children.push(fc_node.with_status(PayloadStatus::Full)); - } - children - } else { - self.proto_array - .nodes - .iter() - .enumerate() - .filter(|(child_index, child_node)| { - viable_indices.contains(child_index) - && child_node.parent() == Some(fc_node.proto_node_index) - && child_node.get_parent_payload_status() == fc_node.payload_status - }) - .map(|(child_index, child_node)| IndexedForkChoiceNode { - root: child_node.root(), - proto_node_index: child_index, - payload_status: PayloadStatus::Pending, - }) - .collect() - } - } else { - self.proto_array - .nodes - .iter() - .enumerate() - .filter(|(child_index, child_node)| { - viable_indices.contains(child_index) - && child_node.parent() == Some(fc_node.proto_node_index) - }) - .map(|(child_index, child_node)| IndexedForkChoiceNode { - root: child_node.root(), - proto_node_index: child_index, - payload_status: PayloadStatus::Pending, - }) - .collect() - }; - - if children.is_empty() { - let weight = self - .proto_array - .get_weight::( - &fc_node, - proto_node, - apply_proposer_boost, - proposer_boost_root, - current_slot, - justified_balances, - spec, - ) - .map_err(|e| format!("get_weight failed: {e:?}"))?; - leaves.push((fc_node.root, fc_node.payload_status, weight)); - } else { - stack.extend(children); - } - } - - Ok(leaves) + .map_err(|e| format!("filtered_block_tree_leaves_and_weights failed: {e:?}")) } /// Returns the payload status of the head node based on accumulated weights and tiebreaker.