From 52e397f8c17ed8bc1d71ed46013318d1851d4ba8 Mon Sep 17 00:00:00 2001 From: Michael Sproul Date: Tue, 24 Mar 2026 15:58:54 +1100 Subject: [PATCH] Refactoring fork choice to look more like the spec --- consensus/proto_array/src/proto_array.rs | 233 +++++++++--------- .../src/proto_array_fork_choice.rs | 7 + 2 files changed, 130 insertions(+), 110 deletions(-) diff --git a/consensus/proto_array/src/proto_array.rs b/consensus/proto_array/src/proto_array.rs index ac5f5be525..5ab879fac4 100644 --- a/consensus/proto_array/src/proto_array.rs +++ b/consensus/proto_array/src/proto_array.rs @@ -153,6 +153,23 @@ pub struct ProtoNode { pub proposer_index: u64, } +impl ProtoNode { + /// Generic version of spec's `parent_payload_status` that works for pre-Gloas nodes by + /// considering their parents Empty. + fn parent_payload_status(&self) -> PayloadStatus { + self.parent_payload_status().unwrap_or(PayloadStatus::Empty) + } + + fn attestation_score(&self, payload_status: PayloadStatus) -> u64 { + match payload_status { + // TODO(gloas): rename weight and remove proposer boost from it? + PayloadStatus::Pending => self.weight, + PayloadStatus::Empty => self.empty_payload_weight().unwrap_or(0), + PayloadStatus::Full => self.full_payload_weight().unwrap_or(0), + } + } +} + #[derive(PartialEq, Debug, Encode, Decode, Serialize, Deserialize, Copy, Clone)] pub struct ProposerBoost { pub root: Hash256, @@ -715,6 +732,19 @@ impl ProtoArray { Ok(()) } + fn is_head_weak(&self, head_node: &ProtoNode, justified_balances: &JustifiedBalances) -> bool { + let reorg_threshold = calculate_committee_fraction::( + justified_balances, + spec.reorg_head_weight_threshold.unwrap_or(20), + ) + .unwrap_or(0); + + let head_weight = head_node.attestation_score(); + + // TODO(gloas): missing equivocating weight from spec + head_weight < reorg_threshold + } + /// Spec's `should_apply_proposer_boost` for Gloas. /// /// Returns `true` if the proposer boost should be kept. Returns `false` if the @@ -746,14 +776,6 @@ impl ProtoArray { return Ok(true); } - // Check if the parent is "weak" (low attestation weight). - // Parent weight currently includes the back-propagated boost, so subtract it. - let reorg_threshold = calculate_committee_fraction::( - justified_balances, - spec.reorg_head_weight_threshold.unwrap_or(20), - ) - .unwrap_or(0); - let parent_weight_without_boost = parent.weight().saturating_sub(proposer_score); if parent_weight_without_boost >= reorg_threshold { return Ok(true); // Parent is not weak — apply. @@ -1089,7 +1111,7 @@ impl ProtoArray { current_slot: Slot, best_justified_checkpoint: Checkpoint, best_finalized_checkpoint: Checkpoint, - ) -> Result { + ) -> Result<(Hash256, PayloadStatus), Error> { let justified_index = self .indices .get(justified_root) @@ -1112,58 +1134,15 @@ impl ProtoArray { }); } - // For V29 (Gloas) justified nodes, use the virtual tree walk directly. - if justified_node.as_v29().is_ok() { - return self.find_head_v29_walk::( - justified_index, - current_slot, - best_justified_checkpoint, - best_finalized_checkpoint, - ); - } - - // Pre-Gloas justified node, but descendants may be V29. - // Walk via best_child chain; switch to V29 walk when we hit one. - if justified_node.best_child().is_some() || justified_node.best_descendant().is_some() { - let mut current_index = justified_index; - loop { - let node = self - .nodes - .get(current_index) - .ok_or(Error::InvalidNodeIndex(current_index))?; - - // Hit a V29 node — switch to virtual tree walk. - if node.as_v29().is_ok() { - return self.find_head_v29_walk::( - current_index, - current_slot, - best_justified_checkpoint, - best_finalized_checkpoint, - ); - } - - // V17 node: follow best_child. - if let Some(bc_idx) = node.best_child() { - current_index = bc_idx; - } else { - break; - } - } - - let head_node = self - .nodes - .get(current_index) - .ok_or(Error::InvalidNodeIndex(current_index))?; - return Ok(head_node.root()); - } - - // Pre-Gloas fallback: use best_descendant directly. - let best_descendant_index = justified_node.best_descendant().unwrap_or(justified_index); - - let best_node = self - .nodes - .get(best_descendant_index) - .ok_or(Error::InvalidBestDescendant(best_descendant_index))?; + // In the post-Gloas worlld, always use a virtual tree walk. + // + // Best child/best descendant is dead. + let best_node = self.find_head_walk::( + justified_index, + current_slot, + best_justified_checkpoint, + best_finalized_checkpoint, + ); // Perform a sanity check that the node is indeed valid to be the head. if !self.node_is_viable_for_head::( @@ -1186,67 +1165,30 @@ impl ProtoArray { Ok(best_node.root()) } - /// V29 virtual tree walk for `find_head`. + /// Virtual tree walk for `find_head`. /// - /// At each V29 node, determine the preferred payload direction (FULL or EMPTY) - /// by comparing weights. If `best_child` matches the preferred direction, follow - /// it directly. Otherwise, scan all nodes to find the best child matching + /// At each node, determine the preferred payload direction (FULL or EMPTY) + /// by comparing weights. Scan all nodes to find the best child matching /// the preferred direction. - fn find_head_v29_walk( + fn find_head_walk( &self, start_index: usize, current_slot: Slot, best_justified_checkpoint: Checkpoint, best_finalized_checkpoint: Checkpoint, - ) -> Result { + ) -> Result { let ptc_size = E::ptc_size(); let mut current_index = start_index; + let mut head = ForkChoiceNode { + root: best_justified_checkpoint.root, + payload_status: PayloadStatus::Pending, + }; + loop { - let node = self - .nodes - .get(current_index) - .ok_or(Error::InvalidNodeIndex(current_index))?; + let children = self.get_node_children(&head)?; - let Ok(v29) = node.as_v29() else { break }; - - let prefer_full = Self::v29_prefer_full(v29, node.slot(), current_slot, ptc_size); - let preferred_status = if prefer_full { - PayloadStatus::Full - } else { - PayloadStatus::Empty - }; - - // Fast path: check if best_child already matches the preferred direction. - let next_index = if let Some(best_child_index) = node.best_child() { - let best_child_node = self - .nodes - .get(best_child_index) - .ok_or(Error::InvalidNodeIndex(best_child_index))?; - if best_child_node - .as_v29() - .is_ok_and(|v| v.parent_payload_status == preferred_status) - { - Some(best_child_index) - } else { - // best_child is on the wrong direction. Scan for the best matching child. - self.find_best_child_with_status::( - current_index, - preferred_status, - current_slot, - best_justified_checkpoint, - best_finalized_checkpoint, - )? - } - } else { - None - }; - - if let Some(child_index) = next_index { - current_index = child_index; - } else { - break; - } + head = children.max_by_key(|child| (child.weight)) } let head_node = self @@ -1256,6 +1198,77 @@ impl ProtoArray { Ok(head_node.root()) } + fn get_weight( + &self, + fc_node: &ForkChoiceNode, + proto_node: &ProtoNode, + current_slot: Slot, + ) -> u64 { + if fc_node.payload_status == PayloadStatus::Pending + || proto_node.slot.saturating_add(1) != current_slot + { + let attestation_score = proto_node.attestation_score(); + + // TODO(gloas): implement proposer boost + // + } + } + + fn get_node_children( + node: &IndexedForkChoiceNode, + ) -> Result, Error> { + if node.payload_status == PayloadStatus::Pending { + let proto_node = self + .nodes + .get(node.node_index) + .ok_or(Error::InvalidNodeIndex(node.node_index))? + .clone(); + let mut children = vec![( + IndexedForkChoiceNode { + root: node.root, + node_index: usize, + payload_status: PayloadStatus::Empty, + }, + proto_node.clone(), + )]; + if proto_node.payload_exists().is_ok_and(|exists| exists) { + children.push(( + IndexedForkChoiceNode { + root: node.root, + node_index: usize, + payload_status: PayloadStatus::Full, + }, + proto_node, + )); + } + Ok(children) + } else { + let children = self + .nodes + .get(node.node_index..) + .ok_or(Error::InvalidNodeIndex(node.node_index))? + .iter() + .enumerate() + .filter(|(_, child_node)| { + child_node.parent_root == node.root + && node.payload_status == child_node.parent_payload_status() + }) + .map(|i, child_node| { + let child_index = node.node_index.saturating_add(i); + ( + IndexedForkChoiceNode { + root: child_node.root, + node_index: child_index, + payload_status: PayloadStatus::Pending, + }, + child_node.clone(), + ) + }) + .collect(); + Ok(children) + } + } + /// Find the best viable child of `parent_index` whose `parent_payload_status` matches /// `target_status`. Returns `None` if no matching viable child exists. fn find_best_child_with_status( diff --git a/consensus/proto_array/src/proto_array_fork_choice.rs b/consensus/proto_array/src/proto_array_fork_choice.rs index 64ec5a8549..061e8a7287 100644 --- a/consensus/proto_array/src/proto_array_fork_choice.rs +++ b/consensus/proto_array/src/proto_array_fork_choice.rs @@ -68,6 +68,13 @@ pub enum PayloadStatus { Pending = 2, } +/// Spec's `ForkChoiceNode` augmented with ProtoNode index. +pub struct IndexedForkChoiceNode { + root: Hash256, + node_index: usize, + payload_status: PayloadStatus, +} + impl ExecutionStatus { pub fn is_execution_enabled(&self) -> bool { !matches!(self, ExecutionStatus::Irrelevant(_))