Fix NonFinalizedTree::finality_checkpoints and add caches

lib/src/chain/blocks_tree.rs

@@ -65,7 +65,12 @@ use crate::{
     header,
 };
 
-use alloc::{collections::BTreeMap, format, sync::Arc, vec::Vec};
+use alloc::{
+    collections::{BTreeMap, BTreeSet},
+    format,
+    sync::Arc,
+    vec::Vec,
+};
 use core::{cmp, fmt, mem, num::NonZeroU64, ops, time::Duration};
 use hashbrown::HashMap;
 
@@ -113,6 +118,8 @@ pub struct NonFinalizedTree<T> {
     finalized_block_header: Vec<u8>,
     /// Hash of [`NonFinalizedTree::finalized_block_header`].
     finalized_block_hash: [u8; 32],
+    /// Number of [`NonFinalizedTree::finalized_block_header`].
+    finalized_block_number: u64,
     /// State of the chain finality engine.
     finality: Finality,
     /// State of the consensus of the finalized block.
@@ -136,6 +143,10 @@ pub struct NonFinalizedTree<T> {
     /// Blocks indexed by the value in [`Block::best_score`]. The best block is the one with the
     /// highest score.
     blocks_by_best_score: BTreeMap<BestScore, fork_tree::NodeIndex>,
+    /// Subset of [`NonFinalizedTree::blocks`] whose [`BlockFinality::Grandpa::triggers_change`]
+    /// is `true`, indexed by the value in
+    /// [`BlockFinality::Grandpa::prev_auth_change_trigger_number`].
+    blocks_trigger_gp_change: BTreeSet<(Option<u64>, fork_tree::NodeIndex)>,
     /// See [`Config::block_number_bytes`].
     block_number_bytes: usize,
     /// See [`Config::allow_unknown_consensus_engines`].
@@ -153,15 +164,14 @@ impl<T> NonFinalizedTree<T> {
         let chain_information: chain_information::ChainInformation =
             config.chain_information.into();
 
-        let finalized_block_hash = chain_information
-            .finalized_block_header
-            .hash(config.block_number_bytes);
-
         NonFinalizedTree {
+            finalized_block_number: chain_information.finalized_block_header.number,
+            finalized_block_hash: chain_information
+                .finalized_block_header
+                .hash(config.block_number_bytes),
             finalized_block_header: chain_information
                 .finalized_block_header
                 .scale_encoding_vec(config.block_number_bytes),
-            finalized_block_hash,
             finality: match chain_information.finality {
                 chain_information::ChainInformationFinality::Outsourced => Finality::Outsourced,
                 chain_information::ChainInformationFinality::Grandpa {
@@ -210,6 +220,7 @@ impl<T> NonFinalizedTree<T> {
                 Default::default(),
             ),
             blocks_by_best_score: BTreeMap::new(),
+            blocks_trigger_gp_change: BTreeSet::new(),
             block_number_bytes: config.block_number_bytes,
             allow_unknown_consensus_engines: config.allow_unknown_consensus_engines,
         }
@@ -220,6 +231,7 @@ impl<T> NonFinalizedTree<T> {
         self.blocks.clear();
         self.blocks_by_hash.clear();
         self.blocks_by_best_score.clear();
+        self.blocks_trigger_gp_change.clear();
     }
 
     /// Returns true if there isn't any non-finalized block in the chain.
@@ -560,6 +572,8 @@ struct Block<T> {
     /// Cache of the hash of the block. Always equal to the hash of the header stored in this
     /// same struct.
     hash: [u8; 32],
+    /// Number contained in [`Block::header`].
+    number: u64,
     /// Changes to the consensus made by the block.
     consensus: BlockConsensus,
     /// Information about finality attached to each block.

lib/src/chain/blocks_tree/finality.rs

@@ -20,7 +20,7 @@
 use super::*;
 use crate::finality::{grandpa, justification};
 
-use core::{cmp, iter};
+use core::cmp;
 
 impl<T> NonFinalizedTree<T> {
     /// Returns a list of blocks (by their height and hash) that need to be finalized before any
@@ -31,38 +31,30 @@ impl<T> NonFinalizedTree<T> {
     /// [`NonFinalizedTree::verify_grandpa_commit_message`], unless they descend from any of the
     /// blocks returned by this function, in which case that block must be finalized beforehand.
     pub fn finality_checkpoints(&self) -> impl Iterator<Item = (u64, &[u8; 32])> {
-        match &self.finality {
-            Finality::Outsourced => {
-                // No checkpoint means all blocks allowed.
-                either::Left(iter::empty())
-            }
-            Finality::Grandpa { .. } => {
-                // TODO: O(n), could add a cache to make it O(1)
-                let iter = self
-                    .blocks
-                    .iter_unordered()
-                    .filter(move |(_, block)| {
-                        if let BlockFinality::Grandpa {
-                            triggers_change, ..
-                        } = &block.finality
-                        {
-                            *triggers_change
-                        } else {
-                            unreachable!()
-                        }
-                    })
-                    .map(|(_, block)| {
-                        (
-                            header::decode(&block.header, self.block_number_bytes)
-                                .unwrap()
-                                .number,
-                            &block.hash,
-                        )
-                    });
+        // Note that the code below assumes that GrandPa is the only finality algorithm currently
+        // supported.
+        debug_assert!(
+            self.blocks_trigger_gp_change.is_empty()
+                || !matches!(self.finality, Finality::Outsourced)
+        );
 
-                either::Right(iter)
-            }
-        }
+        self.blocks_trigger_gp_change
+            .range((
+                ops::Bound::Excluded((
+                    Some(self.finalized_block_number),
+                    fork_tree::NodeIndex::max_value(),
+                )),
+                ops::Bound::Unbounded,
+            ))
+            .map(|(_prev_auth_change_trigger_number, block_index)| {
+                debug_assert!(_prev_auth_change_trigger_number
+                    .map_or(false, |n| n > self.finalized_block_number));
+                let block = self
+                    .blocks
+                    .get(*block_index)
+                    .unwrap_or_else(|| unreachable!());
+                (block.number, &block.hash)
+            })
     }
 
     /// Verifies the given justification.
@@ -252,12 +244,7 @@ impl<T> NonFinalizedTree<T> {
                 finalized_scheduled_change,
                 finalized_triggered_authorities,
             } => {
-                let finalized_block_number =
-                    header::decode(&self.finalized_block_header, self.block_number_bytes)
-                        .unwrap()
-                        .number;
-
-                match target_number.cmp(&finalized_block_number) {
+                match target_number.cmp(&self.finalized_block_number) {
                     cmp::Ordering::Equal if *target_hash == self.finalized_block_hash => {
                         return Err(FinalityVerifyError::EqualToFinalized)
                     }
@@ -288,7 +275,7 @@ impl<T> NonFinalizedTree<T> {
                 } = self.blocks.get(block_index).unwrap().finality
                 {
                     if let Some(prev_auth_change_trigger_number) = prev_auth_change_trigger_number {
-                        if *prev_auth_change_trigger_number > finalized_block_number {
+                        if *prev_auth_change_trigger_number > self.finalized_block_number {
                             return Err(FinalityVerifyError::TooFarAhead {
                                 justification_block_number: target_number,
                                 justification_block_hash: *target_hash,
@@ -354,11 +341,7 @@ impl<T> NonFinalizedTree<T> {
                 );
                 debug_assert!(scheduled_change
                     .as_ref()
-                    .map(|(n, _)| *n
-                        > header::decode(&new_finalized_block.header, self.block_number_bytes)
-                            .unwrap()
-                            .number)
-                    .unwrap_or(true));
+                    .map_or(true, |(n, _)| *n > new_finalized_block.number));
 
                 *after_finalized_block_authorities_set_id = *after_block_authorities_set_id;
                 *finalized_triggered_authorities = triggered_authorities.clone();
@@ -422,22 +405,24 @@ impl<T> NonFinalizedTree<T> {
             _ => unreachable!(),
         }
 
-        // Update `self.finalized_block_header`, `self.finalized_block_hash`, and
-        // `self.finalized_best_score`.
+        // Update `self.finalized_block_header`, `self.finalized_block_hash`,
+        // `self.finalized_block_number`, and `self.finalized_best_score`.
         mem::swap(
             &mut self.finalized_block_header,
             &mut new_finalized_block.header,
         );
-        self.finalized_block_hash =
-            header::hash_from_scale_encoded_header(&self.finalized_block_header);
+        self.finalized_block_hash = new_finalized_block.hash;
+        self.finalized_block_number = new_finalized_block.number;
         self.finalized_best_score = new_finalized_block.best_score;
 
         debug_assert_eq!(self.blocks.len(), self.blocks_by_hash.len());
         debug_assert_eq!(self.blocks.len(), self.blocks_by_best_score.len());
+        debug_assert!(self.blocks.len() >= self.blocks_trigger_gp_change.len());
         SetFinalizedBlockIter {
             iter: self.blocks.prune_ancestors(block_index_to_finalize),
             blocks_by_hash: &mut self.blocks_by_hash,
             blocks_by_best_score: &mut self.blocks_by_best_score,
+            blocks_trigger_gp_change: &mut self.blocks_trigger_gp_change,
             updates_best_block,
         }
     }
@@ -575,6 +560,7 @@ pub struct SetFinalizedBlockIter<'a, T> {
     iter: fork_tree::PruneAncestorsIter<'a, Block<T>>,
     blocks_by_hash: &'a mut HashMap<[u8; 32], fork_tree::NodeIndex, fnv::FnvBuildHasher>,
     blocks_by_best_score: &'a mut BTreeMap<BestScore, fork_tree::NodeIndex>,
+    blocks_trigger_gp_change: &'a mut BTreeSet<(Option<u64>, fork_tree::NodeIndex)>,
     updates_best_block: bool,
 }
 
@@ -597,6 +583,17 @@ impl<'a, T> Iterator for SetFinalizedBlockIter<'a, T> {
                 .blocks_by_best_score
                 .remove(&pruned.user_data.best_score);
             debug_assert_eq!(_removed, Some(pruned.index));
+            if let BlockFinality::Grandpa {
+                prev_auth_change_trigger_number,
+                triggers_change: true,
+                ..
+            } = pruned.user_data.finality
+            {
+                let _removed = self
+                    .blocks_trigger_gp_change
+                    .remove(&(prev_auth_change_trigger_number, pruned.index));
+                debug_assert!(_removed);
+            }
             break Some(RemovedBlock {
                 block_hash: pruned.user_data.hash,
                 scale_encoded_header: pruned.user_data.header,

lib/src/chain/blocks_tree/verify.rs

@@ -455,6 +455,7 @@ impl<T> NonFinalizedTree<T> {
 
         Ok(HeaderVerifySuccess::Verified {
             verified_header: VerifiedHeader {
+                number: decoded_header.number,
                 scale_encoded_header,
                 consensus_update,
                 finality_update,
@@ -498,11 +499,23 @@ impl<T> NonFinalizedTree<T> {
             insertion_counter: self.blocks_insertion_counter,
         };
 
+        let prev_auth_change_trigger_number_if_trigger = if let BlockFinality::Grandpa {
+            prev_auth_change_trigger_number,
+            triggers_change: true,
+            ..
+        } = verified_header.finality_update
+        {
+            Some(prev_auth_change_trigger_number)
+        } else {
+            None
+        };
+
         let new_node_index = self.blocks.insert(
             parent_tree_index,
             Block {
                 header: verified_header.scale_encoded_header,
                 hash: verified_header.hash,
+                number: verified_header.number,
                 consensus: verified_header.consensus_update,
                 finality: verified_header.finality_update,
                 best_score,
@@ -518,6 +531,11 @@ impl<T> NonFinalizedTree<T> {
 
         self.blocks_by_best_score.insert(best_score, new_node_index);
 
+        if let Some(prev_auth_change_trigger_number) = prev_auth_change_trigger_number_if_trigger {
+            self.blocks_trigger_gp_change
+                .insert((prev_auth_change_trigger_number, new_node_index));
+        }
+
         // An overflow here would break the logic of the module. It is better to panic than to
         // continue running.
         self.blocks_insertion_counter = self.blocks_insertion_counter.checked_add(1).unwrap();
@@ -532,6 +550,7 @@ pub struct VerifiedHeader {
     best_score_num_primary_slots: u64,
     best_score_num_secondary_slots: u64,
     hash: [u8; 32],
+    number: u64,
 }
 
 impl VerifiedHeader {

lib/src/chain/fork_tree.rs

@@ -676,6 +676,18 @@ pub struct PrunedNode<T> {
 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
 pub struct NodeIndex(usize);
 
+impl NodeIndex {
+    /// Returns the value that compares inferior or equal to any possible [`NodeIndex`̀].
+    pub fn min_value() -> Self {
+        NodeIndex(usize::min_value())
+    }
+
+    /// Returns the value that compares superior or equal to any possible [`NodeIndex`̀].
+    pub fn max_value() -> Self {
+        NodeIndex(usize::max_value())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::ForkTree;

wasm-node/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Fixed
+
+- Justifications are no longer downloaded for blocks that can't be finalized because an earlier block needs to be finalized first. ([#1127](https://github.com/smol-dot/smoldot/pull/1127))
+
 ## 2.0.1 - 2023-09-08
 
 ### Fixed