miden_client/store/data_store/

mod.rs

use alloc::boxed::Box;
use alloc::collections::{BTreeMap, BTreeSet};
use alloc::sync::Arc;
use alloc::vec::Vec;

use miden_protocol::account::{
    Account,
    AccountCode,
    AccountId,
    PartialAccount,
    StorageMapKey,
    StorageMapWitness,
    StorageSlot,
    StorageSlotContent,
    StorageSlotName,
};
use miden_protocol::asset::{AssetId, AssetVault, AssetWitness};
use miden_protocol::block::{BlockHeader, BlockNumber};
use miden_protocol::crypto::merkle::MerklePath;
use miden_protocol::crypto::merkle::mmr::{InOrderIndex, MmrPeaks, PartialMmr};
use miden_protocol::note::{NoteScript, NoteScriptRoot};
use miden_protocol::protocol_config::ProtocolConfig;
use miden_protocol::transaction::{AccountInputs, PartialBlockchain};
use miden_protocol::vm::FutureMaybeSend;
use miden_protocol::{Word, ZERO};
use miden_tx::{
    DataStore,
    DataStoreError,
    LoadedMastForest,
    MastForestStore,
    TransactionMastStore,
};

use super::{AccountStorageFilter, PartialBlockchainFilter, Store};
use crate::rpc::domain::account::{
    AccountStorageRequirements,
    GetAccountRequest,
    StorageMapEntries,
    StorageMapFetch,
    VaultFetch,
};
use crate::rpc::{AccountStateAt, NodeRpcClient};
use crate::store::StoreError;
use crate::transaction::{ChainAnchor, ChainAnchorError, fetch_public_account_inputs};

mod cache;
use cache::DataStoreCache;

// DATA STORE
// ================================================================================================

/// Wrapper structure that implements [`DataStore`] over any [`Store`].
pub struct ClientDataStore {
    /// Local database containing information about the accounts managed by this client.
    store: alloc::sync::Arc<dyn Store>,
    /// In-memory state served to the executor for the duration of the execution session.
    cache: DataStoreCache,
    /// RPC client used to lazy-load foreign account data on cache miss.
    rpc_api: Arc<dyn NodeRpcClient>,
    /// When set, chain data (reference block header and partial blockchain) is served from this
    /// anchor instead of being rebuilt at the store's sync height. Boxed to keep the data store
    /// small: it is held inline by every execution future.
    anchor: Option<Box<ChainAnchor>>,
}

impl ClientDataStore {
    pub fn new(store: alloc::sync::Arc<dyn Store>, rpc_api: Arc<dyn NodeRpcClient>) -> Self {
        Self {
            store,
            cache: DataStoreCache::new(),
            rpc_api,
            anchor: None,
        }
    }

    /// Serves chain data from the provided [`ChainAnchor`] instead of rebuilding it at the store's
    /// sync height, pinning execution to the anchor's reference block.
    ///
    /// The store's account data is still used as-is: only the reference block header and the
    /// partial blockchain come from the anchor. The anchor must track all blocks required by the
    /// request and all creation blocks of authenticated input notes. Otherwise,
    /// `get_transaction_inputs` fails.
    #[must_use]
    pub fn with_chain_anchor(mut self, anchor: ChainAnchor) -> Self {
        self.anchor = Some(Box::new(anchor));
        self
    }

    /// Enables memoization of `get_transaction_inputs` and `get_vault_asset_witnesses` for the
    /// lifetime of this data store.
    ///
    /// This is only correct when the account state served to the executor does not change between
    /// executions, as is the case while the [`crate::note::NoteScreener`] runs trial executions
    /// against the same accounts and reference block. It must stay disabled for real transaction
    /// execution, where account state evolves between executions.
    #[must_use]
    pub fn with_execution_input_cache(mut self) -> Self {
        self.cache.enable_execution_input_cache();
        self
    }

    pub fn mast_store(&self) -> Arc<TransactionMastStore> {
        self.cache.mast_store.clone()
    }

    /// Stores the provided foreign account inputs so they can be served to the executor upon
    /// request.
    pub fn register_foreign_account_inputs(
        &self,
        foreign_accounts: impl IntoIterator<Item = AccountInputs>,
    ) {
        self.cache.replace_foreign_account_inputs(foreign_accounts);
    }

    /// Stores the blocks that the current transaction must be able to authenticate.
    pub(crate) fn register_block_numbers(
        &self,
        block_numbers: impl IntoIterator<Item = BlockNumber>,
    ) {
        self.cache.replace_block_numbers(block_numbers);
    }

    /// Registers note scripts so they can be served to the executor upon request.
    ///
    /// Scripts accumulate across calls (they are not cleared) so that a data store reused for
    /// several executions — e.g. by [`crate::transaction::BatchBuilder`] — keeps serving the
    /// scripts registered for earlier transactions.
    pub fn register_note_scripts(&self, note_scripts: impl IntoIterator<Item = NoteScript>) {
        self.cache.insert_note_scripts(note_scripts);
    }

    /// Attempts to resolve a storage map witness from the local store.
    ///
    /// This covers any account present in the store (local or foreign) as well as any foreign
    /// account previously cached in `foreign_account_inputs`.
    ///
    /// Returns `Ok(None)` when the map is not found locally.
    async fn get_local_storage_map_witness(
        &self,
        account_id: AccountId,
        map_root: Word,
        map_key: StorageMapKey,
    ) -> Result<Option<StorageMapWitness>, DataStoreError> {
        match self
            .store
            .get_account_storage(account_id, AccountStorageFilter::Root(map_root))
            .await
        {
            Ok(account_storage) => {
                match account_storage.slots().first().map(StorageSlot::content) {
                    Some(StorageSlotContent::Map(map)) => Ok(Some(map.open(&map_key))),
                    Some(StorageSlotContent::Value(value)) => Err(DataStoreError::other(format!(
                        "found StorageSlotContent::Value with {value} as its value."
                    ))),
                    _ => Ok(None),
                }
            },
            Err(err) => {
                tracing::debug!(
                    %account_id,
                    %err,
                    "storage map not found locally, will try remote fetch"
                );
                Ok(None)
            },
        }
    }

    /// Lazily fetches a foreign account's inputs from the network, loads its code into the MAST
    /// store, and caches the result in [`Self::foreign_account_inputs`].
    async fn fetch_and_cache_foreign_account(
        &self,
        account_id: AccountId,
        account_state_at: AccountStateAt,
    ) -> Result<AccountInputs, DataStoreError> {
        let account_inputs = fetch_public_account_inputs(
            &self.store,
            &self.rpc_api,
            account_id,
            AccountStorageRequirements::default(),
            account_state_at,
        )
        .await
        .map_err(|err| {
            DataStoreError::other_with_source("failed to fetch foreign account inputs", err)
        })?;

        self.cache.mast_store.load_account_code(account_inputs.code());
        self.cache.insert_foreign_account_inputs(account_inputs.clone());

        Ok(account_inputs)
    }

    /// Fetches a storage map witness for a specific key from the network via RPC and caches it.
    ///
    /// Anchored at the transaction reference block: a chain-tip query would return proofs for a
    /// newer map root whenever the account changed after the caller's last sync.
    async fn fetch_and_cache_storage_map_witness(
        &self,
        account_id: AccountId,
        map_root: Word,
        slot_name: StorageSlotName,
        map_key: StorageMapKey,
        known_code: AccountCode,
    ) -> Result<StorageMapWitness, DataStoreError> {
        let account_state_at =
            self.cache.ref_block().map_or(AccountStateAt::ChainTip, AccountStateAt::Block);

        let storage_requirements = AccountStorageRequirements::new([(slot_name, &[map_key])]);
        let (_, account_proof): (BlockNumber, _) = self
            .rpc_api
            .get_account(
                account_id,
                GetAccountRequest::new()
                    .with_storage(StorageMapFetch::Slots(storage_requirements))
                    .with_known_code(Some(known_code))
                    .at(account_state_at),
            )
            .await
            .map_err(|err| {
                DataStoreError::other_with_source("failed to fetch storage map via RPC", err)
            })?;

        let (_, account_details) = account_proof.into_parts();
        let details = account_details.ok_or_else(|| {
            DataStoreError::other(format!(
                "RPC returned no account details for account {account_id}"
            ))
        })?;

        let map_detail =
            details.storage_details.map_details.into_iter().next().ok_or_else(|| {
                DataStoreError::other(format!(
                    "RPC returned no storage map details for account {account_id}"
                ))
            })?;

        let StorageMapEntries::PartialMap { partial_smt, .. } = map_detail.entries else {
            return Err(DataStoreError::other(
                "expected a partial storage map in response to a specific-key request",
            ));
        };

        // Reject a wrong-root response here rather than as an opaque merkle error inside the VM.
        // The whole tree shares one root, so this covers every opening taken from it.
        let map_detail_root = partial_smt.root();
        if map_detail_root != map_root {
            return Err(DataStoreError::other(format!(
                "storage map fetched for account {account_id} verifies against root \
                 {map_detail_root} but the executor requires root {map_root}"
            )));
        }

        let proof = partial_smt.open(&map_key.hash().as_word()).map_err(|err| {
            DataStoreError::other_with_source("failed to open the requested storage map key", err)
        })?;

        let witness = StorageMapWitness::new(proof, [map_key]).map_err(|err| {
            DataStoreError::other_with_source("failed to create storage map witness", err)
        })?;
        self.cache.insert_storage_map_witness(map_root, map_key, witness.clone());
        Ok(witness)
    }

    /// Fetches an account's full vault via RPC — anchored at the transaction reference block — and
    /// verifies it against the vault root the executor requires. Fallback for vault reads the local
    /// store cannot serve, typically foreign accounts whose [`AccountInputs`] carry only their
    /// vault root.
    async fn fetch_vault_via_rpc(
        &self,
        account_id: AccountId,
        vault_root: Word,
    ) -> Result<AssetVault, DataStoreError> {
        let account_state_at =
            self.cache.ref_block().map_or(AccountStateAt::ChainTip, AccountStateAt::Block);

        // The cached foreign inputs hold the code, letting the node omit it from the response.
        let known_code = self
            .cache
            .with_foreign_account_inputs(account_id, |inputs| inputs.code().clone());

        let (block_num, mut account_proof) = self
            .rpc_api
            .get_account(
                account_id,
                GetAccountRequest::new()
                    .at(account_state_at)
                    .with_known_code(known_code)
                    .with_vault(VaultFetch::Always),
            )
            .await
            .map_err(|err| {
                DataStoreError::other_with_source("failed to fetch account vault via RPC", err)
            })?;

        let details = account_proof.details_mut().ok_or_else(|| {
            DataStoreError::other(format!(
                "RPC returned no account details for account {account_id}"
            ))
        })?;

        self.rpc_api
            .resolve_oversize_vault(account_id, block_num, details)
            .await
            .map_err(|err| {
                DataStoreError::other_with_source("failed to resolve oversize vault via RPC", err)
            })?;

        let vault = AssetVault::new(&details.vault_details.assets).map_err(|err| {
            DataStoreError::other_with_source("failed to build the fetched vault", err)
        })?;

        if vault.root() != vault_root {
            return Err(DataStoreError::other(format!(
                "vault fetched for account {account_id} has root {} but the executor requires \
                 root {vault_root}",
                vault.root()
            )));
        }

        Ok(vault)
    }
}

impl DataStore for ClientDataStore {
    async fn get_transaction_inputs(
        &self,
        account_id: AccountId,
        mut block_refs: BTreeSet<BlockNumber>,
    ) -> Result<(PartialAccount, BlockHeader, ProtocolConfig, PartialBlockchain), DataStoreError>
    {
        // Last block is used as reference (it does not need to be authenticated manually)
        let ref_block = *block_refs.last().ok_or(DataStoreError::other("block set is empty"))?;

        for block_num in self.cache.block_numbers() {
            if block_num > ref_block {
                return Err(DataStoreError::other(format!(
                    "requested block {block_num} is after transaction reference block {ref_block}"
                )));
            }
            block_refs.insert(block_num);
        }

        // Cache the reference block so lazy-loading methods can use it
        self.cache.set_ref_block(ref_block);

        let partial_account =
            if let Some(partial_account) = self.cache.get_partial_account(account_id) {
                partial_account
            } else {
                let partial_account_record = self
                    .store
                    .get_minimal_partial_account(account_id)
                    .await?
                    .ok_or(DataStoreError::AccountNotFound(account_id))?;

                // New accounts (nonce == 0) need full storage maps as advice inputs for the kernel
                // to validate during account creation. For these, fetch the full account and
                // convert to PartialAccount (which includes full storage for new accounts).
                // Existing accounts use the minimal partial record directly.
                let partial_account: PartialAccount = if partial_account_record.nonce() == ZERO {
                    let full_record = self
                        .store
                        .get_account(account_id)
                        .await?
                        .ok_or(DataStoreError::AccountNotFound(account_id))?;
                    let account: Account = full_record
                        .try_into()
                        .map_err(|_| DataStoreError::AccountNotFound(account_id))?;
                    PartialAccount::from(&account)
                } else {
                    partial_account_record
                        .try_into()
                        .map_err(|_| DataStoreError::AccountNotFound(account_id))?
                };

                self.cache.insert_partial_account(&partial_account);
                partial_account
            };

        let (block_header, partial_blockchain) = if let Some(anchor) = &self.anchor {
            // Anchored execution: serve the pinned chain data. The executor-derived reference block
            // must match the anchor. The anchor's partial blockchain must track every other block
            // in the set.
            if ref_block != anchor.block_num() {
                return Err(DataStoreError::other_with_source(
                    "anchored data store cannot serve the requested reference block",
                    ChainAnchorError::ReferenceBlockMismatch {
                        requested: ref_block,
                        anchor: anchor.block_num(),
                    },
                ));
            }

            for block_num in block_refs.iter().filter(|block_num| **block_num != ref_block) {
                if !anchor.partial_blockchain().contains_block(*block_num) {
                    return Err(DataStoreError::other_with_source(
                        "anchored data store cannot serve an untracked block",
                        ChainAnchorError::BlockNotTracked { block_num: *block_num },
                    ));
                }
            }

            (anchor.header().clone(), anchor.partial_blockchain().clone())
        } else if let Some((block_header, partial_blockchain)) =
            self.cache.get_blockchain(&block_refs)
        {
            (block_header, partial_blockchain)
        } else {
            // The full set identifies the served blockchain, so keep it as the cache key before the
            // reference block is removed from it below.
            let cache_key = block_refs.clone();
            block_refs.remove(&ref_block);

            let current_peaks = self.store.get_current_blockchain_peaks().await?;

            let (block_header, _had_notes) = self
                .store
                .get_block_header_by_num(ref_block)
                .await?
                .ok_or(DataStoreError::BlockNotFound(ref_block))?;

            // TODO: the client stores only the peaks of the MMR at the current sync height, so we
            // are not actually following the block_ref here. If the block_ref !=
            // current_sync_height, this would return an invalid partial blockchain.
            let (partial_mmr, block_headers) = build_partial_mmr_and_headers_with_fallback(
                &self.store,
                &self.rpc_api,
                current_peaks,
                &block_refs,
            )
            .await?;

            let partial_blockchain =
                PartialBlockchain::new(partial_mmr, block_headers).map_err(|err| {
                    DataStoreError::other_with_source(
                        "error creating PartialBlockchain from internal data",
                        err,
                    )
                })?;

            self.cache.insert_blockchain(cache_key, &block_header, &partial_blockchain);
            (block_header, partial_blockchain)
        };

        let protocol_config = crate::protocol_config::load_protocol_config(
            self.store.as_ref(),
            block_header.protocol_config_commitment(),
        )
        .await?;
        Ok((partial_account, block_header, protocol_config, partial_blockchain))
    }

    /// Retrieves witnesses for the requested assets from the local store, falling back to a single
    /// RPC vault fetch when the store cannot serve the requested root.
    ///
    /// Assets absent from the vault are served too: the store returns an emptiness proof for them,
    /// which the executor needs when an asset is being added to the vault.
    async fn get_vault_asset_witnesses(
        &self,
        account_id: AccountId,
        vault_root: Word,
        asset_ids: BTreeSet<AssetId>,
    ) -> Result<Vec<AssetWitness>, DataStoreError> {
        if let Some(witnesses) = self.cache.get_vault_asset_witnesses(vault_root, &asset_ids) {
            return Ok(witnesses);
        }

        let asset_witnesses = match self
            .store
            .get_vault_asset_witnesses(account_id, vault_root, asset_ids.clone())
            .await
        {
            Ok(witnesses) => witnesses,
            Err(err) => {
                tracing::debug!(
                    %account_id,
                    requested_root = %vault_root,
                    %err,
                    "local store cannot serve the requested vault root, will fetch it via RPC"
                );
                let vault = self.fetch_vault_via_rpc(account_id, vault_root).await?;
                asset_ids.iter().copied().map(|asset_id| vault.open(asset_id)).collect()
            },
        };

        self.cache
            .insert_vault_asset_witnesses(vault_root, &asset_ids, &asset_witnesses);
        Ok(asset_witnesses)
    }

    /// Retrieves the [`StorageMapWitness`] requested from the store. Alternatively fetching it from
    /// the RPC if not available locally. Witnesses fetched via RPC are cached in memory so that
    /// repeated accesses to the same map entry within a transaction avoid additional RPC calls.
    async fn get_storage_map_witness(
        &self,
        account_id: AccountId,
        map_root: Word,
        map_key: StorageMapKey,
    ) -> Result<StorageMapWitness, DataStoreError> {
        // Check the in-memory witness cache first.
        if let Some(witness) = self.cache.get_storage_map_witness(map_root, map_key) {
            return Ok(witness);
        }

        // Try the local store.
        if let Some(witness) =
            self.get_local_storage_map_witness(account_id, map_root, map_key).await?
        {
            return Ok(witness);
        }

        // Resolve against the cached account inputs (without cloning them), fetching and caching
        // the account first if it isn't cached yet.
        let resolution = if let Some(resolution) =
            self.cache.with_foreign_account_inputs(account_id, |inputs| {
                resolve_witness_from_inputs(inputs, map_root, map_key)
            }) {
            resolution?
        } else {
            let account_state_at = self
                .cache
                .ref_block()
                .map(AccountStateAt::Block)
                .expect("reference block should be set");
            let inputs = self.fetch_and_cache_foreign_account(account_id, account_state_at).await?;
            resolve_witness_from_inputs(&inputs, map_root, map_key)?
        };

        match resolution {
            WitnessResolution::Witness(witness) => Ok(witness),
            WitnessResolution::FetchParams(slot_name, known_code) => {
                self.fetch_and_cache_storage_map_witness(
                    account_id, map_root, slot_name, map_key, known_code,
                )
                .await
            },
        }
    }

    /// Returns the [`AccountInputs`] for the given foreign account from the cache or alternatively
    /// fetching them from the RPC if not available locally.
    async fn get_foreign_account_inputs(
        &self,
        foreign_account_id: AccountId,
        ref_block: BlockNumber,
    ) -> Result<AccountInputs, DataStoreError> {
        // Fast path: check the cache first.
        if let Some(inputs) = self.cache.get_foreign_account_inputs(foreign_account_id) {
            return Ok(inputs);
        }

        self.fetch_and_cache_foreign_account(foreign_account_id, AccountStateAt::Block(ref_block))
            .await
    }

    /// Returns the [`NoteScript`] for the given script root from the registered session scripts,
    /// the store, or alternatively fetching it from the RPC if not available locally.
    fn get_note_script(
        &self,
        script_root: NoteScriptRoot,
    ) -> impl FutureMaybeSend<Result<Option<NoteScript>, DataStoreError>> {
        let registered_script = self.cache.get_note_script(script_root.into());
        let store = self.store.clone();
        let rpc_api = self.rpc_api.clone();

        async move {
            // Fastest path: scripts registered for the in-flight transaction request.
            if let Some(note_script) = registered_script {
                return Ok(Some(note_script));
            }

            // Fast path: check the local store first.
            match store.get_note_script(script_root.into()).await {
                Ok(note_script) => return Ok(Some(note_script)),
                Err(StoreError::NoteScriptNotFound(_)) => {},
                Err(err) => {
                    return Err(DataStoreError::other_with_source(
                        format!("failed to get note script {script_root} from store"),
                        err,
                    ));
                },
            }

            // Store miss, fetch from the network via RPC.
            let Some(note_script) =
                rpc_api.get_note_script_by_root(script_root.into()).await.map_err(|err| {
                    DataStoreError::other_with_source("failed to fetch note script via RPC", err)
                })?
            else {
                return Ok(None);
            };

            // Persist for future lookups.
            if let Err(err) = store.upsert_note_scripts(core::slice::from_ref(&note_script)).await {
                tracing::warn!(
                    %err,
                    "Failed to persist fetched note script to store"
                );
            }

            Ok(Some(note_script))
        }
    }
}

// MAST FOREST STORE
// ================================================================================================

impl MastForestStore for ClientDataStore {
    fn get(&self, procedure_hash: &Word) -> Option<LoadedMastForest> {
        self.cache.mast_store.get(procedure_hash)
    }
}

// HELPER FUNCTIONS
// ================================================================================================

/// Outcome of resolving a storage map witness against an account's inputs: either the witness
/// itself, or the parameters needed to fetch it via RPC.
enum WitnessResolution {
    Witness(StorageMapWitness),
    /// The [`AccountCode`] is not needed to build the witness: it is only sent along with the RPC
    /// request so the node can omit the account code from its response.
    FetchParams(StorageSlotName, AccountCode),
}

/// Tries to open the witness from the inputs' partial storage maps (this can miss if the account's
/// storage is too big); on a miss, resolves the slot name and account code needed to fetch the
/// witness via RPC.
fn resolve_witness_from_inputs(
    inputs: &AccountInputs,
    map_root: Word,
    map_key: StorageMapKey,
) -> Result<WitnessResolution, DataStoreError> {
    if let Some(partial_map) = inputs.storage().maps().find(|m| m.root() == map_root)
        && let Ok(witness) = partial_map.open(&map_key)
    {
        return Ok(WitnessResolution::Witness(witness));
    }

    let account_id = inputs.id();
    let slot_name = inputs
        .storage()
        .header()
        .slots()
        .find(|slot| slot.slot_type().is_map() && slot.value() == map_root)
        .map(|slot| slot.name().clone())
        .ok_or_else(|| {
            DataStoreError::other(format!(
                "did not find map slot with root {map_root} for foreign account {account_id}"
            ))
        })?;

    Ok(WitnessResolution::FetchParams(slot_name, inputs.code().clone()))
}

/// Builds a [`PartialMmr`] from the given peaks and a list of blocks that should be authenticated
/// against them.
///
/// `authenticated_blocks` must not contain the block whose forest matches `peaks`. For that block
/// the kernel extends the MMR itself, so an authentication path is not needed.
pub(crate) async fn build_partial_mmr_with_paths(
    store: &alloc::sync::Arc<dyn Store>,
    rpc_api: &Arc<dyn NodeRpcClient>,
    peaks: MmrPeaks,
    authenticated_blocks: &[BlockHeader],
) -> Result<PartialMmr, DataStoreError> {
    let mut partial_mmr: PartialMmr = PartialMmr::from_peaks(peaks);

    let block_nums: Vec<BlockNumber> =
        authenticated_blocks.iter().map(BlockHeader::block_num).collect();

    let authentication_paths =
        get_authentication_path_for_blocks(store, &block_nums, partial_mmr.forest().num_leaves())
            .await?;

    for (header, local_path) in authenticated_blocks.iter().zip(authentication_paths.iter()) {
        if partial_mmr
            .track(header.block_num().as_usize(), header.commitment(), local_path)
            .is_ok()
        {
            continue;
        }

        fetch_and_track_block_header(rpc_api, &mut partial_mmr, header.block_num(), Some(header))
            .await?;
    }

    Ok(partial_mmr)
}

/// Builds a [`PartialMmr`] and returns the authenticated block headers.
///
/// Headers in the store use local MMR nodes when possible. If a local path is incomplete, the node
/// returns the header and its proof in one call. A missing header always uses this combined call.
pub(crate) async fn build_partial_mmr_and_headers_with_fallback(
    store: &alloc::sync::Arc<dyn Store>,
    rpc_api: &Arc<dyn NodeRpcClient>,
    peaks: MmrPeaks,
    block_numbers: &BTreeSet<BlockNumber>,
) -> Result<(PartialMmr, Vec<BlockHeader>), DataStoreError> {
    let mut headers: BTreeMap<BlockNumber, BlockHeader> = store
        .get_block_headers(block_numbers)
        .await?
        .into_iter()
        .map(|(header, _has_notes)| (header.block_num(), header))
        .collect();

    let local_headers: Vec<BlockHeader> = headers.values().cloned().collect();
    let mut partial_mmr =
        build_partial_mmr_with_paths(store, rpc_api, peaks, &local_headers).await?;

    for &block_num in block_numbers {
        if headers.contains_key(&block_num) {
            continue;
        }

        let header =
            fetch_and_track_block_header(rpc_api, &mut partial_mmr, block_num, None).await?;
        headers.insert(block_num, header);
    }

    Ok((partial_mmr, headers.into_values().collect()))
}

/// Fetches a block header with its MMR proof and tracks the authenticated header.
///
/// If `expected_header` is set, the fetched header must match it.
async fn fetch_and_track_block_header(
    rpc_api: &Arc<dyn NodeRpcClient>,
    partial_mmr: &mut PartialMmr,
    block_num: BlockNumber,
    expected_header: Option<&BlockHeader>,
) -> Result<BlockHeader, DataStoreError> {
    let (header, proof) = rpc_api.get_block_header_with_proof(block_num).await.map_err(|err| {
        DataStoreError::other_with_source(
            format!("failed to fetch block header and MMR proof for block {block_num}"),
            err,
        )
    })?;

    if header.block_num() != block_num {
        return Err(DataStoreError::other(format!(
            "node returned block header {} for requested block {block_num}",
            header.block_num()
        )));
    }
    if expected_header.is_some_and(|expected| header != *expected) {
        return Err(DataStoreError::other(format!(
            "node returned a different header for block {block_num}"
        )));
    }
    if proof.leaf() != header.commitment() {
        return Err(DataStoreError::other(format!(
            "node returned an invalid MMR proof for block {block_num}"
        )));
    }

    let proof = proof.with_forest(partial_mmr.forest()).map_err(|err| {
        DataStoreError::other(format!("failed to adjust MMR proof for block {block_num}: {err}"))
    })?;
    partial_mmr
        .track(block_num.as_usize(), header.commitment(), proof.merkle_path())
        .map_err(|err| DataStoreError::other(format!("error constructing MMR: {err}")))?;

    Ok(header)
}

/// Retrieves all Partial Blockchain nodes required for authenticating the set of blocks, and then
/// constructs the path for each of them.
///
/// This function assumes `block_nums` doesn't contain values above or equal to `forest`. If there
/// are any such values, the function will panic when calling `mmr_merkle_path_len()`.
async fn get_authentication_path_for_blocks(
    store: &alloc::sync::Arc<dyn Store>,
    block_nums: &[BlockNumber],
    forest: usize,
) -> Result<Vec<MerklePath>, StoreError> {
    let mut node_indices = BTreeSet::new();

    // Calculate all needed nodes indices for generating the paths
    for block_num in block_nums {
        let path_depth = mmr_merkle_path_len(block_num.as_usize(), forest);

        let mut idx = InOrderIndex::from_leaf_pos(block_num.as_usize());

        for _ in 0..path_depth {
            node_indices.insert(idx.sibling());
            idx = idx.parent();
        }
    }

    let node_indices: Vec<InOrderIndex> = node_indices.into_iter().collect();

    let filter = PartialBlockchainFilter::List(node_indices);
    let mmr_nodes = store.get_partial_blockchain_nodes(filter).await?;

    // Construct authentication paths
    let mut authentication_paths = vec![];
    for block_num in block_nums {
        let mut merkle_nodes = vec![];
        let mut idx = InOrderIndex::from_leaf_pos(block_num.as_usize());

        while let Some(node) = mmr_nodes.get(&idx.sibling()) {
            merkle_nodes.push(*node);
            idx = idx.parent();
        }
        let path = MerklePath::new(merkle_nodes);
        authentication_paths.push(path);
    }

    Ok(authentication_paths)
}

/// Calculates the merkle path length for an MMR of a specific forest and a leaf index `leaf_index`
/// is a 0-indexed leaf number and `forest` is the total amount of leaves in the MMR at this point.
fn mmr_merkle_path_len(leaf_index: usize, forest: usize) -> usize {
    let before: usize = forest & leaf_index;
    let after = forest ^ before;

    after.ilog2() as usize
}