Reuse fork-choice child traversal for compliance leaves

This commit is contained in:
dapplion
2026-06-08 15:34:16 +02:00
parent 0223d7ac43
commit 2a938715eb
2 changed files with 81 additions and 97 deletions

View File

@@ -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<E: EthSpec>(
&self,
justified_root: &Hash256,
current_slot: Slot,
justified_checkpoint: Checkpoint,
finalized_checkpoint: Checkpoint,
proposer_boost_root: Hash256,
justified_balances: &JustifiedBalances,
spec: &ChainSpec,
) -> Result<Vec<(Hash256, PayloadStatus, u64)>, Error> {
let start_index = self
.indices
.get(justified_root)
.copied()
.ok_or(Error::NodeUnknown(*justified_root))?;
let viable_nodes = self.get_filtered_block_tree::<E>(
start_index,
current_slot,
justified_checkpoint,
finalized_checkpoint,
)?;
let apply_proposer_boost =
self.should_apply_proposer_boost::<E>(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::<E>(
&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<E: EthSpec>(

View File

@@ -1133,107 +1133,17 @@ impl ProtoArrayForkChoice {
justified_balances: &JustifiedBalances,
spec: &ChainSpec,
) -> Result<Vec<(Hash256, PayloadStatus, u64)>, 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::<E>(
start_index,
self.proto_array
.filtered_block_tree_leaves_and_weights::<E>(
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::<E>(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::<E>(
&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.