Build children index once in find_head to fix O(n²)

Add a build_children_index() helper that creates a parent→children
mapping in O(n), then pass it through find_head_walk, filter_block_tree,
and get_node_children to replace per-node full scans with O(1) lookups.

Convert filter_block_tree from recursive to an iterative reverse pass.
Proto_array nodes are stored in insertion order (children always have
higher indices than parents), so a single reverse iteration processes
every child before its parent — matching the spec's recursive post-order
semantics without recursion or an explicit stack. This avoids stack
overflow on long chains without finality (500k+ blocks).

Adds criterion benchmarks for find_head with chain lengths up to 518k.

Author: dapplion

Cargo.lock

@@ -19,3 +19,11 @@ superstruct = { workspace = true }
 typenum = { workspace = true }
 types = { workspace = true }
 yaml_serde = { workspace = true }
+
+[dev-dependencies]
+criterion = { workspace = true }
+fixed_bytes = { workspace = true }
+
+[[bench]]
+name = "find_head"
+harness = false

consensus/proto_array/Cargo.toml

@@ -0,0 +1,142 @@
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use fixed_bytes::FixedBytesExtended;
+use proto_array::{Block, ExecutionStatus, JustifiedBalances, ProtoArrayForkChoice};
+use std::collections::BTreeSet;
+use std::time::Duration;
+use types::{
+    AttestationShufflingId, Checkpoint, Epoch, EthSpec, ExecutionBlockHash, Hash256,
+    MainnetEthSpec, Slot,
+};
+
+fn get_root(i: u64) -> Hash256 {
+    Hash256::from_low_u64_be(i)
+}
+
+fn get_hash(i: u64) -> ExecutionBlockHash {
+    ExecutionBlockHash::from_root(get_root(i))
+}
+
+/// Build a linear chain of `num_blocks` blocks.
+fn build_chain(num_blocks: u64, gloas: bool) -> (ProtoArrayForkChoice, types::ChainSpec) {
+    let mut spec = MainnetEthSpec::default_spec();
+    spec.proposer_score_boost = Some(50);
+    let gloas_fork_slot = 32;
+    if gloas {
+        spec.gloas_fork_epoch = Some(Epoch::new(1));
+    }
+
+    let finalized_checkpoint = Checkpoint {
+        epoch: Epoch::new(0),
+        root: get_root(0),
+    };
+    let junk_shuffling_id = AttestationShufflingId::from_components(Epoch::new(0), Hash256::zero());
+
+    let mut fork_choice = ProtoArrayForkChoice::new::<MainnetEthSpec>(
+        Slot::new(0),
+        Slot::new(0),
+        Hash256::zero(),
+        finalized_checkpoint,
+        finalized_checkpoint,
+        junk_shuffling_id.clone(),
+        junk_shuffling_id.clone(),
+        ExecutionStatus::Optimistic(ExecutionBlockHash::zero()),
+        None,
+        None,
+        0,
+        &spec,
+    )
+    .expect("should create fork choice");
+
+    for i in 1..=num_blocks {
+        let is_gloas = gloas && i >= gloas_fork_slot;
+        let block = Block {
+            slot: Slot::new(i),
+            root: get_root(i),
+            parent_root: Some(get_root(i - 1)),
+            state_root: Hash256::zero(),
+            target_root: get_root(0),
+            current_epoch_shuffling_id: junk_shuffling_id.clone(),
+            next_epoch_shuffling_id: junk_shuffling_id.clone(),
+            justified_checkpoint: finalized_checkpoint,
+            finalized_checkpoint,
+            execution_status: ExecutionStatus::Optimistic(ExecutionBlockHash::zero()),
+            unrealized_justified_checkpoint: Some(finalized_checkpoint),
+            unrealized_finalized_checkpoint: Some(finalized_checkpoint),
+            execution_payload_parent_hash: if is_gloas {
+                Some(get_hash(i - 1))
+            } else {
+                None
+            },
+            execution_payload_block_hash: if is_gloas { Some(get_hash(i)) } else { None },
+            proposer_index: Some(0),
+        };
+
+        fork_choice
+            .process_block::<MainnetEthSpec>(block, Slot::new(i), &spec, Duration::ZERO)
+            .expect("should process block");
+    }
+
+    (fork_choice, spec)
+}
+
+fn bench_find_head(c: &mut Criterion) {
+    let mut group = c.benchmark_group("find_head");
+    let equivocating_indices = BTreeSet::new();
+
+    // Must survive extended non-finality (500k+ blocks).
+    for &num_blocks in &[100, 1_000, 10_000, 50_000, 216_000, 518_000] {
+        let (mut fork_choice, spec) = build_chain(num_blocks, false);
+        let finalized_checkpoint = Checkpoint {
+            epoch: Epoch::new(0),
+            root: get_root(0),
+        };
+        let balances = JustifiedBalances::from_effective_balances(vec![1; 64]).unwrap();
+
+        group.bench_function(BenchmarkId::new("pre_gloas", num_blocks), |b| {
+            b.iter(|| {
+                fork_choice
+                    .find_head::<MainnetEthSpec>(
+                        finalized_checkpoint,
+                        finalized_checkpoint,
+                        &balances,
+                        Hash256::zero(),
+                        &equivocating_indices,
+                        Slot::new(num_blocks),
+                        &spec,
+                    )
+                    .expect("should find head")
+            });
+        });
+    }
+
+    // 216k = ~1 month non-finality mainnet, 518k = ~1 month non-finality Gnosis
+    for &num_blocks in &[100, 1_000, 10_000, 50_000, 216_000, 518_000] {
+        let (mut fork_choice, spec) = build_chain(num_blocks, true);
+        let finalized_checkpoint = Checkpoint {
+            epoch: Epoch::new(0),
+            root: get_root(0),
+        };
+        let balances = JustifiedBalances::from_effective_balances(vec![1; 64]).unwrap();
+
+        group.bench_function(BenchmarkId::new("gloas", num_blocks), |b| {
+            b.iter(|| {
+                fork_choice
+                    .find_head::<MainnetEthSpec>(
+                        finalized_checkpoint,
+                        finalized_checkpoint,
+                        &balances,
+                        Hash256::zero(),
+                        &equivocating_indices,
+                        Slot::new(num_blocks),
+                        &spec,
+                    )
+                    .expect("should find head")
+            });
+        });
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_find_head);
+criterion_main!(benches);

consensus/proto_array/benches/find_head.rs

@@ -1108,6 +1108,20 @@ impl ProtoArray {
         Ok((best_fc_node.root, best_fc_node.payload_status))
     }
 
+    /// Build parent→children index in O(n). Returns a vec where `result[i]` contains
+    /// the indices of all children of node `i`.
+    fn build_children_index(&self) -> Vec<Vec<usize>> {
+        let mut children = vec![Vec::new(); self.nodes.len()];
+        for (i, node) in self.nodes.iter().enumerate() {
+            if let Some(parent_idx) = node.parent()
+                && parent_idx < children.len()
+            {
+                children[parent_idx].push(i);
+            }
+        }
+        children
+    }
+
     /// Spec: `get_filtered_block_tree`.
     ///
     /// Returns the set of node indices on viable branches — those with at least
@@ -1118,6 +1132,7 @@ impl ProtoArray {
         current_slot: Slot,
         best_justified_checkpoint: Checkpoint,
         best_finalized_checkpoint: Checkpoint,
+        children_index: &[Vec<usize>],
     ) -> HashSet<usize> {
         let mut viable = HashSet::new();
         self.filter_block_tree::<E>(
@@ -1126,71 +1141,66 @@ impl ProtoArray {
             best_justified_checkpoint,
             best_finalized_checkpoint,
             &mut viable,
+            children_index,
         );
         viable
     }
 
     /// Spec: `filter_block_tree`.
+    ///
+    /// Proto_array stores nodes in insertion order — children always have higher
+    /// indices than their parents. A single reverse pass therefore processes every
+    /// child before its parent, matching the spec's recursive post-order semantics
+    /// without recursion (required to survive 500k+ blocks of non-finality).
     fn filter_block_tree<E: EthSpec>(
         &self,
-        node_index: usize,
+        start_index: usize,
         current_slot: Slot,
         best_justified_checkpoint: Checkpoint,
         best_finalized_checkpoint: Checkpoint,
         viable: &mut HashSet<usize>,
-    ) -> bool {
-        let Some(node) = self.nodes.get(node_index) else {
-            return false;
-        };
-
-        // Skip invalid children — they aren't in store.blocks in the spec.
-        let children: Vec<usize> = self
-            .nodes
-            .iter()
-            .enumerate()
-            .filter(|(_, child)| {
-                child.parent() == Some(node_index)
-                    && !child
-                        .execution_status()
-                        .is_ok_and(|status| status.is_invalid())
-            })
-            .map(|(i, _)| i)
-            .collect();
+        children_index: &[Vec<usize>],
+    ) {
+        for node_index in (start_index..self.nodes.len()).rev() {
+            let Some(node) = self.nodes.get(node_index) else {
+                continue;
+            };
 
-        if !children.is_empty() {
-            // Evaluate ALL children (no short-circuit) to mark all viable branches.
-            let any_viable = children
+            // Spec: children = [root for root in blocks if blocks[root].parent_root == block_root]
+            // Skip execution-invalid children (not in store.blocks in the spec).
+            let children = children_index
+                .get(node_index)
+                .map(|c| c.as_slice())
+                .unwrap_or(&[]);
+            let valid_children: Vec<usize> = children
                 .iter()
-                .map(|&child_index| {
-                    self.filter_block_tree::<E>(
-                        child_index,
-                        current_slot,
-                        best_justified_checkpoint,
-                        best_finalized_checkpoint,
-                        viable,
-                    )
+                .copied()
+                .filter(|&i| {
+                    !self.nodes.get(i).is_some_and(|child| {
+                        child
+                            .execution_status()
+                            .is_ok_and(|status| status.is_invalid())
+                    })
                 })
-                .collect::<Vec<_>>()
-                .into_iter()
-                .any(|v| v);
-            if any_viable {
-                viable.insert(node_index);
-                return true;
-            }
-            return false;
-        }
+                .collect();
 
-        // Leaf node: check viability.
-        if self.node_is_viable_for_head::<E>(
-            node,
-            current_slot,
-            best_justified_checkpoint,
-            best_finalized_checkpoint,
-        ) {
-            viable.insert(node_index);
-            return true;
+            if !valid_children.is_empty() {
+                // Spec: if any(children): if any(filter_block_tree_result): blocks[block_root] = block
+                if valid_children.iter().any(|c| viable.contains(c)) {
+                    viable.insert(node_index);
+                }
+            } else {
+                // Spec: leaf — check correct_justified and correct_finalized
+                if self.node_is_viable_for_head::<E>(
+                    node,
+                    current_slot,
+                    best_justified_checkpoint,
+                    best_finalized_checkpoint,
+                ) {
+                    viable.insert(node_index);
+                }
+            }
         }
-        false
     }
 
     /// Spec: `get_head`.
@@ -1211,12 +1221,15 @@ impl ProtoArray {
             payload_status: PayloadStatus::Pending,
         };
 
+        let children_index = self.build_children_index();
+
         // Spec: `get_filtered_block_tree`.
         let viable_nodes = self.get_filtered_block_tree::<E>(
             start_index,
             current_slot,
             best_justified_checkpoint,
             best_finalized_checkpoint,
+            &children_index,
         );
 
         // Compute once rather than per-child per-level.
@@ -1225,7 +1238,7 @@ impl ProtoArray {
 
         loop {
             let children: Vec<_> = self
-                .get_node_children(&head)?
+                .get_node_children(&head, &children_index)?
                 .into_iter()
                 .filter(|(fc_node, _)| viable_nodes.contains(&fc_node.proto_node_index))
                 .collect();
@@ -1384,6 +1397,7 @@ impl ProtoArray {
     fn get_node_children(
         &self,
         node: &IndexedForkChoiceNode,
+        children_index: &[Vec<usize>],
     ) -> Result<Vec<(IndexedForkChoiceNode, ProtoNode)>, Error> {
         if node.payload_status == PayloadStatus::Pending {
             let proto_node = self
@@ -1397,25 +1411,29 @@ impl ProtoArray {
             }
             Ok(children)
         } else {
-            Ok(self
-                .nodes
-                .iter()
-                .enumerate()
-                .filter(|(_, child_node)| {
-                    child_node.parent() == Some(node.proto_node_index)
-                        && child_node.get_parent_payload_status() == node.payload_status
-                })
-                .map(|(child_index, child_node)| {
-                    (
-                        IndexedForkChoiceNode {
-                            root: child_node.root(),
-                            proto_node_index: child_index,
-                            payload_status: PayloadStatus::Pending,
-                        },
-                        child_node.clone(),
-                    )
+            Ok(children_index
+                .get(node.proto_node_index)
+                .map(|indices| {
+                    indices
+                        .iter()
+                        .filter_map(|&child_index| {
+                            let child_node = self.nodes.get(child_index)?;
+                            if child_node.get_parent_payload_status() == node.payload_status {
+                                Some((
+                                    IndexedForkChoiceNode {
+                                        root: child_node.root(),
+                                        proto_node_index: child_index,
+                                        payload_status: PayloadStatus::Pending,
+                                    },
+                                    child_node.clone(),
+                                ))
+                            } else {
+                                None
+                            }
+                        })
+                        .collect()
                 })
-                .collect())
+                .unwrap_or_default())
         }
     }