Fix O(n²) find_head and stack overflow in filter_block_tree (#9090)

Co-Authored-By: dapplion <35266934+dapplion@users.noreply.github.com>

Co-Authored-By: Michael Sproul <michaelsproul@users.noreply.github.com>

consensus/proto_array/src/proto_array.rs

@@ -391,6 +391,10 @@ pub struct ProtoArray {
     pub prune_threshold: usize,
     pub nodes: Vec<ProtoNode>,
     pub indices: HashMap<Hash256, usize>,
+    /// Cached parent→children index. `children[i]` holds the node indices of all children of
+    /// node `i`. Maintained incrementally by `on_block` and `maybe_prune`.
+    #[serde(skip)]
+    pub children: Vec<Vec<usize>>,
 }
 
 impl ProtoArray {
@@ -673,6 +677,16 @@ impl ProtoArray {
         self.indices.insert(node.root(), node_index);
         self.nodes.push(node.clone());
 
+        // Maintain cached children index. `parent_index` is already bounds-checked above
+        // against `self.nodes`, and `self.children` is kept in lockstep with `self.nodes`.
+        self.children.push(Vec::new());
+        if let Some(parent_index) = node.parent() {
+            self.children
+                .get_mut(parent_index)
+                .ok_or(Error::InvalidNodeIndex(parent_index))?
+                .push(node_index);
+        }
+
         if let Some(parent_index) = node.parent()
             && matches!(block.execution_status, ExecutionStatus::Valid(_))
         {
@@ -1095,6 +1109,22 @@ impl ProtoArray {
         Ok((best_fc_node.root, best_fc_node.payload_status))
     }
 
+    /// Rebuild the cached `self.children` index from `self.nodes`. Called once after
+    /// deserialization to populate the transient field.
+    pub fn rebuild_children_index(&mut self) -> Result<(), Error> {
+        let mut children = vec![Vec::new(); self.nodes.len()];
+        for (i, node) in self.nodes.iter().enumerate() {
+            if let Some(parent_idx) = node.parent() {
+                children
+                    .get_mut(parent_idx)
+                    .ok_or(Error::InvalidNodeIndex(parent_idx))?
+                    .push(i);
+            }
+        }
+        self.children = children;
+        Ok(())
+    }
+
     /// Spec: `get_filtered_block_tree`.
     ///
     /// Returns the set of node indices on viable branches — those with at least
@@ -1105,7 +1135,7 @@ impl ProtoArray {
         current_slot: Slot,
         best_justified_checkpoint: Checkpoint,
         best_finalized_checkpoint: Checkpoint,
-    ) -> HashSet<usize> {
+    ) -> Result<HashSet<usize>, Error> {
         let mut viable = HashSet::new();
         self.filter_block_tree::<E>(
             start_index,
@@ -1113,71 +1143,88 @@ impl ProtoArray {
             best_justified_checkpoint,
             best_finalized_checkpoint,
             &mut viable,
-        );
-        viable
+        )?;
+        Ok(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).
+    ///
+    /// The spec removes execution-invalid blocks (and their entire subtrees) from
+    /// `store.blocks` before running. We replicate that here with a forward pass
+    /// propagating `excluded` from parent to child — V29 children of an invalidated
+    /// V17 ancestor are excluded transitively, since V29 nodes carry no
+    /// `execution_status` of their own.
     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;
-        };
+    ) -> Result<(), Error> {
+        // Forward pass: a node is "excluded" if it (or any ancestor down to
+        // `start_index`) has an invalid execution status.
+        let mut excluded = vec![false; self.nodes.len()];
+        for i in (start_index + 1)..self.nodes.len() {
+            let node = self.nodes.get(i).ok_or(Error::InvalidNodeIndex(i))?;
+            let parent_excluded = match node.parent() {
+                Some(p) => *excluded.get(p).ok_or(Error::InvalidNodeIndex(p))?,
+                None => false,
+            };
+            let self_invalid = node.execution_status().is_ok_and(|s| s.is_invalid());
+            excluded[i] = parent_excluded || self_invalid;
+        }
 
-        // 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();
-
-        if !children.is_empty() {
-            // Evaluate ALL children (no short-circuit) to mark all viable branches.
-            let any_viable = children
-                .iter()
-                .map(|&child_index| {
-                    self.filter_block_tree::<E>(
-                        child_index,
-                        current_slot,
-                        best_justified_checkpoint,
-                        best_finalized_checkpoint,
-                        viable,
-                    )
-                })
-                .collect::<Vec<_>>()
-                .into_iter()
-                .any(|v| v);
-            if any_viable {
-                viable.insert(node_index);
-                return true;
+        for node_index in (start_index..self.nodes.len()).rev() {
+            // Spec: invalid subtree removed from `store.blocks` — skip entirely.
+            if *excluded
+                .get(node_index)
+                .ok_or(Error::InvalidNodeIndex(node_index))?
+            {
+                continue;
             }
-            return false;
-        }
+            let node = self
+                .nodes
+                .get(node_index)
+                .ok_or(Error::InvalidNodeIndex(node_index))?;
 
-        // 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;
+            // Spec: children = [root for root in blocks if blocks[root].parent_root == block_root]
+            let valid_children: Vec<usize> = self
+                .children
+                .get(node_index)
+                .ok_or(Error::InvalidNodeIndex(node_index))?
+                .iter()
+                .copied()
+                .filter_map(|i| match excluded.get(i) {
+                    Some(false) => Some(Ok(i)),
+                    Some(true) => None,
+                    None => Some(Err(Error::InvalidNodeIndex(i))),
+                })
+                .collect::<Result<_, _>>()?;
+
+            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
+        Ok(())
     }
 
     /// Spec: `get_head`.
@@ -1204,7 +1251,7 @@ impl ProtoArray {
             current_slot,
             best_justified_checkpoint,
             best_finalized_checkpoint,
-        );
+        )?;
 
         // Compute once rather than per-child per-level.
         let apply_proposer_boost =
@@ -1468,25 +1515,35 @@ impl ProtoArray {
             }
             Ok(children)
         } else {
-            Ok(self
-                .nodes
+            // Spec: [root for root in blocks.keys() if blocks[root].parent_root == node.root ...]
+            // (cached `self.children[i]` is the same set as the spec's filtered scan).
+            let indices = self
+                .children
+                .get(node.proto_node_index)
+                .ok_or(Error::InvalidNodeIndex(node.proto_node_index))?;
+            indices
                 .iter()
-                .enumerate()
-                .filter(|(_, child_node)| {
-                    child_node.parent() == Some(node.proto_node_index)
-                        && child_node.get_parent_payload_status() == node.payload_status
+                .copied()
+                .filter_map(|i| {
+                    self.nodes
+                        .get(i)
+                        .ok_or(Error::InvalidNodeIndex(i))
+                        .map(|child| {
+                            // Spec: node.payload_status == get_parent_payload_status(store, blocks[root])
+                            (child.get_parent_payload_status() == node.payload_status).then(|| {
+                                (
+                                    IndexedForkChoiceNode {
+                                        root: child.root(),
+                                        proto_node_index: i,
+                                        payload_status: PayloadStatus::Pending,
+                                    },
+                                    child.clone(),
+                                )
+                            })
+                        })
+                        .transpose()
                 })
-                .map(|(child_index, child_node)| {
-                    (
-                        IndexedForkChoiceNode {
-                            root: child_node.root(),
-                            proto_node_index: child_index,
-                            payload_status: PayloadStatus::Pending,
-                        },
-                        child_node.clone(),
-                    )
-                })
-                .collect())
+                .collect()
         }
     }
 
@@ -1617,6 +1674,19 @@ impl ProtoArray {
         // Drop all the nodes prior to finalization.
         self.nodes = self.nodes.split_off(finalized_index);
 
+        // Drop pruned entries from children index and shift all remaining indices down.
+        // Invariant: child_index > parent_index, and all parents we kept have
+        // index >= finalized_index, so every remaining child_index is also
+        // >= finalized_index.
+        self.children = self.children.split_off(finalized_index);
+        for children in self.children.iter_mut() {
+            for child_index in children.iter_mut() {
+                *child_index = child_index
+                    .checked_sub(finalized_index)
+                    .ok_or(Error::IndexOverflow("children"))?;
+            }
+        }
+
         // Adjust the indices map.
         for (_root, index) in self.indices.iter_mut() {
             *index = index