// Smoldot
// Copyright (C) 2019-2022  Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

//! Warp syncing.
//!
//! Warp syncing makes it possible to very quickly reach the currently finalized block of a chain.
//! No attempt is made at verifying blocks. The algorithm assumes that the currently finalized
//! block is valid, as the chain is likely bricked if this is not the case.
//!
//! # Long range attack vulnerability
//!
//! Warp syncing is particularly vulnerable to what is called long range attacks.
//!
//! The authorities allowed to finalize blocks can generate multiple proofs of finality for
//! multiple different blocks of the same height. In other words, they can finalize more than one
//! chain at a time.
//!
//! Finalizing multiple chains at the same time is called an equivocation. Client implementations
//! detect equivocations and report them to the runtime. While this is out of scope of the client,
//! the runtime then typically provides a mechanism that punishes the validators that have
//! equivocated.
//! However, this defense mechanism is flawed in case when a long time passes between the
//! multiple finality proofs generated by the same validators. Clients cannot hold an unlimited
//! amount of information in memory, so they might not detect the equivocation, and even if it is
//! detected, the punishment might not be enforceable because validators have moved all their
//! funds.
//!
//! In other words, it is possible for two thirds of the validators that were active at a certain
//! past block N to collude and decide to finalize a different block N, even when N has been
//! finalized for the first time several weeks or months in the past. When a client then warp
//! syncs, it can be tricked to consider this alternative block N as the finalized one.
//!
//! There is no fool-proof defense against this attack. However, consider the extremely high
//! investment and high risk for the malicious validators, and the difficulty of pulling off this
//! attack, it is extremely unlikely to happen in reality.
//! The aforementioned punishment system is the only defense against this attack, and in order to
//! be effective, the starting point of the warp syncing shouldn't be too far in the past. How
//! far exactly depends on the logic of the runtime of the chain.
//!
//! # Overview
//!
//! The warp syncing algorithm works only if the chain uses Grandpa for its finality.
//! It consists in the following steps:
//!
//! - Downloading a warp sync proof from a source. This proof contains a list of *fragments*. Each
//! fragment represents a change in the list of Grandpa authorities, and a list of signatures of
//! the previous authorities that certify that this change is correct.
//! - Verifying the fragments. Each fragment that is successfully verified progresses towards
//! the head of the chain. Even if one fragment is invalid, all the previously-verified
//! fragments can still be kept, and the warp syncing can resume from there.
//! - Downloading from a source the runtime code of the final block of the proof.
//! - Performing some runtime calls in order to obtain the current consensus-related parameters
//! of the chain. This might require obtaining some storage items, in which case they must also
//! be downloaded from a source.
//!
//! At the end of the syncing, a [`ValidChainInformation`] corresponding to the head of the chain
//! is yielded.
//!
//! # Usage
//!
//! Use the [`start_warp_sync()`] function to start a Grandpa warp syncing state machine.
//!
//! At any given moment, this state machine holds a list of *sources* that it might use to
//! download the warp sync proof or the runtime code. Sources must be added and removed by the API
//! user by calling one of the various `add_source` and `remove_source` functions.
//!
//! Sources are identified through a [`SourceId`]. Each source has an opaque so-called "user data"
//! of type `TSrc` associated to it. The content of this "user data" is at the discretion of the
//! API user.
//!
//! Similarly, at any given moment, this state machine holds a list of requests that concern these
//! sources. Use [`WarpSync::desired_requests`] to determine which requests will be useful to the
//! progress of the warp syncing, then use [`WarpSync::add_request`] to update the state machine
//! with a newly-started request.
//!
//! Use [`WarpSync::process_one`] in order to run verifications of the payloads that have
//! previously been downloaded.
//!

use crate::{
    chain::chain_information::{
        self, ChainInformationConsensusRef, ChainInformationFinality, ChainInformationFinalityRef,
        ValidChainInformation, ValidChainInformationRef,
    },
    executor::{
        self,
        host::{self, HostVmPrototype},
        vm::ExecHint,
    },
    finality::{decode, verify},
    header,
    informant::HashDisplay,
    trie::{self, proof_decode},
};

use alloc::{
    borrow::{Cow, ToOwned as _},
    collections::{BTreeSet, VecDeque},
    vec,
    vec::Vec,
};
use core::{cmp, fmt, iter, mem, ops};

pub use trie::Nibble;

/// The configuration for [`start_warp_sync()`].
#[derive(Debug)]
pub struct Config {
    /// The chain information of the starting point of the warp syncing.
    pub start_chain_information: ValidChainInformation,

    /// Number of bytes used when encoding/decoding the block number. Influences how various data
    /// structures should be parsed.
    pub block_number_bytes: usize,

    /// The initial capacity of the list of sources.
    pub sources_capacity: usize,

    /// The initial capacity of the list of requests.
    pub requests_capacity: usize,

    /// Known valid Merkle value and storage value combination for the `:code` key.
    ///
    /// If provided, the warp syncing algorithm will first fetch the Merkle value of `:code`, and
    /// if it matches the Merkle value provided in the hint, use the storage value in the hint
    /// instead of downloading it. If the hint doesn't match, an extra round-trip will be needed,
    /// but if the hint matches it saves a big download.
    pub code_trie_node_hint: Option<ConfigCodeTrieNodeHint>,

    /// Number of warp sync fragments after which the state machine will pause downloading new
    /// ones until the ones that have been downloaded are verified.
    ///
    /// A too low value will cause stalls, while a high value will use more memory and runs the
    /// risk of wasting more bandwidth in case the downloaded fragments need to be thrown away.
    pub num_download_ahead_fragments: usize,

    /// If the height of the current local finalized block is `N`, the warp sync state machine
    /// will not attempt to warp sync to blocks whose height inferior or equal to `N + k` where
    /// `k` is the value in this field.
    ///
    /// Because warp syncing is a relatively expensive process, it is not worth performing it
    /// between two blocks that are too close to each other.
    ///
    /// The ideal value of this field depends on the block production rate and the time it takes
    /// to answer requests.
    pub warp_sync_minimum_gap: usize,

    /// If `true`, the body of the warp sync target will be downloaded before the warp sync
    /// finishes.
    /// Determines whether [`RuntimeInformation::finalized_body`] is `Some`.
    pub download_block_body: bool,

    /// If `true`, all the storage proofs and call proofs necessary in order to compute the chain
    /// information of the warp synced block will be downloaded. If `false`, the finality
    /// information of the warp synced block is inferred from the warp sync fragments instead.
    pub download_all_chain_information_storage_proofs: bool,
}

/// See [`Config::code_trie_node_hint`].
#[derive(Debug)]
pub struct ConfigCodeTrieNodeHint {
    /// Potential Merkle value of the `:code` key.
    pub merkle_value: Vec<u8>,

    /// Storage value corresponding to [`ConfigCodeTrieNodeHint::merkle_value`].
    pub storage_value: Vec<u8>,

    /// Closest ancestor of the `:code` key except for `:code` itself.
    pub closest_ancestor_excluding: Vec<Nibble>,
}

/// Initializes the warp sync state machine.
///
/// On error, returns the [`ValidChainInformation`] that was provided in the configuration.
pub fn start_warp_sync<TSrc, TRq>(
    config: Config,
) -> Result<WarpSync<TSrc, TRq>, (ValidChainInformation, WarpSyncInitError)> {
    match config.start_chain_information.as_ref().finality {
        // TODO: we make sure that `finalized_scheduled_change` is `None` because it seems complicated to support, but ideally it would be supported
        ChainInformationFinalityRef::Grandpa {
            finalized_scheduled_change: None,
            ..
        } => {}
        _ => {
            return Err((
                config.start_chain_information,
                WarpSyncInitError::NotGrandpa,
            ));
        }
    }

    match config.start_chain_information.as_ref().consensus {
        ChainInformationConsensusRef::Babe { .. } | ChainInformationConsensusRef::Aura { .. } => {}
        ChainInformationConsensusRef::Unknown => {
            return Err((
                config.start_chain_information,
                WarpSyncInitError::UnknownConsensus,
            ));
        }
    }

    let warped_header = config
        .start_chain_information
        .as_ref()
        .finalized_block_header
        .scale_encoding_vec(config.block_number_bytes);

    Ok(WarpSync {
        warped_header_number: config
            .start_chain_information
            .as_ref()
            .finalized_block_header
            .number,
        warped_header_state_root: *config
            .start_chain_information
            .as_ref()
            .finalized_block_header
            .state_root,
        warped_header_extrinsics_root: *config
            .start_chain_information
            .as_ref()
            .finalized_block_header
            .extrinsics_root,
        warped_header_hash: header::hash_from_scale_encoded_header(&warped_header),
        warped_header,
        warped_finality: config.start_chain_information.as_ref().finality.into(),
        warped_block_ty: WarpedBlockTy::AlreadyVerified,
        runtime_calls: runtime_calls_default_value(
            config.start_chain_information.as_ref().consensus,
        ),
        verified_chain_information: config.start_chain_information,
        code_trie_node_hint: config.code_trie_node_hint,
        num_download_ahead_fragments: config.num_download_ahead_fragments,
        warp_sync_minimum_gap: config.warp_sync_minimum_gap,
        block_number_bytes: config.block_number_bytes,
        download_all_chain_information_storage_proofs: config
            .download_all_chain_information_storage_proofs,
        sources: slab::Slab::with_capacity(config.sources_capacity),
        sources_by_finalized_height: BTreeSet::new(),
        in_progress_requests: slab::Slab::with_capacity(config.requests_capacity),
        in_progress_requests_by_source: BTreeSet::new(),
        warp_sync_fragments_download: None,
        verify_queue: VecDeque::new(),
        runtime_download: RuntimeDownload::NotStarted {
            hint_doesnt_match: false,
        },
        body_download: if config.download_block_body {
            BodyDownload::NotStarted
        } else {
            BodyDownload::NotNeeded
        },
    })
}

/// Error potentially returned by [`start_warp_sync()`].
#[derive(Debug, derive_more::Display, derive_more::Error, Clone)]
pub enum WarpSyncInitError {
    /// Chain doesn't use the Grandpa finality algorithm.
    NotGrandpa,
    /// Chain uses an unrecognized consensus mechanism.
    UnknownConsensus,
}

/// Identifier for a source in the [`WarpSync`].
//
// Implementation note: this represents the index within the `Slab` used for the list of sources.
#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub struct SourceId(usize);

impl SourceId {
    /// Smallest possible [`SourceId`]. It is always inferior or equal to any other.
    pub const MIN: Self = SourceId(usize::MIN);
}

// TODO: consider removing this entirely
pub struct RuntimeInformation {
    /// The runtime constructed in `VirtualMachineParamsGet`. Corresponds to the runtime of the
    /// finalized block of [`WarpSync::as_chain_information`].
    pub finalized_runtime: HostVmPrototype,

    /// List of SCALE-encoded extrinsics of the body of the finalized block.
    /// `Some` if and only if [`Config::download_block_body`] was `true`.
    pub finalized_body: Option<Vec<Vec<u8>>>,

    /// Storage value at the `:code` key of the finalized block.
    pub finalized_storage_code: Option<Vec<u8>>,

    /// Storage value at the `:heappages` key of the finalized block.
    pub finalized_storage_heap_pages: Option<Vec<u8>>,

    /// Merkle value of the `:code` trie node of the finalized block.
    pub finalized_storage_code_merkle_value: Option<Vec<u8>>,

    /// Closest ancestor of the `:code` trie node of the finalized block excluding `:code` itself.
    pub finalized_storage_code_closest_ancestor_excluding: Option<Vec<Nibble>>,
}

/// Fragment to be verified.
#[derive(Debug)]
pub struct WarpSyncFragment {
    /// Header of a block in the chain.
    pub scale_encoded_header: Vec<u8>,

    /// Justification that proves the finality of [`WarpSyncFragment::scale_encoded_header`].
    pub scale_encoded_justification: Vec<u8>,
}

/// Warp syncing process state machine.
pub struct WarpSync<TSrc, TRq> {
    /// SCALE-encoded header of the finalized block of the chain we warp synced to. Initially
    /// identical to the value in [`WarpSync::verified_chain_information`].
    warped_header: Vec<u8>,
    /// Hash of the block in [`WarpSync::warped_header`].
    warped_header_hash: [u8; 32],
    /// State trie root hash of the block in [`WarpSync::warped_header`].
    warped_header_state_root: [u8; 32],
    /// Extrinsics trie root hash of the block in [`WarpSync::warped_header`].
    warped_header_extrinsics_root: [u8; 32],
    /// Number of the block in [`WarpSync::warped_header`].
    warped_header_number: u64,
    /// Information about the finality of the chain at the point where we warp synced to.
    /// Initially identical to the value in [`WarpSync::verified_chain_information`].
    warped_finality: ChainInformationFinality,
    /// Information about the block described by [`WarpSync::warped_header`] and
    /// [`WarpSync::warped_finality`].
    warped_block_ty: WarpedBlockTy,
    /// See [`Config::code_trie_node_hint`].
    code_trie_node_hint: Option<ConfigCodeTrieNodeHint>,
    /// Starting point of the warp syncing as provided to [`start_warp_sync`], or latest chain
    /// information that was warp synced to.
    verified_chain_information: ValidChainInformation,
    /// See [`Config::num_download_ahead_fragments`].
    num_download_ahead_fragments: usize,
    /// See [`Config::warp_sync_minimum_gap`].
    warp_sync_minimum_gap: usize,
    /// See [`Config::block_number_bytes`].
    block_number_bytes: usize,
    /// See [`Config::download_all_chain_information_storage_proofs`].
    download_all_chain_information_storage_proofs: bool,
    /// List of requests that have been added using [`WarpSync::add_source`].
    sources: slab::Slab<Source<TSrc>>,
    /// Subset of the entries as [`WarpSync::sources`] whose [`Source::finalized_block_height`]
    /// is `Ok`. Indexed by [`Source::finalized_block_height`].
    sources_by_finalized_height: BTreeSet<(u64, SourceId)>,
    /// List of requests that have been added using [`WarpSync::add_request`].
    in_progress_requests: slab::Slab<(SourceId, TRq, RequestDetail)>,
    /// Identical to [`WarpSync::in_progress_requests`], but indexed differently.
    in_progress_requests_by_source: BTreeSet<(SourceId, RequestId)>,
    /// Request that is downloading warp sync fragments, if any has been started yet.
    warp_sync_fragments_download: Option<RequestId>,
    /// Queue of fragments that have been downloaded and need to be verified.
    verify_queue: VecDeque<PendingVerify>,
    /// State of the download of the runtime and chain information call proofs.
    runtime_download: RuntimeDownload,
    /// State of the download of the body of the block.
    body_download: BodyDownload,
    /// For each call required by the chain information builder, whether it has been downloaded yet.
    runtime_calls:
        hashbrown::HashMap<chain_information::build::RuntimeCall, CallProof, fnv::FnvBuildHasher>,
}

/// See [`WarpSync::sources`].
#[derive(Debug, Copy, Clone)]
struct Source<TSrc> {
    /// User data chosen by the API user.
    user_data: TSrc,
    /// Height of the finalized block of the source, as reported by the source.
    finalized_block_height: u64,
}

/// See [`WarpSync::warped_block_ty`].
enum WarpedBlockTy {
    /// Block is equal to the finalized block in [`WarpSync::verified_chain_information`].
    AlreadyVerified,
    /// Block is known to not be warp-syncable due to an incompatibility between smoldot and
    /// the chain.
    KnownBad,
    /// Block is expected to be warp syncable.
    Normal,
}

/// See [`WarpSync::runtime_download`].
enum RuntimeDownload {
    NotStarted {
        hint_doesnt_match: bool,
    },
    Downloading {
        hint_doesnt_match: bool,
        request_id: RequestId,
    },
    NotVerified {
        /// Source the runtime has been obtained from. `None` if the source has been removed.
        downloaded_source: Option<SourceId>,
        hint_doesnt_match: bool,
        trie_proof: Vec<u8>,
    },
    Verified {
        downloaded_runtime: DownloadedRuntime,
        chain_info_builder: chain_information::build::ChainInformationBuild,
    },
}

/// See [`WarpSync::body_download`].
enum BodyDownload {
    NotNeeded,
    NotStarted,
    Downloading {
        request_id: RequestId,
    },
    Downloaded {
        /// Source the body has been obtained from. `None` if the source has been removed.
        downloaded_source: Option<SourceId>,
        body: Vec<Vec<u8>>,
    },
}

/// See [`WarpSync::verify_queue`].
struct PendingVerify {
    /// Source the fragments have been obtained from. `None` if the source has been removed.
    downloaded_source: Option<SourceId>,
    /// `true` if the source has indicated that there is no more fragment afterwards, in other
    /// words that the last fragment corresponds to the current finalized block of the chain.
    final_set_of_fragments: bool,
    /// List of fragments to verify. Can be empty.
    fragments: Vec<WarpSyncFragment>,
    /// Number of fragments at the start of [`PendingVerify::fragments`] that have already been
    /// verified. Must always be strictly inferior to `fragments.len()`, unless the list of
    /// fragments is empty.
    next_fragment_to_verify_index: usize,
}

/// See [`RuntimeDownload::Verified`].
struct DownloadedRuntime {
    /// Storage item at the `:code` key. `None` if there is no entry at that key.
    storage_code: Option<Vec<u8>>,
    /// Storage item at the `:heappages` key. `None` if there is no entry at that key.
    storage_heap_pages: Option<Vec<u8>>,
    /// Merkle value of the `:code` trie node. `None` if there is no entry at that key.
    code_merkle_value: Option<Vec<u8>>,
    /// Closest ancestor of the `:code` key except for `:code` itself.
    closest_ancestor_excluding: Option<Vec<Nibble>>,
}

/// See [`WarpSync::runtime_calls`].
enum CallProof {
    NotStarted,
    Downloading(RequestId),
    Downloaded {
        /// Source the proof has been obtained from. `None` if the source has been removed.
        downloaded_source: Option<SourceId>,
        proof: Vec<u8>,
    },
}

/// Returns the default value for [`WarpSync::runtime_calls`].
///
/// Contains the list of calls that we anticipate the chain information builder will make. This
/// assumes that the runtime is the latest version available.
fn runtime_calls_default_value(
    verified_chain_information_consensus: chain_information::ChainInformationConsensusRef,
) -> hashbrown::HashMap<chain_information::build::RuntimeCall, CallProof, fnv::FnvBuildHasher> {
    let mut list = hashbrown::HashMap::with_capacity_and_hasher(8, Default::default());
    match verified_chain_information_consensus {
        ChainInformationConsensusRef::Aura { .. } => {
            list.insert(
                chain_information::build::RuntimeCall::AuraApiAuthorities,
                CallProof::NotStarted,
            );
            list.insert(
                chain_information::build::RuntimeCall::AuraApiSlotDuration,
                CallProof::NotStarted,
            );
        }
        ChainInformationConsensusRef::Babe { .. } => {
            list.insert(
                chain_information::build::RuntimeCall::BabeApiCurrentEpoch,
                CallProof::NotStarted,
            );
            list.insert(
                chain_information::build::RuntimeCall::BabeApiNextEpoch,
                CallProof::NotStarted,
            );
            list.insert(
                chain_information::build::RuntimeCall::BabeApiConfiguration,
                CallProof::NotStarted,
            );
        }
        ChainInformationConsensusRef::Unknown => {}
    }
    list
}

/// See [`WarpSync::status`].
#[derive(Debug)]
pub enum Status<'a, TSrc> {
    /// Warp syncing algorithm is downloading Grandpa warp sync fragments containing a finality
    /// proof.
    Fragments {
        /// Source from which the fragments are currently being downloaded, if any.
        source: Option<(SourceId, &'a TSrc)>,
        /// Hash of the highest block that is proven to be finalized.
        ///
        /// This isn't necessarily the same block as returned by
        /// [`WarpSync::as_chain_information`], as this function first has to download
        /// extra information compared to just the finalized block.
        finalized_block_hash: [u8; 32],
        /// Height of the block indicated by [`Status::ChainInformation::finalized_block_hash`].
        finalized_block_number: u64,
    },
    /// Warp syncing algorithm has reached the head of the finalized chain and is downloading and
    /// building the chain information.
    ChainInformation {
        /// Hash of the highest block that is proven to be finalized.
        ///
        /// This isn't necessarily the same block as returned by
        /// [`WarpSync::as_chain_information`], as this function first has to download
        /// extra information compared to just the finalized block.
        finalized_block_hash: [u8; 32],
        /// Height of the block indicated by [`Status::ChainInformation::finalized_block_hash`].
        finalized_block_number: u64,
    },
}

impl<TSrc, TRq> WarpSync<TSrc, TRq> {
    /// Returns the value that was initially passed in [`Config::block_number_bytes`].
    pub fn block_number_bytes(&self) -> usize {
        self.block_number_bytes
    }

    /// Returns the chain information that is considered verified.
    pub fn as_chain_information(&'_ self) -> ValidChainInformationRef<'_> {
        // Note: after verifying a warp sync fragment, we are certain that the header targeted by
        // this fragment is indeed part of the chain. However, this is not enough in order to
        // produce a full chain information struct. Such struct can only be produced after the
        // entire warp syncing has succeeded. If if it still in progress, all we can return is
        // the starting point.
        (&self.verified_chain_information).into()
    }

    /// Modifies the chain information known to be valid.
    pub fn set_chain_information(&mut self, chain_information: ValidChainInformationRef) {
        // The implementation simply resets the state machine, apart from the fragment download
        // and verification queue.
        // TODO: what if the new chain doesn't support grandpa?
        if self.warped_header_number <= chain_information.as_ref().finalized_block_header.number {
            self.warped_header = chain_information
                .as_ref()
                .finalized_block_header
                .scale_encoding_vec(self.block_number_bytes);
            self.warped_header_hash = chain_information
                .as_ref()
                .finalized_block_header
                .hash(self.block_number_bytes);
            self.warped_header_state_root =
                *chain_information.as_ref().finalized_block_header.state_root;
            self.warped_header_extrinsics_root = *chain_information
                .as_ref()
                .finalized_block_header
                .extrinsics_root;
            self.warped_header_number = chain_information.as_ref().finalized_block_header.number;
            self.warped_finality = chain_information.as_ref().finality.into();
            self.warped_block_ty = WarpedBlockTy::AlreadyVerified;

            self.verified_chain_information = chain_information.into();
            self.runtime_calls =
                runtime_calls_default_value(self.verified_chain_information.as_ref().consensus);

            self.runtime_download = RuntimeDownload::NotStarted {
                hint_doesnt_match: false,
            };

            if !matches!(self.body_download, BodyDownload::NotNeeded) {
                self.body_download = BodyDownload::NotStarted;
            }
        }
    }

    /// Returns the current status of the warp syncing.
    pub fn status(&'_ self) -> Status<'_, TSrc> {
        match &self.runtime_download {
            RuntimeDownload::NotStarted { .. } => {
                let finalized_block_hash = self.warped_header_hash;

                let source_id =
                    if let Some(warp_sync_fragments_download) = self.warp_sync_fragments_download {
                        Some(
                            self.in_progress_requests
                                .get(warp_sync_fragments_download.0)
                                .unwrap()
                                .0,
                        )
                    } else {
                        self.verify_queue.back().and_then(|f| f.downloaded_source)
                    };

                Status::Fragments {
                    source: source_id.map(|id| (id, &self.sources[id.0].user_data)),
                    finalized_block_hash,
                    finalized_block_number: self.warped_header_number,
                }
            }
            _ => Status::ChainInformation {
                finalized_block_hash: self.warped_header_hash,
                finalized_block_number: self.warped_header_number,
            },
        }
    }

    /// Returns a list of all known sources stored in the state machine.
    pub fn sources(&self) -> impl Iterator<Item = SourceId> {
        self.sources.iter().map(|(id, _)| SourceId(id))
    }

    /// Returns the number of ongoing requests that concern this source.
    ///
    /// # Panic
    ///
    /// Panics if the [`SourceId`] is invalid.
    ///
    pub fn source_num_ongoing_requests(&self, source_id: SourceId) -> usize {
        assert!(self.sources.contains(source_id.0));
        self.in_progress_requests_by_source
            .range((source_id, RequestId(usize::MIN))..=(source_id, RequestId(usize::MAX)))
            .count()
    }

    /// Add a source to the list of sources.
    ///
    /// The source has a finalized block height of 0, which should later be updated using
    /// [`WarpSync::set_source_finality_state`].
    pub fn add_source(&mut self, user_data: TSrc) -> SourceId {
        let source_id = SourceId(self.sources.insert(Source {
            user_data,
            finalized_block_height: 0,
        }));

        let _inserted = self.sources_by_finalized_height.insert((0, source_id));
        debug_assert!(_inserted);
        debug_assert!(self.sources.len() >= self.sources_by_finalized_height.len());

        source_id
    }

    /// Removes a source from the list of sources. In addition to the user data associated to this
    /// source, also returns a list of requests that were in progress concerning this source. These
    /// requests are now considered obsolete.
    ///
    /// # Panic
    ///
    /// Panics if the [`SourceId`] is invalid.
    ///
    pub fn remove_source(
        &mut self,
        to_remove: SourceId,
    ) -> (TSrc, impl Iterator<Item = (RequestId, TRq)>) {
        debug_assert!(self.sources.contains(to_remove.0));
        let removed = self.sources.remove(to_remove.0);
        let _was_in = self
            .sources_by_finalized_height
            .remove(&(removed.finalized_block_height, to_remove));
        debug_assert!(_was_in);
        debug_assert!(self.sources.len() >= self.sources_by_finalized_height.len());

        // We make sure to not leave invalid source IDs in the state of `self`.
        // TODO: O(n)
        for item in &mut self.verify_queue {
            if item.downloaded_source == Some(to_remove) {
                item.downloaded_source = None;
            }
        }
        if let RuntimeDownload::NotVerified {
            downloaded_source, ..
        } = &mut self.runtime_download
        {
            if *downloaded_source == Some(to_remove) {
                *downloaded_source = None;
            }
        }
        if let BodyDownload::Downloaded {
            downloaded_source, ..
        } = &mut self.body_download
        {
            if *downloaded_source == Some(to_remove) {
                *downloaded_source = None;
            }
        }
        for (_, call_proof) in &mut self.runtime_calls {
            if let CallProof::Downloaded {
                downloaded_source, ..
            } = call_proof
            {
                if *downloaded_source == Some(to_remove) {
                    *downloaded_source = None;
                }
            }
        }

        let obsolete_requests_indices = self
            .in_progress_requests_by_source
            .range((to_remove, RequestId(usize::MIN))..=(to_remove, RequestId(usize::MAX)))
            .map(|(_, rq_id)| rq_id.0)
            .collect::<Vec<_>>();
        let mut obsolete_requests = Vec::with_capacity(obsolete_requests_indices.len());
        for index in obsolete_requests_indices {
            let (_, user_data, _) = self.in_progress_requests.remove(index);
            self.in_progress_requests_by_source
                .remove(&(to_remove, RequestId(index)));
            if self.warp_sync_fragments_download == Some(RequestId(index)) {
                self.warp_sync_fragments_download = None;
            }
            for call in self.runtime_calls.values_mut() {
                if matches!(call, CallProof::Downloading(rq_id) if *rq_id == RequestId(index)) {
                    *call = CallProof::NotStarted;
                }
            }
            if let RuntimeDownload::Downloading {
                request_id,
                hint_doesnt_match,
            } = &mut self.runtime_download
            {
                if *request_id == RequestId(index) {
                    self.runtime_download = RuntimeDownload::NotStarted {
                        hint_doesnt_match: *hint_doesnt_match,
                    };
                }
            }
            if let BodyDownload::Downloading { request_id } = &mut self.body_download {
                if *request_id == RequestId(index) {
                    self.body_download = BodyDownload::NotStarted;
                }
            }
            obsolete_requests.push((RequestId(index), user_data));
        }

        (removed.user_data, obsolete_requests.into_iter())
    }

    /// Sets the finalized block height of the given source.
    ///
    /// # Panic
    ///
    /// Panics if `source_id` is invalid.
    ///
    pub fn set_source_finality_state(&mut self, source_id: SourceId, finalized_block_height: u64) {
        let stored_height = &mut self.sources[source_id.0].finalized_block_height;

        // Small optimization. No need to do anything more if the block doesn't actuall change.
        if *stored_height == finalized_block_height {
            return;
        }

        // Note that if the new finalized block is below the former one (which is not something
        // that is ever supposed to happen), we should in principle cancel the requests
        // targeting that source that require a specific block height. In practice, however,
        // we don't care as again this isn't supposed to ever happen. While ongoing requests
        // might fail as a result, this is handled the same way as a regular request failure.

        let _was_in = self
            .sources_by_finalized_height
            .remove(&(*stored_height, source_id));
        debug_assert!(_was_in);
        let _inserted = self
            .sources_by_finalized_height
            .insert((finalized_block_height, source_id));
        debug_assert!(_inserted);

        *stored_height = finalized_block_height;
    }

    /// Gets the finalized block height of the given source.
    ///
    /// Equal to 0 if [`WarpSync::set_source_finality_state`] hasn't been called.
    ///
    /// # Panic
    ///
    /// Panics if `source_id` is invalid.
    ///
    pub fn source_finality_state(&self, source_id: SourceId) -> u64 {
        self.sources[source_id.0].finalized_block_height
    }

    /// Returns a list of requests that should be started in order to drive the warp syncing
    /// process to completion.
    ///
    /// Once a request that matches a desired request is added through
    /// [`WarpSync::add_request`], it is no longer returned by this function.
    pub fn desired_requests(&self) -> impl Iterator<Item = (SourceId, &TSrc, DesiredRequest)> {
        // If we are in the fragments download phase, return a fragments download request.
        let mut desired_warp_sync_request = if self.warp_sync_fragments_download.is_none() {
            if self.verify_queue.iter().fold(0, |sum, entry| {
                sum + entry.fragments.len() - entry.next_fragment_to_verify_index
            }) < self.num_download_ahead_fragments
            {
                // Block hash to request.
                let start_block_hash = self
                    .verify_queue
                    .back()
                    .and_then(|entry| entry.fragments.last())
                    .map(|fragment| {
                        header::hash_from_scale_encoded_header(&fragment.scale_encoded_header)
                    })
                    .unwrap_or(self.warped_header_hash);

                // Calculate the block number at the tail of the verify queue.
                // Contains `None` if the verify queue has a problem such as an indecodable header.
                // In that situation, we don't start any new request and wait for the verify
                // queue to empty itself.
                let verify_queue_tail_block_number = self
                    .verify_queue
                    .back()
                    .map(|entry| {
                        entry
                            .fragments
                            .last()
                            .and_then(|fragment| {
                                header::decode(
                                    &fragment.scale_encoded_header,
                                    self.block_number_bytes,
                                )
                                .ok()
                            })
                            .map(|header| header.number)
                    })
                    .unwrap_or(Some(self.warped_header_number));
                let warp_sync_minimum_gap = self.warp_sync_minimum_gap;

                if let Some(verify_queue_tail_block_number) = verify_queue_tail_block_number {
                    // Combine the request with every single available source.
                    either::Left(self.sources.iter().filter_map(move |(src_id, src)| {
                        if src.finalized_block_height
                            <= verify_queue_tail_block_number.saturating_add(
                                u64::try_from(warp_sync_minimum_gap).unwrap_or(u64::MAX),
                            )
                        {
                            return None;
                        }

                        Some((
                            SourceId(src_id),
                            &src.user_data,
                            DesiredRequest::WarpSyncRequest {
                                block_hash: start_block_hash,
                            },
                        ))
                    }))
                } else {
                    either::Right(iter::empty())
                }
            } else {
                either::Right(iter::empty())
            }
        } else {
            either::Right(iter::empty())
        }
        .peekable();

        // If we are in the appropriate phase, and we are not currently downloading the runtime,
        // return a runtime download request.
        let desired_runtime_parameters_get = if let (
            WarpedBlockTy::Normal,
            RuntimeDownload::NotStarted { hint_doesnt_match },
            None,
            true,
            None,
        ) = (
            &self.warped_block_ty,
            &self.runtime_download,
            self.warp_sync_fragments_download,
            self.verify_queue.is_empty(),
            desired_warp_sync_request.peek(),
        ) {
            let code_key_to_request = if let (false, Some(hint)) =
                (*hint_doesnt_match, self.code_trie_node_hint.as_ref())
            {
                Cow::Owned(
                    trie::nibbles_to_bytes_truncate(
                        hint.closest_ancestor_excluding.iter().copied(),
                    )
                    .collect::<Vec<_>>(),
                )
            } else {
                Cow::Borrowed(&b":code"[..])
            };

            // Sources are ordered by increasing finalized block height, in order to
            // have the highest chance for the block to not be pruned.
            let sources_with_block = self
                .sources_by_finalized_height
                .range((self.warped_header_number, SourceId(usize::MIN))..)
                .map(|(_, src_id)| src_id);

            either::Left(sources_with_block.map(move |source_id| {
                (
                    *source_id,
                    &self.sources[source_id.0].user_data,
                    DesiredRequest::StorageGetMerkleProof {
                        block_hash: self.warped_header_hash,
                        state_trie_root: self.warped_header_state_root,
                        keys: vec![code_key_to_request.to_vec(), b":heappages".to_vec()],
                    },
                )
            }))
        } else {
            either::Right(iter::empty())
        };

        // If we are in the appropriate phase, and we are not currently downloading the body of
        // the block, return a runtime download request.
        let desired_body_download =
            if let (WarpedBlockTy::Normal, BodyDownload::NotStarted, None, true, None) = (
                &self.warped_block_ty,
                &self.body_download,
                self.warp_sync_fragments_download,
                self.verify_queue.is_empty(),
                desired_warp_sync_request.peek(),
            ) {
                // Sources are ordered by increasing finalized block height, in order to
                // have the highest chance for the block to not be pruned.
                let sources_with_block = self
                    .sources_by_finalized_height
                    .range((self.warped_header_number, SourceId(usize::MIN))..)
                    .map(|(_, src_id)| src_id);

                either::Left(sources_with_block.map(move |source_id| {
                    (
                        *source_id,
                        &self.sources[source_id.0].user_data,
                        DesiredRequest::BlockBodyDownload {
                            block_hash: self.warped_header_hash,
                            block_number: self.warped_header_number,
                            extrinsics_root: self.warped_header_extrinsics_root,
                        },
                    )
                }))
            } else {
                either::Right(iter::empty())
            };

        // Return the list of runtime calls indicated by the chain information builder state
        // machine.
        let desired_call_proofs = if matches!(self.warped_block_ty, WarpedBlockTy::Normal)
            && self.warp_sync_fragments_download.is_none()
            && self.verify_queue.is_empty()
            && desired_warp_sync_request.peek().is_none()
        {
            either::Left(
                self.runtime_calls
                    .iter()
                    .filter(|(_, v)| matches!(v, CallProof::NotStarted))
                    .map(|(call, _)| DesiredRequest::RuntimeCallMerkleProof {
                        block_hash: self.warped_header_hash,
                        function_name: call.function_name().into(),
                        parameter_vectored: Cow::Owned(call.parameter_vectored_vec()),
                    })
                    .flat_map(move |request_detail| {
                        // Sources are ordered by increasing finalized block height, in order to
                        // have the highest chance for the block to not be pruned.
                        let sources_with_block = self
                            .sources_by_finalized_height
                            .range((self.warped_header_number, SourceId(usize::MIN))..)
                            .map(|(_, src_id)| src_id);

                        sources_with_block.map(move |source_id| {
                            (
                                *source_id,
                                &self.sources[source_id.0].user_data,
                                request_detail.clone(),
                            )
                        })
                    }),
            )
        } else {
            either::Right(iter::empty())
        };

        // Chain all these demanded requests together.
        desired_warp_sync_request
            .chain(desired_runtime_parameters_get)
            .chain(desired_body_download)
            .chain(desired_call_proofs)
    }

    /// Inserts a new request in the data structure.
    ///
    /// > **Note**: The request doesn't necessarily have to match a request returned by
    /// >           [`WarpSync::desired_requests`].
    ///
    /// # Panic
    ///
    /// Panics if the [`SourceId`] is out of range.
    ///
    pub fn add_request(
        &mut self,
        source_id: SourceId,
        user_data: TRq,
        detail: RequestDetail,
    ) -> RequestId {
        assert!(self.sources.contains(source_id.0));

        let request_slot = self.in_progress_requests.vacant_entry();
        let request_id = RequestId(request_slot.key());

        match (&detail, &mut self.runtime_download, &mut self.body_download) {
            (RequestDetail::WarpSyncRequest { block_hash }, _, _)
                if self.warp_sync_fragments_download.is_none()
                    && *block_hash
                        == self
                            .verify_queue
                            .back()
                            .and_then(|entry| entry.fragments.last())
                            .map(|fragment| {
                                header::hash_from_scale_encoded_header(
                                    &fragment.scale_encoded_header,
                                )
                            })
                            .unwrap_or(self.warped_header_hash) =>
            {
                self.warp_sync_fragments_download = Some(request_id);
            }
            (
                RequestDetail::BlockBodyDownload {
                    block_hash,
                    block_number,
                },
                _,
                BodyDownload::NotStarted,
            ) => {
                if self.sources[source_id.0].finalized_block_height >= self.warped_header_number
                    && *block_number == self.warped_header_number
                    && *block_hash == self.warped_header_hash
                {
                    self.body_download = BodyDownload::Downloading { request_id };
                }
            }
            (
                RequestDetail::StorageGetMerkleProof { block_hash, keys },
                RuntimeDownload::NotStarted { hint_doesnt_match },
                _,
            ) => {
                let code_key_to_request = if let (false, Some(hint)) =
                    (*hint_doesnt_match, self.code_trie_node_hint.as_ref())
                {
                    Cow::Owned(
                        trie::nibbles_to_bytes_truncate(
                            hint.closest_ancestor_excluding.iter().copied(),
                        )
                        .collect::<Vec<_>>(),
                    )
                } else {
                    Cow::Borrowed(&b":code"[..])
                };

                if self.sources[source_id.0].finalized_block_height >= self.warped_header_number
                    && *block_hash == self.warped_header_hash
                    && keys.iter().any(|k| *k == *code_key_to_request)
                    && keys.iter().any(|k| k == b":heappages")
                {
                    self.runtime_download = RuntimeDownload::Downloading {
                        hint_doesnt_match: *hint_doesnt_match,
                        request_id,
                    };
                }
            }
            (
                RequestDetail::RuntimeCallMerkleProof {
                    block_hash,
                    function_name,
                    parameter_vectored,
                },
                _,
                _,
            ) => {
                for (info, status) in &mut self.runtime_calls {
                    if matches!(status, CallProof::NotStarted)
                        && self.sources[source_id.0].finalized_block_height
                            >= self.warped_header_number
                        && *block_hash == self.warped_header_hash
                        && function_name == info.function_name()
                        && parameters_equal(parameter_vectored, info.parameter_vectored())
                    {
                        *status = CallProof::Downloading(request_id);
                        break;
                    }
                }
            }
            _ => {}
        }

        request_slot.insert((source_id, user_data, detail));
        let _was_inserted = self
            .in_progress_requests_by_source
            .insert((source_id, request_id));
        debug_assert!(_was_inserted);
        request_id
    }

    /// Removes the given request from the state machine. Returns the user data that was associated
    /// to it.
    ///
    /// > **Note**: The state machine might want to re-start the same request again. It is out of
    /// >           the scope of this module to keep track of requests that don't succeed.
    ///
    /// # Panic
    ///
    /// Panics if the [`RequestId`] is invalid.
    ///
    pub fn remove_request(&mut self, id: RequestId) -> TRq {
        if self.warp_sync_fragments_download == Some(id) {
            self.warp_sync_fragments_download = None;
        }

        for call in self.runtime_calls.values_mut() {
            if matches!(call, CallProof::Downloading(rq_id) if *rq_id == id) {
                *call = CallProof::NotStarted;
            }
        }

        if let RuntimeDownload::Downloading {
            request_id,
            hint_doesnt_match,
        } = &mut self.runtime_download
        {
            if *request_id == id {
                self.runtime_download = RuntimeDownload::NotStarted {
                    hint_doesnt_match: *hint_doesnt_match,
                }
            }
        }

        if let BodyDownload::Downloading { request_id } = &mut self.body_download {
            if *request_id == id {
                self.body_download = BodyDownload::NotStarted;
            }
        }

        let (source_id, user_data, _) = self.in_progress_requests.remove(id.0);
        let _was_removed = self.in_progress_requests_by_source.remove(&(source_id, id));
        debug_assert!(_was_removed);
        user_data
    }

    /// Returns the [`SourceId`] that is expected to fulfill the given request.
    ///
    /// # Panic
    ///
    /// Panics if the [`RequestId`] is invalid.
    ///
    pub fn request_source_id(&self, request_id: RequestId) -> SourceId {
        self.in_progress_requests[request_id.0].0
    }

    /// Injects a body download and removes the given request from the state machine.
    /// Returns the user data that was associated to it.
    ///
    /// # Panic
    ///
    /// Panics if the [`RequestId`] is invalid.
    /// Panics if the [`RequestId`] doesn't correspond to a body download request.
    ///
    pub fn body_download_response(&mut self, id: RequestId, body: Vec<Vec<u8>>) -> TRq {
        // Remove the request from the list, obtaining its user data.
        // If the request corresponds to the runtime parameters we're looking for, the function
        // continues below, otherwise we return early.
        let (source_id, user_data) =
            match (self.in_progress_requests.remove(id.0), &self.body_download) {
                ((source_id, user_data, _), BodyDownload::Downloading { request_id })
                    if *request_id == id =>
                {
                    (source_id, user_data)
                }
                ((source_id, user_data, RequestDetail::BlockBodyDownload { .. }), _) => {
                    let _was_removed = self.in_progress_requests_by_source.remove(&(source_id, id));
                    debug_assert!(_was_removed);
                    return user_data;
                }
                (
                    (
                        _,
                        _,
                        RequestDetail::RuntimeCallMerkleProof { .. }
                        | RequestDetail::WarpSyncRequest { .. }
                        | RequestDetail::StorageGetMerkleProof { .. },
                    ),
                    _,
                ) => panic!(),
            };

        self.body_download = BodyDownload::Downloaded {
            downloaded_source: Some(source_id),
            body,
        };

        let _was_removed = self.in_progress_requests_by_source.remove(&(source_id, id));
        debug_assert!(_was_removed);

        user_data
    }

    /// Injects a Merkle proof and removes the given request from the state machine.
    /// Returns the user data that was associated to it.
    ///
    /// # Panic
    ///
    /// Panics if the [`RequestId`] is invalid.
    /// Panics if the [`RequestId`] doesn't correspond to a storage get request.
    ///
    pub fn storage_get_response(&mut self, id: RequestId, merkle_proof: Vec<u8>) -> TRq {
        // Remove the request from the list, obtaining its user data.
        // If the request corresponds to the runtime parameters we're looking for, the function
        // continues below, otherwise we return early.
        let (source_id, hint_doesnt_match, user_data) = match (
            self.in_progress_requests.remove(id.0),
            &self.runtime_download,
        ) {
            (
                (source_id, user_data, _),
                RuntimeDownload::Downloading {
                    request_id,
                    hint_doesnt_match,
                },
            ) if *request_id == id => (source_id, *hint_doesnt_match, user_data),
            ((source_id, user_data, RequestDetail::StorageGetMerkleProof { .. }), _) => {
                let _was_removed = self.in_progress_requests_by_source.remove(&(source_id, id));
                debug_assert!(_was_removed);
                return user_data;
            }
            (
                (
                    _,
                    _,
                    RequestDetail::RuntimeCallMerkleProof { .. }
                    | RequestDetail::WarpSyncRequest { .. }
                    | RequestDetail::BlockBodyDownload { .. },
                ),
                _,
            ) => panic!(),
        };

        self.runtime_download = RuntimeDownload::NotVerified {
            downloaded_source: Some(source_id),
            hint_doesnt_match,
            trie_proof: merkle_proof,
        };

        let _was_removed = self.in_progress_requests_by_source.remove(&(source_id, id));
        debug_assert!(_was_removed);

        user_data
    }

    /// Injects a response and removes the given request from the state machine. Returns
    /// the user data that was associated to it.
    ///
    /// # Panic
    ///
    /// Panics if the [`RequestId`] is invalid.
    /// Panics if the [`RequestId`] doesn't correspond to a runtime Merkle call proof request.
    ///
    pub fn runtime_call_merkle_proof_response(
        &mut self,
        request_id: RequestId,
        response: Vec<u8>,
    ) -> TRq {
        let (source_id, user_data, RequestDetail::RuntimeCallMerkleProof { .. }) =
            self.in_progress_requests.remove(request_id.0)
        else {
            // Wrong request type.
            panic!()
        };

        for call in self.runtime_calls.values_mut() {
            if matches!(call, CallProof::Downloading(rq_id) if *rq_id == request_id) {
                *call = CallProof::Downloaded {
                    downloaded_source: Some(source_id),
                    proof: response,
                };
                break;
            }
        }

        let _was_removed = self
            .in_progress_requests_by_source
            .remove(&(source_id, request_id));
        debug_assert!(_was_removed);

        user_data
    }

    /// Injects a response and removes the given request from the state machine. Returns
    /// the user data that was associated to it.
    ///
    /// If the header of the last fragment of the response is decodable, this function updates
    /// the finalized block of the source.
    ///
    /// # Panic
    ///
    /// Panics if the [`RequestId`] is invalid.
    /// Panics if the [`RequestId`] doesn't correspond to a warp sync request.
    ///
    // TODO: more zero cost API w.r.t. the fragments
    pub fn warp_sync_request_response(
        &mut self,
        request_id: RequestId,
        fragments: Vec<WarpSyncFragment>,
        final_set_of_fragments: bool,
    ) -> TRq {
        let (rq_source_id, user_data) = match self.in_progress_requests.remove(request_id.0) {
            (rq_source_id, user_data, RequestDetail::WarpSyncRequest { .. }) => {
                (rq_source_id, user_data)
            }
            (_, _, _) => panic!(),
        };

        debug_assert!(self.sources.contains(rq_source_id.0));

        // Since we send requests only to sources with an appropriate finalized block, we make
        // sure that the finalized block of the source that sent the response matches the
        // fragments that it sent.
        // If we didn't do that, it would be possible for example to warp sync to block 200 while
        // believing that the source is only at block 199, and thus the warp syncing would stall.
        if let Some(last_header) = fragments
            .last()
            .and_then(|h| header::decode(&h.scale_encoded_header, self.block_number_bytes).ok())
        {
            let src_finalized_height = &mut self.sources[rq_source_id.0].finalized_block_height;

            let new_height = if final_set_of_fragments {
                // If the source indicated that this is the last fragment, then we know that
                // it's also equal to their finalized block.
                last_header.number
            } else {
                // If this is not the last fragment, we know that the finalized block of the
                // source is *at least* the one provided.
                // TODO: could maybe do + gap or something?
                cmp::max(*src_finalized_height, last_header.number.saturating_add(1))
            };

            if *src_finalized_height != new_height {
                let _was_in = self
                    .sources_by_finalized_height
                    .remove(&(*src_finalized_height, rq_source_id));
                debug_assert!(_was_in);

                *src_finalized_height = new_height;

                let _inserted = self
                    .sources_by_finalized_height
                    .insert((*src_finalized_height, rq_source_id));
                debug_assert!(_inserted);
            }
        }

        if self.warp_sync_fragments_download == Some(request_id) {
            self.warp_sync_fragments_download = None;

            self.verify_queue.push_back(PendingVerify {
                final_set_of_fragments,
                downloaded_source: Some(rq_source_id),
                fragments,
                next_fragment_to_verify_index: 0,
            });
        }

        let _was_removed = self
            .in_progress_requests_by_source
            .remove(&(rq_source_id, request_id));
        debug_assert!(_was_removed);

        user_data
    }

    /// Start processing one CPU operation.
    ///
    /// This function takes ownership of `self` and yields it back after the operation is finished.
    // TODO: take a `&mut self` instead of `self` ; requires many changes in all.rs
    pub fn process_one(self) -> ProcessOne<TSrc, TRq> {
        // If we've downloaded everything that was needed, switch to "build chain information"
        // mode.
        if matches!(self.runtime_download, RuntimeDownload::Verified { .. })
            && matches!(
                self.body_download,
                BodyDownload::NotNeeded | BodyDownload::Downloaded { .. }
            )
            && self
                .runtime_calls
                .values()
                .all(|c| matches!(c, CallProof::Downloaded { .. }))
        {
            return ProcessOne::BuildChainInformation(BuildChainInformation { inner: self });
        }

        if let RuntimeDownload::NotVerified { .. } = &self.runtime_download {
            return ProcessOne::BuildRuntime(BuildRuntime { inner: self });
        }

        if !self.verify_queue.is_empty() {
            return ProcessOne::VerifyWarpSyncFragment(VerifyWarpSyncFragment { inner: self });
        }

        ProcessOne::Idle(self)
    }
}

impl<TSrc, TRq> ops::Index<SourceId> for WarpSync<TSrc, TRq> {
    type Output = TSrc;

    #[track_caller]
    fn index(&self, source_id: SourceId) -> &TSrc {
        debug_assert!(self.sources.contains(source_id.0));
        &self.sources[source_id.0].user_data
    }
}

impl<TSrc, TRq> ops::IndexMut<SourceId> for WarpSync<TSrc, TRq> {
    #[track_caller]
    fn index_mut(&mut self, source_id: SourceId) -> &mut TSrc {
        debug_assert!(self.sources.contains(source_id.0));
        &mut self.sources[source_id.0].user_data
    }
}

impl<TSrc, TRq> ops::Index<RequestId> for WarpSync<TSrc, TRq> {
    type Output = TRq;

    #[track_caller]
    fn index(&self, request_id: RequestId) -> &TRq {
        debug_assert!(self.in_progress_requests.contains(request_id.0));
        &self.in_progress_requests[request_id.0].1
    }
}

impl<TSrc, TRq> ops::IndexMut<RequestId> for WarpSync<TSrc, TRq> {
    #[track_caller]
    fn index_mut(&mut self, request_id: RequestId) -> &mut TRq {
        debug_assert!(self.in_progress_requests.contains(request_id.0));
        &mut self.in_progress_requests[request_id.0].1
    }
}

/// Information about a request that the warp sync state machine would like to start.
///
/// See [`WarpSync::desired_requests`].
#[derive(Debug, Clone)]
pub enum DesiredRequest {
    /// A warp sync request should be started.
    WarpSyncRequest {
        /// Starting point of the warp syncing. The first fragment of the response should be the
        /// of the epoch that the starting point is in.
        block_hash: [u8; 32],
    },
    /// A block body download should be started.
    BlockBodyDownload {
        /// Hash of the block whose body to download.
        block_hash: [u8; 32],
        /// Height of the block whose body to download.
        block_number: u64,
        /// Extrinsics trie root hash found in the header of the block.
        extrinsics_root: [u8; 32],
    },
    /// A storage request of the runtime code and heap pages should be started.
    StorageGetMerkleProof {
        /// Hash of the block to request the parameters against.
        block_hash: [u8; 32],
        /// State trie root hash found in the header of the block.
        state_trie_root: [u8; 32],
        /// Keys whose values are requested.
        // TODO: consider Cow<'static, [u8]> instead
        keys: Vec<Vec<u8>>,
    },
    /// A call proof should be started.
    RuntimeCallMerkleProof {
        /// Hash of the header of the block the call should be made against.
        block_hash: [u8; 32],
        /// Name of the function of the call proof.
        function_name: Cow<'static, str>,
        /// Parameters of the call.
        parameter_vectored: Cow<'static, [u8]>,
    },
}

/// Information about a request to add to the state machine.
///
/// See [`WarpSync::add_request`].
#[derive(Debug, Clone)]
pub enum RequestDetail {
    /// See [`DesiredRequest::WarpSyncRequest`].
    WarpSyncRequest {
        /// See [`DesiredRequest::WarpSyncRequest::block_hash`].
        block_hash: [u8; 32],
    },
    /// See [`DesiredRequest::BlockBodyDownload`].
    BlockBodyDownload {
        /// See [`DesiredRequest::BlockBodyDownload::block_hash`].
        block_hash: [u8; 32],
        /// See [`DesiredRequest::BlockBodyDownload::block_number`].
        // TODO: remove this field as it's inappropriate, but this causes issues in all.rs
        block_number: u64,
    },
    /// See [`DesiredRequest::StorageGetMerkleProof`].
    StorageGetMerkleProof {
        /// See [`DesiredRequest::StorageGetMerkleProof::block_hash`].
        block_hash: [u8; 32],
        /// See [`DesiredRequest::StorageGetMerkleProof::keys`].
        keys: Vec<Vec<u8>>,
    },
    /// See [`DesiredRequest::RuntimeCallMerkleProof`].
    RuntimeCallMerkleProof {
        /// See [`DesiredRequest::RuntimeCallMerkleProof::block_hash`].
        block_hash: [u8; 32],
        /// See [`DesiredRequest::RuntimeCallMerkleProof::function_name`].
        function_name: Cow<'static, str>,
        /// See [`DesiredRequest::RuntimeCallMerkleProof::parameter_vectored`].
        parameter_vectored: Cow<'static, [u8]>,
    },
}

/// Identifier for a request in the warp sync state machine.
#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash)]
pub struct RequestId(usize);

/// Return value of [`WarpSync::process_one`].
pub enum ProcessOne<TSrc, TRq> {
    /// Nothing to verify at the moment. The state machine is yielded back.
    Idle(WarpSync<TSrc, TRq>),
    /// Ready to verify a warp sync fragment.
    ///
    /// > **Note**: In case where a source has sent an empty list of fragment, which is invalid,
    /// >           this variant will "verify" the list and produce an error.
    VerifyWarpSyncFragment(VerifyWarpSyncFragment<TSrc, TRq>),
    /// Ready to build the runtime of the chain..
    BuildRuntime(BuildRuntime<TSrc, TRq>),
    /// Ready to verify the parameters of the chain against the finalized block.
    BuildChainInformation(BuildChainInformation<TSrc, TRq>),
}

/// Ready to verify a warp sync fragment.
///
/// > **Note**: In case where a source has sent an empty list of fragment, which is invalid,
/// >           this variant will "verify" the list and produce an error.
pub struct VerifyWarpSyncFragment<TSrc, TRq> {
    inner: WarpSync<TSrc, TRq>,
}

impl<TSrc, TRq> VerifyWarpSyncFragment<TSrc, TRq> {
    /// Returns the source that has sent the fragments that we are about to verify, and its user
    /// data.
    ///
    /// Returns `None` if the source has been removed since the fragments have been downloaded.
    pub fn proof_sender(&self) -> Option<(SourceId, &TSrc)> {
        let entry_to_verify = self.inner.verify_queue.front().unwrap();
        let source_id = entry_to_verify.downloaded_source?;
        Some((source_id, &self.inner.sources[source_id.0].user_data))
    }

    /// Verify one warp sync fragment.
    ///
    /// Must be passed a randomly-generated value that is used by the verification process. Note
    /// that the verification is still deterministic.
    ///
    /// On success, returns the block hash and height that have been verified as being part of
    /// the chain.
    /// On error, returns why the verification has failed. The warp syncing process still
    /// continues.
    pub fn verify(
        mut self,
        randomness_seed: [u8; 32],
    ) -> (
        WarpSync<TSrc, TRq>,
        Result<([u8; 32], u64), VerifyFragmentError>,
    ) {
        // A `VerifyWarpSyncFragment` is only ever created if `verify_queue` is non-empty.
        debug_assert!(!self.inner.verify_queue.is_empty());
        let fragments_to_verify = self
            .inner
            .verify_queue
            .front_mut()
            .unwrap_or_else(|| unreachable!());

        // The source has sent an empty list of fragments. This is invalid.
        if fragments_to_verify.fragments.is_empty() {
            self.inner.verify_queue.pop_front().unwrap();
            return (self.inner, Err(VerifyFragmentError::EmptyProof));
        }

        // Given that the list of fragments is non-empty, we are assuming that there are still
        // fragments to verify, otherwise this entry should have been removed in a previous
        // iteration.
        let fragment_to_verify = fragments_to_verify
            .fragments
            .get(fragments_to_verify.next_fragment_to_verify_index)
            .unwrap_or_else(|| unreachable!());

        // It has been checked at the warp sync initialization that the finality algorithm is
        // indeed Grandpa.
        let chain_information::ChainInformationFinality::Grandpa {
            after_finalized_block_authorities_set_id,
            finalized_triggered_authorities,
            .. // TODO: support finalized_scheduled_change? difficult to implement
        } = &mut self.inner.warped_finality
        else {
            unreachable!()
        };

        // Decode the header and justification of the fragment.
        let fragment_header_hash =
            header::hash_from_scale_encoded_header(&fragment_to_verify.scale_encoded_header);
        let fragment_decoded_header = match header::decode(
            &fragment_to_verify.scale_encoded_header,
            self.inner.block_number_bytes,
        ) {
            Ok(j) => j,
            Err(err) => {
                self.inner.verify_queue.clear();
                self.inner.warp_sync_fragments_download = None;
                return (self.inner, Err(VerifyFragmentError::InvalidHeader(err)));
            }
        };
        let fragment_decoded_justification = match decode::decode_grandpa_justification(
            &fragment_to_verify.scale_encoded_justification,
            self.inner.block_number_bytes,
        ) {
            Ok(j) => j,
            Err(err) => {
                self.inner.verify_queue.clear();
                self.inner.warp_sync_fragments_download = None;
                return (
                    self.inner,
                    Err(VerifyFragmentError::InvalidJustification(err)),
                );
            }
        };

        // Make sure that the header would actually advance the warp sync process forward.
        if fragment_decoded_header.number <= self.inner.warped_header_number {
            self.inner.verify_queue.clear();
            self.inner.warp_sync_fragments_download = None;
            return (
                self.inner,
                Err(VerifyFragmentError::BlockNumberNotIncrementing),
            );
        }

        // Make sure that the justification indeed corresponds to the header.
        if *fragment_decoded_justification.target_hash != fragment_header_hash
            || fragment_decoded_justification.target_number != fragment_decoded_header.number
        {
            let error = VerifyFragmentError::TargetHashMismatch {
                justification_target_hash: *fragment_decoded_justification.target_hash,
                justification_target_height: fragment_decoded_justification.target_number,
                header_hash: fragment_header_hash,
                header_height: fragment_decoded_header.number,
            };
            self.inner.verify_queue.clear();
            self.inner.warp_sync_fragments_download = None;
            return (self.inner, Err(error));
        }

        // Check whether the justification is valid.
        if let Err(err) = verify::verify_justification(verify::JustificationVerifyConfig {
            justification: &fragment_to_verify.scale_encoded_justification,
            block_number_bytes: self.inner.block_number_bytes,
            authorities_list: finalized_triggered_authorities
                .iter()
                .map(|a| &a.public_key[..]),
            authorities_set_id: *after_finalized_block_authorities_set_id,
            randomness_seed,
        }) {
            self.inner.verify_queue.clear();
            self.inner.warp_sync_fragments_download = None;
            return (
                self.inner,
                Err(VerifyFragmentError::JustificationVerify(err)),
            );
        }

        // Try to grab the new list of authorities from the header.
        let new_authorities_list = fragment_decoded_header
            .digest
            .logs()
            .find_map(|log_item| match log_item {
                header::DigestItemRef::GrandpaConsensus(grandpa_log_item) => match grandpa_log_item
                {
                    header::GrandpaConsensusLogRef::ScheduledChange(change)
                    | header::GrandpaConsensusLogRef::ForcedChange { change, .. } => {
                        Some(change.next_authorities)
                    }
                    _ => None,
                },
                _ => None,
            })
            .map(|next_authorities| {
                next_authorities
                    .map(header::GrandpaAuthority::from)
                    .collect()
            });

        // Fragments must only include headers containing an update to the list of authorities,
        // unless it's the very head of the chain.
        if new_authorities_list.is_none()
            && (!fragments_to_verify.final_set_of_fragments
                || fragments_to_verify.next_fragment_to_verify_index
                    != fragments_to_verify.fragments.len() - 1)
        {
            self.inner.verify_queue.clear();
            self.inner.warp_sync_fragments_download = None;
            return (self.inner, Err(VerifyFragmentError::NonMinimalProof));
        }

        // Verification of the fragment has succeeded 🎉. We can now update `self`.
        fragments_to_verify.next_fragment_to_verify_index += 1;
        self.inner.warped_header_number = fragment_decoded_header.number;
        self.inner.warped_header_state_root = *fragment_decoded_header.state_root;
        self.inner.warped_header_extrinsics_root = *fragment_decoded_header.extrinsics_root;
        self.inner.warped_header_hash = fragment_header_hash;
        self.inner.warped_header = fragment_to_verify.scale_encoded_header.clone(); // TODO: figure out how to remove this clone()
        self.inner.warped_block_ty = WarpedBlockTy::Normal;
        self.inner.runtime_download = RuntimeDownload::NotStarted {
            hint_doesnt_match: false,
        };
        if !matches!(self.inner.body_download, BodyDownload::NotNeeded) {
            self.inner.body_download = BodyDownload::NotStarted;
        }
        self.inner.runtime_calls =
            runtime_calls_default_value(self.inner.verified_chain_information.as_ref().consensus);
        if let Some(new_authorities_list) = new_authorities_list {
            *finalized_triggered_authorities = new_authorities_list;
            *after_finalized_block_authorities_set_id += 1;
        }
        if fragments_to_verify.next_fragment_to_verify_index == fragments_to_verify.fragments.len()
        {
            self.inner.verify_queue.pop_front().unwrap();
        }

        // Returning.
        let result = Ok((
            self.inner.warped_header_hash,
            self.inner.warped_header_number,
        ));
        (self.inner, result)
    }
}

/// Error potentially returned by [`VerifyWarpSyncFragment::verify`].
#[derive(Debug)]
pub enum VerifyFragmentError {
    /// Justification found within the fragment is invalid.
    JustificationVerify(verify::JustificationVerifyError),
    /// Mismatch between the block targeted by the justification and the header.
    TargetHashMismatch {
        /// Hash of the block the justification targets.
        justification_target_hash: [u8; 32],
        /// Height of the block the justification targets.
        justification_target_height: u64,
        /// Hash of the header.
        header_hash: [u8; 32],
        /// Height of the header.
        header_height: u64,
    },
    /// Warp sync fragment doesn't contain an authorities list change when it should.
    NonMinimalProof,
    /// Header does not actually advance the warp syncing process. This means that a source has
    /// sent a header below the requested hash.
    BlockNumberNotIncrementing,
    /// Warp sync proof is empty.
    EmptyProof,
    /// Failed to decode header.
    InvalidHeader(header::Error),
    /// Failed to decode justification.
    InvalidJustification(decode::JustificationDecodeError),
}

impl fmt::Display for VerifyFragmentError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            VerifyFragmentError::JustificationVerify(err) => fmt::Display::fmt(err, f),
            VerifyFragmentError::TargetHashMismatch {
                justification_target_hash,
                justification_target_height,
                header_hash,
                header_height,
            } => {
                write!(
                    f,
                    "Justification target (hash: {}, height: {}) doesn't match the associated header (hash: {}, height: {})",
                    HashDisplay(justification_target_hash),
                    justification_target_height,
                    HashDisplay(header_hash),
                    header_height,
                )
            }
            VerifyFragmentError::NonMinimalProof => write!(
                f,
                "Warp sync proof fragment doesn't contain an authorities list change"
            ),
            VerifyFragmentError::BlockNumberNotIncrementing => write!(
                f,
                "Warp sync proof header doesn't advance the warp syncing process"
            ),
            VerifyFragmentError::EmptyProof => write!(f, "Warp sync proof is empty"),
            VerifyFragmentError::InvalidHeader(err) => write!(f, "Failed to decode header: {err}"),
            VerifyFragmentError::InvalidJustification(err) => {
                write!(f, "Failed to decode justification: {err}")
            }
        }
    }
}

/// Problem encountered during a call to [`BuildRuntime::build`] or
/// [`BuildChainInformation::build`] that can be attributed to the source sending invalid data.
#[derive(Debug, derive_more::Display, derive_more::Error)]
#[display("{error}")]
pub struct SourceMisbehavior {
    /// Source that committed the felony. `None` if the source has been removed between the moment
    /// when the request has succeeded and when it has been verified.
    pub source_id: Option<SourceId>,
    /// Error that the source made.
    #[error(source)]
    pub error: SourceMisbehaviorTy,
}

/// See [`SourceMisbehavior::error`].
#[derive(Debug, derive_more::Display, derive_more::Error)]
pub enum SourceMisbehaviorTy {
    /// Failed to verify Merkle proof.
    InvalidMerkleProof(proof_decode::Error),
    /// Merkle proof is missing the necessary entries.
    MerkleProofEntriesMissing,
    /// Downloaded block body doesn't match the expected extrinsics root.
    BlockBodyExtrinsicsRootMismatch,
}

/// Problem encountered during a call to [`BuildRuntime::build`].
#[derive(Debug, derive_more::Display, derive_more::Error)]
pub enum BuildRuntimeError {
    /// The chain doesn't include any storage item at `:code`.
    #[display("The chain doesn't include any storage item at `:code`")]
    MissingCode,
    /// The storage item at `:heappages` is in an incorrect format.
    #[display("Invalid heap pages value: {_0}")]
    InvalidHeapPages(executor::InvalidHeapPagesError),
    /// Error building the runtime of the chain.
    #[display("Error building the runtime: {_0}")]
    RuntimeBuild(executor::host::NewErr),
    /// Source that has sent a proof didn't behave properly.
    SourceMisbehavior(SourceMisbehavior),
}

/// Ready to build the runtime of the finalized chain.
pub struct BuildRuntime<TSrc, TRq> {
    inner: WarpSync<TSrc, TRq>,
}

impl<TSrc, TRq> BuildRuntime<TSrc, TRq> {
    /// Build the runtime of the chain.
    ///
    /// Must be passed parameters used for the construction of the runtime: a hint as to whether
    /// the runtime is trusted and/or will be executed again, and whether unresolved function
    /// imports are allowed.
    pub fn build(
        mut self,
        exec_hint: ExecHint,
        allow_unresolved_imports: bool,
    ) -> (WarpSync<TSrc, TRq>, Result<(), BuildRuntimeError>) {
        let RuntimeDownload::NotVerified {
            downloaded_source,
            hint_doesnt_match,
            trie_proof,
        } = &mut self.inner.runtime_download
        else {
            unreachable!()
        };

        let downloaded_runtime = mem::take(trie_proof);
        let decoded_downloaded_runtime =
            match proof_decode::decode_and_verify_proof(proof_decode::Config {
                proof: &downloaded_runtime[..],
            }) {
                Ok(p) => p,
                Err(err) => {
                    let downloaded_source = *downloaded_source;
                    self.inner.runtime_download = RuntimeDownload::NotStarted {
                        hint_doesnt_match: *hint_doesnt_match,
                    };
                    return (
                        self.inner,
                        Err(BuildRuntimeError::SourceMisbehavior(SourceMisbehavior {
                            source_id: downloaded_source,
                            error: SourceMisbehaviorTy::InvalidMerkleProof(err),
                        })),
                    );
                }
            };

        let (
            finalized_storage_code_merkle_value,
            finalized_storage_code_closest_ancestor_excluding,
        ) = {
            let code_nibbles = trie::bytes_to_nibbles(b":code".iter().copied()).collect::<Vec<_>>();
            match decoded_downloaded_runtime.closest_ancestor_in_proof(
                &self.inner.warped_header_state_root,
                code_nibbles.iter().take(code_nibbles.len() - 1).copied(),
            ) {
                Ok(Some(closest_ancestor_key)) => {
                    let closest_ancestor_key = closest_ancestor_key.collect::<Vec<_>>();
                    let next_nibble = code_nibbles[closest_ancestor_key.len()];
                    let merkle_value = decoded_downloaded_runtime
                        .trie_node_info(
                            &self.inner.warped_header_state_root,
                            closest_ancestor_key.iter().copied(),
                        )
                        .unwrap()
                        .children
                        .child(next_nibble)
                        .merkle_value();

                    match merkle_value {
                        Some(mv) => (mv.to_owned(), closest_ancestor_key),
                        None => {
                            self.inner.warped_block_ty = WarpedBlockTy::KnownBad;
                            self.inner.runtime_download = RuntimeDownload::NotStarted {
                                hint_doesnt_match: *hint_doesnt_match,
                            };
                            if !matches!(self.inner.body_download, BodyDownload::NotNeeded) {
                                self.inner.body_download = BodyDownload::NotStarted;
                            }
                            return (self.inner, Err(BuildRuntimeError::MissingCode));
                        }
                    }
                }
                Ok(None) => {
                    self.inner.warped_block_ty = WarpedBlockTy::KnownBad;
                    self.inner.runtime_download = RuntimeDownload::NotStarted {
                        hint_doesnt_match: *hint_doesnt_match,
                    };
                    if !matches!(self.inner.body_download, BodyDownload::NotNeeded) {
                        self.inner.body_download = BodyDownload::NotStarted;
                    }
                    return (self.inner, Err(BuildRuntimeError::MissingCode));
                }
                Err(proof_decode::IncompleteProofError { .. }) => {
                    let downloaded_source = *downloaded_source;
                    self.inner.runtime_download = RuntimeDownload::NotStarted {
                        hint_doesnt_match: *hint_doesnt_match,
                    };
                    if !matches!(self.inner.body_download, BodyDownload::NotNeeded) {
                        self.inner.body_download = BodyDownload::NotStarted;
                    }
                    return (
                        self.inner,
                        Err(BuildRuntimeError::SourceMisbehavior(SourceMisbehavior {
                            source_id: downloaded_source,
                            error: SourceMisbehaviorTy::MerkleProofEntriesMissing,
                        })),
                    );
                }
            }
        };

        let finalized_storage_code = if let (false, Some(hint)) =
            (*hint_doesnt_match, self.inner.code_trie_node_hint.as_ref())
        {
            if hint.merkle_value == finalized_storage_code_merkle_value {
                &hint.storage_value
            } else {
                self.inner.runtime_download = RuntimeDownload::NotStarted {
                    hint_doesnt_match: true,
                };
                return (self.inner, Ok(()));
            }
        } else {
            match decoded_downloaded_runtime
                .storage_value(&self.inner.warped_header_state_root, b":code")
            {
                Ok(Some((code, _))) => code,
                Ok(None) => {
                    self.inner.warped_block_ty = WarpedBlockTy::KnownBad;
                    self.inner.runtime_download = RuntimeDownload::NotStarted {
                        hint_doesnt_match: *hint_doesnt_match,
                    };
                    if !matches!(self.inner.body_download, BodyDownload::NotNeeded) {
                        self.inner.body_download = BodyDownload::NotStarted;
                    }
                    return (self.inner, Err(BuildRuntimeError::MissingCode));
                }
                Err(proof_decode::IncompleteProofError { .. }) => {
                    let downloaded_source = *downloaded_source;
                    return (
                        self.inner,
                        Err(BuildRuntimeError::SourceMisbehavior(SourceMisbehavior {
                            source_id: downloaded_source,
                            error: SourceMisbehaviorTy::MerkleProofEntriesMissing,
                        })),
                    );
                }
            }
        };

        let finalized_storage_heappages = match decoded_downloaded_runtime
            .storage_value(&self.inner.warped_header_state_root, b":heappages")
        {
            Ok(val) => val.map(|(v, _)| v),
            Err(proof_decode::IncompleteProofError { .. }) => {
                let downloaded_source = *downloaded_source;
                return (
                    self.inner,
                    Err(BuildRuntimeError::SourceMisbehavior(SourceMisbehavior {
                        source_id: downloaded_source,
                        error: SourceMisbehaviorTy::MerkleProofEntriesMissing,
                    })),
                );
            }
        };

        let decoded_heap_pages =
            match executor::storage_heap_pages_to_value(finalized_storage_heappages) {
                Ok(hp) => hp,
                Err(err) => {
                    self.inner.warped_block_ty = WarpedBlockTy::KnownBad;
                    self.inner.runtime_download = RuntimeDownload::NotStarted {
                        hint_doesnt_match: *hint_doesnt_match,
                    };
                    if !matches!(self.inner.body_download, BodyDownload::NotNeeded) {
                        self.inner.body_download = BodyDownload::NotStarted;
                    }
                    return (self.inner, Err(BuildRuntimeError::InvalidHeapPages(err)));
                }
            };

        let runtime = match HostVmPrototype::new(host::Config {
            module: &finalized_storage_code,
            heap_pages: decoded_heap_pages,
            exec_hint,
            allow_unresolved_imports,
        }) {
            Ok(runtime) => runtime,
            Err(err) => {
                self.inner.warped_block_ty = WarpedBlockTy::KnownBad;
                self.inner.runtime_download = RuntimeDownload::NotStarted {
                    hint_doesnt_match: *hint_doesnt_match,
                };
                if !matches!(self.inner.body_download, BodyDownload::NotNeeded) {
                    self.inner.body_download = BodyDownload::NotStarted;
                }
                return (self.inner, Err(BuildRuntimeError::RuntimeBuild(err)));
            }
        };

        let chain_info_builder = chain_information::build::ChainInformationBuild::new(
            chain_information::build::Config {
                finalized_block_header: chain_information::build::ConfigFinalizedBlockHeader::Any {
                    scale_encoded_header: self.inner.warped_header.clone(),
                    known_finality: if self.inner.download_all_chain_information_storage_proofs {
                        None
                    } else {
                        Some(self.inner.warped_finality.clone())
                    },
                },
                block_number_bytes: self.inner.block_number_bytes,
                runtime,
            },
        );

        if let chain_information::build::ChainInformationBuild::InProgress(in_progress) =
            &chain_info_builder
        {
            for call in in_progress.remaining_calls() {
                if let hashbrown::hash_map::Entry::Vacant(entry) =
                    self.inner.runtime_calls.entry(call)
                {
                    entry.insert(CallProof::NotStarted);
                }
            }
        }

        self.inner.runtime_download = RuntimeDownload::Verified {
            downloaded_runtime: DownloadedRuntime {
                storage_code: Some(finalized_storage_code.to_vec()),
                storage_heap_pages: finalized_storage_heappages.map(|v| v.to_vec()),
                code_merkle_value: Some(finalized_storage_code_merkle_value),
                closest_ancestor_excluding: Some(finalized_storage_code_closest_ancestor_excluding),
            },
            chain_info_builder,
        };

        (self.inner, Ok(()))
    }
}

/// Problem encountered during a call to [`BuildChainInformation::build`].
#[derive(Debug, derive_more::Display, derive_more::Error)]
pub enum BuildChainInformationError {
    /// Error building the chain information.
    #[display("Error building the chain information: {_0}")]
    ChainInformationBuild(chain_information::build::Error),
    /// Source that has sent a proof didn't behave properly.
    SourceMisbehavior(SourceMisbehavior),
}

/// Ready to verify the parameters of the chain against the finalized block.
pub struct BuildChainInformation<TSrc, TRq> {
    inner: WarpSync<TSrc, TRq>,
}

impl<TSrc, TRq> BuildChainInformation<TSrc, TRq> {
    /// Build the information about the chain.
    pub fn build(
        mut self,
    ) -> (
        WarpSync<TSrc, TRq>,
        Result<RuntimeInformation, BuildChainInformationError>,
    ) {
        let downloaded_body = match &mut self.inner.body_download {
            BodyDownload::NotNeeded => None,
            BodyDownload::Downloaded {
                downloaded_source,
                body,
            } => {
                if header::extrinsics_root(body) != self.inner.warped_header_extrinsics_root {
                    let source_id = *downloaded_source;
                    self.inner.body_download = BodyDownload::NotStarted;
                    return (
                        self.inner,
                        Err(BuildChainInformationError::SourceMisbehavior(
                            SourceMisbehavior {
                                source_id,
                                error: SourceMisbehaviorTy::BlockBodyExtrinsicsRootMismatch,
                            },
                        )),
                    );
                }

                Some(body)
            }
            _ => unreachable!(),
        };

        let RuntimeDownload::Verified {
            mut chain_info_builder,
            downloaded_runtime,
            ..
        } = mem::replace(
            &mut self.inner.runtime_download,
            RuntimeDownload::NotStarted {
                hint_doesnt_match: false,
            },
        )
        else {
            unreachable!()
        };

        let runtime_calls = mem::take(&mut self.inner.runtime_calls);

        debug_assert!(
            runtime_calls
                .values()
                .all(|c| matches!(c, CallProof::Downloaded { .. }))
        );

        // Decode all the Merkle proofs that have been received.
        let calls = {
            let mut decoded_proofs = hashbrown::HashMap::with_capacity_and_hasher(
                runtime_calls.len(),
                fnv::FnvBuildHasher::default(),
            );

            for (call, proof) in runtime_calls {
                let CallProof::Downloaded {
                    proof,
                    downloaded_source,
                } = proof
                else {
                    unreachable!()
                };

                let decoded_proof =
                    match proof_decode::decode_and_verify_proof(proof_decode::Config {
                        proof: proof.into_iter(),
                    }) {
                        Ok(d) => d,
                        Err(err) => {
                            return (
                                self.inner,
                                Err(BuildChainInformationError::SourceMisbehavior(
                                    SourceMisbehavior {
                                        source_id: downloaded_source,
                                        error: SourceMisbehaviorTy::InvalidMerkleProof(err),
                                    },
                                )),
                            );
                        }
                    };

                decoded_proofs.insert(call, (decoded_proof, downloaded_source));
            }

            decoded_proofs
        };

        loop {
            let in_progress = match chain_info_builder {
                chain_information::build::ChainInformationBuild::Finished {
                    result: Ok(chain_information),
                    virtual_machine,
                } => {
                    // This `if` is necessary as in principle we might have continued warp syncing
                    // after downloading everything needed but before building the chain
                    // information.
                    if self.inner.warped_header_number
                        == chain_information.as_ref().finalized_block_header.number
                    {
                        self.inner.warped_block_ty = WarpedBlockTy::AlreadyVerified;
                    }

                    let finalized_body = downloaded_body.map(mem::take);
                    if !matches!(self.inner.body_download, BodyDownload::NotNeeded) {
                        self.inner.body_download = BodyDownload::NotStarted;
                    }

                    self.inner.verified_chain_information = chain_information;
                    self.inner.runtime_calls = runtime_calls_default_value(
                        self.inner.verified_chain_information.as_ref().consensus,
                    );
                    return (
                        self.inner,
                        Ok(RuntimeInformation {
                            finalized_runtime: virtual_machine,
                            finalized_body,
                            finalized_storage_code: downloaded_runtime.storage_code,
                            finalized_storage_heap_pages: downloaded_runtime.storage_heap_pages,
                            finalized_storage_code_merkle_value: downloaded_runtime
                                .code_merkle_value,
                            finalized_storage_code_closest_ancestor_excluding: downloaded_runtime
                                .closest_ancestor_excluding,
                        }),
                    );
                }
                chain_information::build::ChainInformationBuild::Finished {
                    result: Err(err),
                    ..
                } => {
                    self.inner.warped_block_ty = WarpedBlockTy::KnownBad;
                    return (
                        self.inner,
                        Err(BuildChainInformationError::ChainInformationBuild(err)),
                    );
                }
                chain_information::build::ChainInformationBuild::InProgress(in_progress) => {
                    in_progress
                }
            };

            chain_info_builder = match in_progress {
                chain_information::build::InProgress::StorageGet(get) => {
                    // TODO: child tries not supported
                    let (proof, downloaded_source) = calls.get(&get.call_in_progress()).unwrap();
                    let value = match proof
                        .storage_value(&self.inner.warped_header_state_root, get.key().as_ref())
                    {
                        Ok(v) => v,
                        Err(proof_decode::IncompleteProofError { .. }) => {
                            return (
                                self.inner,
                                Err(BuildChainInformationError::SourceMisbehavior(
                                    SourceMisbehavior {
                                        source_id: *downloaded_source,
                                        error: SourceMisbehaviorTy::MerkleProofEntriesMissing,
                                    },
                                )),
                            );
                        }
                    };

                    get.inject_value(value.map(|(val, ver)| (iter::once(val), ver)))
                }
                chain_information::build::InProgress::NextKey(nk) => {
                    // TODO: child tries not supported
                    let (proof, downloaded_source) = calls.get(&nk.call_in_progress()).unwrap();
                    let value = match proof.next_key(
                        &self.inner.warped_header_state_root,
                        nk.key(),
                        nk.or_equal(),
                        nk.prefix(),
                        nk.branch_nodes(),
                    ) {
                        Ok(v) => v,
                        Err(proof_decode::IncompleteProofError { .. }) => {
                            return (
                                self.inner,
                                Err(BuildChainInformationError::SourceMisbehavior(
                                    SourceMisbehavior {
                                        source_id: *downloaded_source,
                                        error: SourceMisbehaviorTy::MerkleProofEntriesMissing,
                                    },
                                )),
                            );
                        }
                    };
                    nk.inject_key(value)
                }
                chain_information::build::InProgress::ClosestDescendantMerkleValue(mv) => {
                    // TODO: child tries not supported
                    let (proof, downloaded_source) = calls.get(&mv.call_in_progress()).unwrap();
                    let value = match proof.closest_descendant_merkle_value(
                        &self.inner.warped_header_state_root,
                        mv.key(),
                    ) {
                        Ok(v) => v,
                        Err(proof_decode::IncompleteProofError { .. }) => {
                            return (
                                self.inner,
                                Err(BuildChainInformationError::SourceMisbehavior(
                                    SourceMisbehavior {
                                        source_id: *downloaded_source,
                                        error: SourceMisbehaviorTy::MerkleProofEntriesMissing,
                                    },
                                )),
                            );
                        }
                    };
                    mv.inject_merkle_value(value)
                }
            };
        }
    }
}

/// Returns `true` if `a` and `b` are equal.
fn parameters_equal(mut a: &[u8], b: impl Iterator<Item = impl AsRef<[u8]>>) -> bool {
    for slice in b {
        let slice = slice.as_ref();

        if a.len() < slice.len() {
            return false;
        }

        if &a[..slice.len()] != slice {
            return false;
        }

        a = &a[slice.len()..];
    }

    true
}
