miden_objects/asset/vault/

partial.rs

use alloc::string::ToString;

use miden_crypto::merkle::{InnerNodeInfo, MerkleError, PartialSmt, SmtLeaf, SmtProof};

use super::AssetVault;
use crate::Word;
use crate::asset::Asset;
use crate::errors::PartialAssetVaultError;
use crate::utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};

/// A partial representation of an [`AssetVault`], containing only proofs for a subset of assets.
///
/// Partial vault is used to provide verifiable access to specific assets in a vault
/// without the need to provide the full vault data. It contains all required data for loading
/// vault data into the transaction kernel for transaction execution.
#[derive(Clone, Debug, PartialEq, Eq, Default)]
pub struct PartialVault {
    /// An SMT with a partial view into an account's full [`AssetVault`].
    partial_smt: PartialSmt,
}

impl PartialVault {
    // CONSTRUCTORS
    // --------------------------------------------------------------------------------------------

    /// Returns a new instance of a partial vault from the provided partial SMT.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - the provided SMT does not track only valid [`Asset`]s.
    /// - the vault key at which the asset is stored does not match the vault key derived from the
    ///   asset.
    pub fn new(partial_smt: PartialSmt) -> Result<Self, PartialAssetVaultError> {
        Self::validate_entries(partial_smt.entries())?;

        Ok(PartialVault { partial_smt })
    }

    // ACCESSORS
    // --------------------------------------------------------------------------------------------

    /// Returns the root of the partial vault.
    pub fn root(&self) -> Word {
        self.partial_smt.root()
    }

    /// Returns an iterator over all inner nodes in the Sparse Merkle Tree proofs.
    ///
    /// This is useful for reconstructing parts of the Sparse Merkle Tree or for
    /// verification purposes.
    pub fn inner_nodes(&self) -> impl Iterator<Item = InnerNodeInfo> + '_ {
        self.partial_smt.inner_nodes()
    }

    /// Returns an iterator over all leaves in the Sparse Merkle Tree proofs.
    ///
    /// Each item returned is a tuple containing the leaf index and a reference to the leaf.
    pub fn leaves(&self) -> impl Iterator<Item = &SmtLeaf> {
        self.partial_smt.leaves().map(|(_, leaf)| leaf)
    }

    /// Returns the [`Asset`] associated with the given `vault_key`.
    ///
    /// The return value is `None` if the asset does not exist in the vault.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - the key is not tracked by this partial SMT.
    pub fn get(&self, vault_key: Word) -> Result<Option<Asset>, MerkleError> {
        self.partial_smt.get_value(&vault_key).map(|word| {
            if word.is_empty() {
                None
            } else {
                // SAFETY: If this returned a non-empty word, then it should be a valid asset,
                // because the vault should only track valid ones.
                Some(Asset::try_from(word).expect("partial vault should only track valid assets"))
            }
        })
    }

    // MUTATORS
    // --------------------------------------------------------------------------------------------

    /// Adds an [`SmtProof`] to this [`PartialVault`].
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - the provided proof does not prove inclusion of valid [`Asset`]s.
    /// - the vault key of the proven asset does not match the vault key derived from the asset.
    /// - the new root after the insertion of the leaf and the path does not match the existing root
    ///   (except when the first leaf is added).
    pub fn add(&mut self, proof: SmtProof) -> Result<(), PartialAssetVaultError> {
        Self::validate_entries(proof.leaf().entries())?;

        self.partial_smt
            .add_proof(proof)
            .map_err(PartialAssetVaultError::FailedToAddProof)
    }

    // HELPER FUNCTIONS
    // --------------------------------------------------------------------------------------------

    /// Validates that the provided entries are valid vault keys and assets.
    ///
    /// For brevity, the error conditions are only mentioned on the public methods that use this
    /// function.
    fn validate_entries<'a>(
        entries: impl IntoIterator<Item = &'a (Word, Word)>,
    ) -> Result<(), PartialAssetVaultError> {
        for (vault_key, asset) in entries {
            let asset = Asset::try_from(asset).map_err(|source| {
                PartialAssetVaultError::InvalidAssetInSmt { entry: *asset, source }
            })?;

            if asset.vault_key() != *vault_key {
                return Err(PartialAssetVaultError::VaultKeyMismatch {
                    expected: asset.vault_key(),
                    actual: *vault_key,
                });
            }
        }

        Ok(())
    }
}

impl From<&AssetVault> for PartialVault {
    fn from(value: &AssetVault) -> Self {
        let vault_partial_smt = value.asset_tree.clone().into();

        PartialVault { partial_smt: vault_partial_smt }
    }
}

impl Serializable for PartialVault {
    fn write_into<W: ByteWriter>(&self, target: &mut W) {
        target.write(&self.partial_smt)
    }
}

impl Deserializable for PartialVault {
    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
        let vault_partial_smt = source.read()?;

        PartialVault::new(vault_partial_smt)
            .map_err(|err| DeserializationError::InvalidValue(err.to_string()))
    }
}

// TESTS
// ================================================================================================

#[cfg(test)]
mod tests {
    use assert_matches::assert_matches;
    use miden_crypto::merkle::Smt;

    use super::*;
    use crate::asset::FungibleAsset;

    #[test]
    fn partial_vault_ensures_asset_validity() -> anyhow::Result<()> {
        let invalid_asset = Word::from([0, 0, 0, 5u32]);
        let smt = Smt::with_entries([(invalid_asset, invalid_asset)])?;
        let proof = smt.open(&invalid_asset);
        let partial_smt = PartialSmt::from_proofs([proof.clone()])?;

        let err = PartialVault::new(partial_smt).unwrap_err();
        assert_matches!(err, PartialAssetVaultError::InvalidAssetInSmt { entry, .. } => {
            assert_eq!(entry, invalid_asset);
        });

        let err = PartialVault::default().add(proof).unwrap_err();
        assert_matches!(err, PartialAssetVaultError::InvalidAssetInSmt { entry, .. } => {
            assert_eq!(entry, invalid_asset);
        });

        Ok(())
    }

    #[test]
    fn partial_vault_ensures_asset_vault_key_matches() -> anyhow::Result<()> {
        let asset = FungibleAsset::mock(500);
        let invalid_vault_key = Word::from([0, 1, 2, 3u32]);
        let smt = Smt::with_entries([(invalid_vault_key, asset.into())])?;
        let proof = smt.open(&invalid_vault_key);
        let partial_smt = PartialSmt::from_proofs([proof.clone()])?;

        let err = PartialVault::new(partial_smt).unwrap_err();
        assert_matches!(err, PartialAssetVaultError::VaultKeyMismatch { expected, actual } => {
            assert_eq!(actual, invalid_vault_key);
            assert_eq!(expected, asset.vault_key());
        });

        let err = PartialVault::default().add(proof).unwrap_err();
        assert_matches!(err, PartialAssetVaultError::VaultKeyMismatch { expected, actual } => {
            assert_eq!(actual, invalid_vault_key);
            assert_eq!(expected, asset.vault_key());
        });

        Ok(())
    }
}