miden_client/sync/

state_sync.rs

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

use async_trait::async_trait;
use miden_protocol::Word;
use miden_protocol::account::{Account, AccountHeader, AccountId};
use miden_protocol::block::{BlockHeader, BlockNumber};
use miden_protocol::crypto::merkle::mmr::{InOrderIndex, MmrDelta, MmrPeaks, PartialMmr};
use miden_protocol::note::{NoteId, NoteTag};
use tracing::info;

use super::state_sync_update::TransactionUpdateTracker;
use super::{AccountUpdates, StateSyncUpdate};
use crate::ClientError;
use crate::note::NoteUpdateTracker;
use crate::rpc::NodeRpcClient;
use crate::rpc::domain::note::CommittedNote;
use crate::rpc::domain::sync::StateSyncInfo;
use crate::rpc::domain::transaction::TransactionInclusion;
use crate::store::{InputNoteRecord, OutputNoteRecord, StoreError};
use crate::transaction::TransactionRecord;

// SYNC CALLBACKS
// ================================================================================================

/// The action to be taken when a note update is received as part of the sync response.
#[allow(clippy::large_enum_variant)]
pub enum NoteUpdateAction {
    /// The note commit update is relevant and the specified note should be marked as committed in
    /// the store, storing its inclusion proof.
    Commit(CommittedNote),
    /// The public note is relevant and should be inserted into the store.
    Insert(InputNoteRecord),
    /// The note update is not relevant and should be discarded.
    Discard,
}

#[async_trait(?Send)]
pub trait OnNoteReceived {
    /// Callback that gets executed when a new note is received as part of the sync response.
    ///
    /// It receives:
    ///
    /// - The committed note received from the network.
    /// - An optional note record that corresponds to the state of the note in the network (only if
    ///   the note is public).
    ///
    /// It returns an enum indicating the action to be taken for the received note update. Whether
    /// the note updated should be committed, new public note inserted, or ignored.
    async fn on_note_received(
        &self,
        committed_note: CommittedNote,
        public_note: Option<InputNoteRecord>,
    ) -> Result<NoteUpdateAction, ClientError>;
}

// STATE SYNC
// ================================================================================================

/// The state sync components encompasses the client's sync logic. It is then used to request
/// updates from the node and apply them to the relevant elements. The updates are then returned and
/// can be applied to the store to persist the changes.
///
/// When created it receives a callback that will be executed when a new note inclusion is received
/// in the sync response.
#[derive(Clone)]
pub struct StateSync {
    /// The RPC client used to communicate with the node.
    rpc_api: Arc<dyn NodeRpcClient>,
    /// Responsible for checking the relevance of notes and executing the
    /// [`OnNoteReceived`] callback when a new note inclusion is received.
    note_screener: Arc<dyn OnNoteReceived>,
    /// The number of blocks that are considered old enough to discard pending transactions. If
    /// `None`, there is no limit and transactions will be kept indefinitely.
    tx_graceful_blocks: Option<u32>,
}

impl StateSync {
    /// Creates a new instance of the state sync component.
    ///
    /// # Arguments
    ///
    /// * `rpc_api` - The RPC client used to communicate with the node.
    /// * `on_note_received` - A callback to be executed when a new note inclusion is received.
    /// * `tx_graceful_blocks` - The number of blocks that are considered old enough to discard.
    /// * `note_screener` - The note screener used to check the relevance of notes.
    pub fn new(
        rpc_api: Arc<dyn NodeRpcClient>,
        note_screener: Arc<dyn OnNoteReceived>,
        tx_graceful_blocks: Option<u32>,
    ) -> Self {
        Self {
            rpc_api,
            note_screener,
            tx_graceful_blocks,
        }
    }

    /// Syncs the state of the client with the chain tip of the node, returning the updates that
    /// should be applied to the store.
    ///
    /// During the sync process, the client will go through the following steps:
    /// 1. A request is sent to the node to get the state updates. This request includes tracked
    ///    account IDs and the tags of notes that might have changed or that might be of interest to
    ///    the client.
    /// 2. A response is received with the current state of the network. The response includes
    ///    information about new and committed notes, updated accounts, and committed transactions.
    /// 3. Tracked public accounts are updated and private accounts are validated against the node
    ///    state.
    /// 4. Tracked notes are updated with their new states. Notes might be committed or nullified
    ///    during the sync processing.
    /// 5. New notes are checked, and only relevant ones are stored. Relevance is determined by the
    ///    [`OnNoteReceived`] callback.
    /// 6. Transactions are updated with their new states. Transactions might be committed or
    ///    discarded.
    /// 7. The MMR is updated with the new peaks and authentication nodes.
    ///
    /// # Arguments
    /// * `current_partial_blockchain` - The current partial view of the blockchain.
    /// * `accounts` - All the headers of tracked accounts.
    /// * `note_tags` - The note tags to be used in the sync state request.
    /// * `unspent_input_notes` - The current state of unspent input notes tracked by the client.
    /// * `unspent_output_notes` - The current state of unspent output notes tracked by the client.
    pub async fn sync_state(
        &self,
        mut current_partial_mmr: PartialMmr,
        accounts: Vec<AccountHeader>,
        note_tags: BTreeSet<NoteTag>,
        unspent_input_notes: Vec<InputNoteRecord>,
        unspent_output_notes: Vec<OutputNoteRecord>,
        uncommitted_transactions: Vec<TransactionRecord>,
    ) -> Result<StateSyncUpdate, ClientError> {
        let block_num = u32::try_from(
            current_partial_mmr.forest().num_leaves().checked_sub(1).unwrap_or_default(),
        )
        .map_err(|_| ClientError::InvalidPartialMmrForest)?
        .into();

        let mut state_sync_update = StateSyncUpdate {
            block_num,
            note_updates: NoteUpdateTracker::new(unspent_input_notes, unspent_output_notes),
            transaction_updates: TransactionUpdateTracker::new(uncommitted_transactions),
            ..Default::default()
        };

        let note_tags = Arc::new(note_tags);
        let account_ids: Vec<AccountId> = accounts.iter().map(AccountHeader::id).collect();
        let mut state_sync_steps = Vec::new();

        while let Some(step) = self
            .sync_state_step(state_sync_update.block_num, &account_ids, &note_tags)
            .await?
        {
            let sync_block_num = step.block_header.block_num();

            let reached_tip = step.chain_tip == sync_block_num;

            state_sync_update.block_num = sync_block_num;
            state_sync_steps.push(step);

            if reached_tip {
                break;
            }
        }

        // TODO: fetch_public_note_details should take an iterator or btreeset down to the RPC call
        // (this would be a breaking change so it should be done separately)
        let public_note_ids: Vec<NoteId> = state_sync_steps
            .iter()
            .flat_map(|s| s.note_inclusions.iter())
            .filter(|n| !n.metadata().is_private())
            .map(|n| *n.note_id())
            .collect();

        let public_note_records = self.fetch_public_note_details(&public_note_ids).await?;

        // Collect account commitment updates across all sync steps. Each account only needs
        // to be checked once since GetAccount always returns the latest state.
        let merged_commitment_updates: Vec<(AccountId, Word)> = state_sync_steps
            .iter()
            .flat_map(|s| s.account_commitment_updates.iter())
            .map(|(id, w)| (*id, *w))
            .collect::<BTreeMap<_, _>>()
            .into_iter()
            .collect();

        self.account_state_sync(
            &mut state_sync_update.account_updates,
            &accounts,
            &merged_commitment_updates,
        )
        .await?;

        // Apply local changes. These involve updating the MMR and applying state transitions
        // to notes based on the received information.
        info!("Applying state transitions locally.");

        for sync_step in state_sync_steps {
            let StateSyncInfo {
                chain_tip,
                block_header,
                mmr_delta,
                note_inclusions,
                transactions,
                ..
            } = sync_step;

            self.transaction_state_sync(
                &mut state_sync_update.transaction_updates,
                &block_header,
                &transactions,
            );

            let found_relevant_note = self
                .note_state_sync(
                    &mut state_sync_update.note_updates,
                    note_inclusions,
                    &block_header,
                    &public_note_records,
                )
                .await?;

            let (new_mmr_peaks, new_authentication_nodes) = apply_mmr_changes(
                &block_header,
                found_relevant_note,
                &mut current_partial_mmr,
                mmr_delta,
            )?;

            let include_block = found_relevant_note || chain_tip == block_header.block_num();
            if include_block {
                state_sync_update.block_updates.insert(
                    block_header,
                    found_relevant_note,
                    new_mmr_peaks,
                    new_authentication_nodes,
                );
            } else {
                // Even though this block header is not stored, `apply_mmr_changes` may
                // produce authentication nodes for already-tracked leaves whose Merkle
                // paths change as the MMR grows. These must be persisted so that the
                // `PartialMmr` can be correctly reconstructed from the store after a
                // client restart.
                state_sync_update
                    .block_updates
                    .extend_authentication_nodes(new_authentication_nodes);
            }
        }
        info!("Syncing nullifiers.");

        self.sync_nullifiers(&mut state_sync_update, block_num).await?;

        Ok(state_sync_update)
    }

    /// Executes a single sync request and returns the node response if the chain advanced.
    ///
    /// This method issues a `/SyncState` call starting from `current_block_num`. If the node
    /// reports the same block number, `None` is returned, signalling that the client is already at
    /// the requested height. Otherwise the full [`StateSyncInfo`] is returned for deferred
    /// processing by the caller.
    async fn sync_state_step(
        &self,
        current_block_num: BlockNumber,
        account_ids: &[AccountId],
        note_tags: &Arc<BTreeSet<NoteTag>>,
    ) -> Result<Option<StateSyncInfo>, ClientError> {
        info!("Performing sync state step.");
        let response = self
            .rpc_api
            .sync_state(current_block_num, account_ids, note_tags.as_ref())
            .await?;

        // We don't need to continue if the chain has not advanced, there are no new changes
        if response.block_header.block_num() == current_block_num {
            return Ok(None);
        }

        Ok(Some(response))
    }

    // HELPERS
    // --------------------------------------------------------------------------------------------

    /// Compares the state of tracked accounts with the updates received from the node. The method
    /// updates the `state_sync_update` field with the details of the accounts that need to be
    /// updated.
    ///
    /// The account updates might include:
    /// * Public accounts that have been updated in the node.
    /// * Network accounts that have been updated in the node and are being tracked by the client.
    /// * Private accounts that have been marked as mismatched because the current commitment
    ///   doesn't match the one received from the node. The client will need to handle these cases
    ///   as they could be a stale account state or a reason to lock the account.
    async fn account_state_sync(
        &self,
        account_updates: &mut AccountUpdates,
        accounts: &[AccountHeader],
        account_commitment_updates: &[(AccountId, Word)],
    ) -> Result<(), ClientError> {
        let (public_accounts, private_accounts): (Vec<_>, Vec<_>) =
            accounts.iter().partition(|account_header| !account_header.id().is_private());

        let updated_public_accounts = self
            .get_updated_public_accounts(account_commitment_updates, &public_accounts)
            .await?;

        let mismatched_private_accounts = account_commitment_updates
            .iter()
            .filter(|(account_id, digest)| {
                private_accounts
                    .iter()
                    .any(|account| account.id() == *account_id && &account.commitment() != digest)
            })
            .copied()
            .collect::<Vec<_>>();

        account_updates
            .extend(AccountUpdates::new(updated_public_accounts, mismatched_private_accounts));

        Ok(())
    }

    /// Queries the node for the latest state of the public accounts that don't match the current
    /// state of the client.
    async fn get_updated_public_accounts(
        &self,
        account_updates: &[(AccountId, Word)],
        current_public_accounts: &[&AccountHeader],
    ) -> Result<Vec<Account>, ClientError> {
        let mut mismatched_public_accounts = vec![];

        for (id, commitment) in account_updates {
            // check if this updated account state is tracked by the client
            if let Some(account) = current_public_accounts
                .iter()
                .find(|acc| *id == acc.id() && *commitment != acc.commitment())
            {
                mismatched_public_accounts.push(*account);
            }
        }

        self.rpc_api
            .get_updated_public_accounts(&mismatched_public_accounts)
            .await
            .map_err(ClientError::RpcError)
    }

    /// Applies the changes received from the sync response to the notes and transactions tracked
    /// by the client and updates the `note_updates` accordingly.
    ///
    /// This method uses the callbacks provided to the [`StateSync`] component to check if the
    /// updates received are relevant to the client.
    ///
    /// The note updates might include:
    /// * New notes that we received from the node and might be relevant to the client.
    /// * Tracked expected notes that were committed in the block.
    /// * Tracked notes that were being processed by a transaction that got committed.
    /// * Tracked notes that were nullified by an external transaction.
    ///
    /// The `public_notes` parameter provides cached public note details for the current sync
    /// iteration so the node is only queried once per batch.
    async fn note_state_sync(
        &self,
        note_updates: &mut NoteUpdateTracker,
        note_inclusions: Vec<CommittedNote>,
        block_header: &BlockHeader,
        public_notes: &BTreeMap<NoteId, InputNoteRecord>,
    ) -> Result<bool, ClientError> {
        // `found_relevant_note` tracks whether we want to persist the block header in the end
        let mut found_relevant_note = false;

        for committed_note in note_inclusions {
            let public_note = (!committed_note.metadata().is_private())
                .then(|| public_notes.get(committed_note.note_id()))
                .flatten()
                .cloned();

            match self.note_screener.on_note_received(committed_note, public_note).await? {
                NoteUpdateAction::Commit(committed_note) => {
                    // Only mark the downloaded block header as relevant if we are talking about
                    // an input note (output notes get marked as committed but we don't need the
                    // block for anything there)
                    found_relevant_note |= note_updates
                        .apply_committed_note_state_transitions(&committed_note, block_header)?;
                },
                NoteUpdateAction::Insert(public_note) => {
                    found_relevant_note = true;

                    note_updates.apply_new_public_note(public_note, block_header)?;
                },
                NoteUpdateAction::Discard => {},
            }
        }

        Ok(found_relevant_note)
    }

    /// Queries the node for all received notes that aren't being locally tracked in the client.
    ///
    /// The client can receive metadata for private notes that it's not tracking. In this case,
    /// notes are ignored for now as they become useless until details are imported.
    async fn fetch_public_note_details(
        &self,
        query_notes: &[NoteId],
    ) -> Result<BTreeMap<NoteId, InputNoteRecord>, ClientError> {
        if query_notes.is_empty() {
            return Ok(BTreeMap::new());
        }

        info!("Getting note details for notes that are not being tracked.");

        let return_notes = self.rpc_api.get_public_note_records(query_notes, None).await?;

        Ok(return_notes.into_iter().map(|note| (note.id(), note)).collect())
    }

    /// Collects the nullifier tags for the notes that were updated in the sync response and uses
    /// the `sync_nullifiers` endpoint to check if there are new nullifiers for these
    /// notes. It then processes the nullifiers to apply the state transitions on the note updates.
    ///
    /// The `state_sync_update` parameter will be updated to track the new discarded transactions.
    async fn sync_nullifiers(
        &self,
        state_sync_update: &mut StateSyncUpdate,
        current_block_num: BlockNumber,
    ) -> Result<(), ClientError> {
        // To receive information about added nullifiers, we reduce them to the higher 16 bits
        // Note that besides filtering by nullifier prefixes, the node also filters by block number
        // (it only returns nullifiers from current_block_num until
        // response.block_header.block_num())

        // Check for new nullifiers for input notes that were updated
        let nullifiers_tags: Vec<u16> = state_sync_update
            .note_updates
            .unspent_nullifiers()
            .map(|nullifier| nullifier.prefix())
            .collect();

        let mut new_nullifiers = self
            .rpc_api
            .sync_nullifiers(&nullifiers_tags, current_block_num, Some(state_sync_update.block_num))
            .await?;

        // Discard nullifiers that are newer than the current block (this might happen if the block
        // changes between the sync_state and the check_nullifier calls)
        new_nullifiers.retain(|update| update.block_num <= state_sync_update.block_num);

        for nullifier_update in new_nullifiers {
            state_sync_update.note_updates.apply_nullifiers_state_transitions(
                &nullifier_update,
                state_sync_update.transaction_updates.committed_transactions(),
            )?;

            // Process nullifiers and track the updates of local tracked transactions that were
            // discarded because the notes that they were processing were nullified by an
            // another transaction.
            state_sync_update
                .transaction_updates
                .apply_input_note_nullified(nullifier_update.nullifier);
        }

        Ok(())
    }

    /// Applies the changes received from the sync response to the transactions tracked by the
    /// client and updates the `transaction_updates` accordingly.
    ///
    /// The transaction updates might include:
    /// * New transactions that were committed in the block.
    /// * Transactions that were discarded because they were stale or expired.
    fn transaction_state_sync(
        &self,
        transaction_updates: &mut TransactionUpdateTracker,
        new_block_header: &BlockHeader,
        transaction_inclusions: &[TransactionInclusion],
    ) {
        for transaction_inclusion in transaction_inclusions {
            transaction_updates.apply_transaction_inclusion(
                transaction_inclusion,
                u64::from(new_block_header.timestamp()),
            ); //TODO: Change timestamps from u64 to u32
        }

        transaction_updates
            .apply_sync_height_update(new_block_header.block_num(), self.tx_graceful_blocks);
    }
}

// HELPERS
// ================================================================================================

/// Applies changes to the current MMR structure, returns the updated [`MmrPeaks`] and the
/// authentication nodes for leaves we track.
fn apply_mmr_changes(
    new_block: &BlockHeader,
    new_block_has_relevant_notes: bool,
    current_partial_mmr: &mut PartialMmr,
    mmr_delta: MmrDelta,
) -> Result<(MmrPeaks, Vec<(InOrderIndex, Word)>), ClientError> {
    // Apply the MMR delta to bring MMR to forest equal to chain tip
    let mut new_authentication_nodes: Vec<(InOrderIndex, Word)> =
        current_partial_mmr.apply(mmr_delta).map_err(StoreError::MmrError)?;

    let new_peaks = current_partial_mmr.peaks();

    new_authentication_nodes
        .append(&mut current_partial_mmr.add(new_block.commitment(), new_block_has_relevant_notes));

    Ok((new_peaks, new_authentication_nodes))
}