Adds child_blocks method and refactors invalidate_block in chain.rs. …

…Refactors invalidate_block in NonFinalizedState.

zebra-state/src/service/non_finalized_state.rs

@@ -18,7 +18,7 @@ use crate::{
     constants::MAX_NON_FINALIZED_CHAIN_FORKS,
     request::{ContextuallyVerifiedBlock, FinalizableBlock},
     service::{check, finalized_state::ZebraDb},
-    HashOrHeight, SemanticallyVerifiedBlock, ValidateContextError,
+    SemanticallyVerifiedBlock, ValidateContextError,
 };
 
 mod chain;
@@ -270,45 +270,33 @@ impl NonFinalizedState {
         Ok(())
     }
 
-    /// Invalidate block with has `block_hash` and all descendants from the non-finalized state. Insert
+    /// Invalidate block with hash `block_hash` and all descendants from the non-finalized state. Insert
     /// the new chain into the chain_set and discard the previous.
     pub fn invalidate_block(&mut self, block_hash: Hash) {
-        if !self.any_chain_contains(&block_hash) {
+        let Some(chain) = self.find_chain(|chain| chain.contains_block_hash(block_hash)) else {
             return;
-        }
-
-        let mut chain = self
-            .find_chain(|chain| chain.contains_block_hash(block_hash))
-            .expect("block hash exist in a chain");
+        };
 
-        // If the non-finalized chain root has the intended hash drop the chain
-        // and return early
-        let root = chain
-            .block(HashOrHeight::Height(chain.non_finalized_root_height()))
-            .unwrap();
-        if root.hash == block_hash {
+        let invalidated_blocks = if chain.non_finalized_root_hash() == block_hash {
             self.chain_set.remove(&chain);
-            self.update_metrics_for_chains();
-            self.update_metrics_bars();
-            return;
-        }
+            chain.blocks.values().cloned().collect()
+        } else {
+            let (new_chain, invalidated_blocks) = chain
+                .invalidate_block(block_hash)
+                .expect("already checked that chain contains hash");
+
+            // Add the new chain fork or updated chain to the set of recent chains, and
+            // remove the chain containing the hash of the block from chain set
+            self.insert_with(Arc::new(new_chain.clone()), |chain_set| {
+                chain_set.retain(|c| !c.contains_block_hash(block_hash))
+            });
 
-        let new_chain = Arc::make_mut(&mut chain);
-        let block_height = new_chain.height_by_hash(block_hash).unwrap();
+            invalidated_blocks
+        };
 
-        let invalidated_descendants = new_chain.invalidate_block_and_descendants(&block_height);
         self.invalidated_blocks
-            .insert(block_hash, Arc::new(invalidated_descendants));
+            .insert(block_hash, Arc::new(invalidated_blocks));
 
-        // If the new chain still contains blocks:
-        // - add the new chain fork or updated chain to the set of recent chains
-        // - remove the chain containing the hash of the block from the chain set
-        if !new_chain.is_empty() {
-            self.insert_with(Arc::new(new_chain.clone()), |chain_set| {
-                chain_set.retain(|c| !c.contains_block_hash(block_hash))
-            });
-        }
-
         self.update_metrics_for_chains();
         self.update_metrics_bars();
     }

zebra-state/src/service/non_finalized_state/chain.rs

@@ -343,18 +343,24 @@ impl Chain {
         (block, treestate)
     }
 
-    pub fn invalidate_block_and_descendants(
-        &mut self,
-        height: &block::Height,
-    ) -> Vec<ContextuallyVerifiedBlock> {
-        // Split the new chain at the the `block_hash` and invalidate the block with the
-        // block_hash along with the block's descendants
-        let invalidated_blocks = self.blocks.split_off(height);
-        for (_, ctx_block) in invalidated_blocks.iter().rev() {
-            self.revert_chain_with(ctx_block, RevertPosition::Tip);
-        }
+    // Returns the block at the provided height and all of its descendant blocks.
+    pub fn child_blocks(&self, block_height: &block::Height) -> Vec<ContextuallyVerifiedBlock> {
+        self.blocks
+            .range(block_height..)
+            .map(|(_h, b)| b.clone())
+            .collect()
+    }
 
-        invalidated_blocks.into_values().collect()
+    // Returns a new chain without the invalidated block or its descendants.
+    pub fn invalidate_block(
+        &self,
+        block_hash: block::Hash,
+    ) -> Option<(Self, Vec<ContextuallyVerifiedBlock>)> {
+        let block_height = self.height_by_hash(block_hash)?;
+        let mut new_chain = self.fork(block_hash)?;
+        new_chain.pop_tip();
+        new_chain.last_fork_height = None;
+        Some((new_chain, self.child_blocks(&block_height)))
     }
 
     /// Returns the height of the chain root.

zebra-state/src/service/non_finalized_state/tests/vectors.rs

@@ -279,15 +279,13 @@ fn invalidate_block_removes_block_and_descendants_from_chain_for_network(
         Network::Mainnet => assert!(
             invalidated_blocks_state_descendants
                 .iter()
-                .find(|block| block.height == block::Height(653601))
-                .is_some(),
+                .any(|block| block.height == block::Height(653601)),
             "invalidated descendants vec should contain block3"
         ),
         Network::Testnet(_parameters) => assert!(
             invalidated_blocks_state_descendants
                 .iter()
-                .find(|block| block.height == block::Height(584001))
-                .is_some(),
+                .any(|block| block.height == block::Height(584001)),
             "invalidated descendants vec should contain block3"
         ),
     }