From cde9f0a5c889f0203331f409869c9a835f168852 Mon Sep 17 00:00:00 2001 From: Oleksandr Pravdyvyi Date: Thu, 15 May 2025 11:38:37 +0300 Subject: [PATCH] feat: sparse trees preparation for serialization --- accounts/src/account_core/mod.rs | 38 ++++-- common/src/nullifier_sparse_merkle_tree.rs | 129 +++++++++++++++++---- node_core/src/chain_storage/mod.rs | 26 ++++- node_core/src/executions/de.rs | 15 ++- node_core/src/executions/private_exec.rs | 15 ++- node_core/src/executions/se.rs | 15 ++- sc_core/src/proofs_circuits.rs | 15 ++- sequencer_core/src/lib.rs | 22 +++- utxo/src/utxo_tree.rs | 113 ++++++++++++++---- 9 files changed, 314 insertions(+), 74 deletions(-) diff --git a/accounts/src/account_core/mod.rs b/accounts/src/account_core/mod.rs index 2d3417c..afa0197 100644 --- a/accounts/src/account_core/mod.rs +++ b/accounts/src/account_core/mod.rs @@ -7,7 +7,7 @@ use log::info; use serde::Serialize; use utxo::{ utxo_core::{UTXOPayload, UTXO}, - utxo_tree::UTXOSparseMerkleTree, + utxo_tree::{UTXOSparseMerkleTree, UTXOTreeInput}, }; use crate::key_management::{ @@ -100,7 +100,7 @@ impl Account { Ok(()) } - pub fn add_new_utxo_outputs(&mut self, utxos: Vec) -> Result<()> { + pub fn add_new_utxo_outputs(&mut self, utxos: Vec) -> Result<()> { Ok(self.utxo_tree.insert_items(utxos)?) } @@ -113,6 +113,9 @@ impl Account { asset: Asset, amount: u128, privacy_flag: bool, + utxo_id: u64, + tx_id: u64, + block_id: u64, ) -> Result<()> { let payload_with_asset = UTXOPayload { owner: self.address, @@ -123,7 +126,14 @@ impl Account { let asset_utxo = UTXO::create_utxo_from_payload(payload_with_asset)?; - self.utxo_tree.insert_item(asset_utxo)?; + let input_utxo = UTXOTreeInput { + utxo_id, + tx_id, + block_id, + utxo: asset_utxo, + }; + + self.utxo_tree.insert_item(input_utxo)?; Ok(()) } @@ -163,14 +173,24 @@ mod tests { UTXONullifier::default() } - fn generate_dummy_utxo(address: TreeHashType, amount: u128) -> anyhow::Result { + fn generate_dummy_utxo( + address: TreeHashType, + amount: u128, + utxo_id: u64, + ) -> anyhow::Result { let payload = UTXOPayload { owner: address, asset: vec![], amount, privacy_flag: false, }; - UTXO::create_utxo_from_payload(payload) + let utxo = UTXO::create_utxo_from_payload(payload); + utxo.map(|utxo| UTXOTreeInput { + utxo_id, + tx_id: 1, + block_id: 1, + utxo, + }) } #[test] @@ -184,7 +204,7 @@ mod tests { #[test] fn test_mark_spent_utxo() { let mut account = Account::new(); - let utxo = generate_dummy_utxo(account.address, 100).unwrap(); + let utxo = generate_dummy_utxo(account.address, 100, 1).unwrap(); account.add_new_utxo_outputs(vec![utxo]).unwrap(); let mut utxo_nullifier_map = HashMap::new(); @@ -199,8 +219,8 @@ mod tests { #[test] fn test_add_new_utxo_outputs() { let mut account = Account::new(); - let utxo1 = generate_dummy_utxo(account.address, 100).unwrap(); - let utxo2 = generate_dummy_utxo(account.address, 200).unwrap(); + let utxo1 = generate_dummy_utxo(account.address, 100, 1).unwrap(); + let utxo2 = generate_dummy_utxo(account.address, 200, 2).unwrap(); let result = account.add_new_utxo_outputs(vec![utxo1.clone(), utxo2.clone()]); @@ -222,7 +242,7 @@ mod tests { let asset = "dummy_asset"; let amount = 1000u128; - let result = account.add_asset(asset, amount, false); + let result = account.add_asset(asset, amount, false, 1, 1, 1); assert!(result.is_ok()); assert_eq!(account.utxo_tree.store.len(), 1); diff --git a/common/src/nullifier_sparse_merkle_tree.rs b/common/src/nullifier_sparse_merkle_tree.rs index 282fccc..c6707a7 100644 --- a/common/src/nullifier_sparse_merkle_tree.rs +++ b/common/src/nullifier_sparse_merkle_tree.rs @@ -1,3 +1,5 @@ +use std::collections::BTreeMap; + use monotree::database::MemoryDB; use monotree::hasher::Blake3; use monotree::{Hasher, Monotree, Proof}; @@ -5,10 +7,31 @@ use monotree::{Hasher, Monotree, Proof}; use crate::merkle_tree_public::TreeHashType; use crate::nullifier::UTXONullifier; +#[derive(Debug, Clone)] +pub struct NullifierTreeInput { + pub nullifier_id: u64, + pub tx_id: u64, + pub block_id: u64, + pub nullifier: UTXONullifier, +} + +#[derive(Debug, Clone)] +pub struct TreeTxWithNullifierId { + pub id: u64, + pub nullifiers: BTreeMap, +} + +#[derive(Debug, Clone)] +pub struct TreeBlockWithTxId { + pub id: u64, + pub txs: BTreeMap, +} + pub struct NullifierSparseMerkleTree { pub curr_root: Option, pub tree: Monotree, pub hasher: Blake3, + pub leafs: BTreeMap, } impl NullifierSparseMerkleTree { @@ -17,30 +40,71 @@ impl NullifierSparseMerkleTree { curr_root: None, tree: Monotree::default(), hasher: Blake3::new(), + leafs: BTreeMap::new(), } } - pub fn insert_item(&mut self, nullifier: UTXONullifier) -> Result<(), monotree::Errors> { + pub fn modify_leavs_with_nullifier_input(&mut self, tree_nullifier: NullifierTreeInput) { + self.leafs + .entry(tree_nullifier.block_id) + .and_modify(|tree_block| { + tree_block + .txs + .entry(tree_nullifier.tx_id) + .and_modify(|tree_tx| { + tree_tx + .nullifiers + .insert(tree_nullifier.nullifier_id, tree_nullifier.nullifier); + }) + .or_insert(TreeTxWithNullifierId { + id: tree_nullifier.tx_id, + nullifiers: BTreeMap::new(), + }); + }) + .or_insert(TreeBlockWithTxId { + id: tree_nullifier.block_id, + txs: BTreeMap::new(), + }); + } + + pub fn insert_item( + &mut self, + tree_nullifier: NullifierTreeInput, + ) -> Result<(), monotree::Errors> { let root = self.curr_root.as_ref(); - let new_root = self - .tree - .insert(root, &nullifier.utxo_hash, &nullifier.utxo_hash)?; + let new_root = self.tree.insert( + root, + &tree_nullifier.nullifier.utxo_hash, + &tree_nullifier.nullifier.utxo_hash, + )?; self.curr_root = new_root; + self.modify_leavs_with_nullifier_input(tree_nullifier); + Ok(()) } - pub fn insert_items(&mut self, nullifiers: Vec) -> Result<(), monotree::Errors> { + pub fn insert_items( + &mut self, + tree_nullifiers: Vec, + ) -> Result<(), monotree::Errors> { let root = self.curr_root.as_ref(); - let hashes: Vec = nullifiers.iter().map(|nu| nu.utxo_hash).collect(); + let hashes: Vec = tree_nullifiers + .iter() + .map(|nu| nu.nullifier.utxo_hash) + .collect(); let new_root = self.tree.inserts(root, &hashes, &hashes)?; self.curr_root = new_root; + for tree_nullifier in tree_nullifiers { + self.modify_leavs_with_nullifier_input(tree_nullifier); + } + Ok(()) } @@ -130,6 +194,20 @@ mod tests { UTXONullifier { utxo_hash: hash } } + fn create_nullifier_input( + hash: TreeHashType, + nullifier_id: u64, + tx_id: u64, + block_id: u64, + ) -> NullifierTreeInput { + NullifierTreeInput { + nullifier_id, + tx_id, + block_id, + nullifier: create_nullifier(hash), + } + } + #[test] fn test_new_tree_initialization() { let tree = NullifierSparseMerkleTree::new(); @@ -139,9 +217,9 @@ mod tests { #[test] fn test_insert_single_item() { let mut tree = NullifierSparseMerkleTree::new(); - let nullifier = create_nullifier([1u8; 32]); // Sample 32-byte hash + let tree_nullifier = create_nullifier_input([1u8; 32], 1, 1, 1); // Sample 32-byte hash - let result = tree.insert_item(nullifier); + let result = tree.insert_item(tree_nullifier); assert!(result.is_ok()); assert!(tree.curr_root.is_some()); } @@ -149,13 +227,13 @@ mod tests { #[test] fn test_insert_multiple_items() { let mut tree = NullifierSparseMerkleTree::new(); - let nullifiers = vec![ - create_nullifier([1u8; 32]), - create_nullifier([2u8; 32]), - create_nullifier([3u8; 32]), + let tree_nullifiers = vec![ + create_nullifier_input([1u8; 32], 1, 1, 1), + create_nullifier_input([2u8; 32], 2, 1, 1), + create_nullifier_input([3u8; 32], 3, 1, 1), ]; - let result = tree.insert_items(nullifiers); + let result = tree.insert_items(tree_nullifiers); assert!(result.is_ok()); assert!(tree.curr_root.is_some()); } @@ -163,9 +241,9 @@ mod tests { #[test] fn test_search_item_inclusion() { let mut tree = NullifierSparseMerkleTree::new(); - let nullifier = create_nullifier([1u8; 32]); + let tree_nullifier = create_nullifier_input([1u8; 32], 1, 1, 1); - tree.insert_item(nullifier.clone()).unwrap(); + tree.insert_item(tree_nullifier.clone()).unwrap(); let result = tree.search_item_inclusion([1u8; 32]); assert!(result.is_ok()); @@ -179,13 +257,13 @@ mod tests { #[test] fn test_search_multiple_item_inclusions() { let mut tree = NullifierSparseMerkleTree::new(); - let nullifiers = vec![ - create_nullifier([1u8; 32]), - create_nullifier([2u8; 32]), - create_nullifier([3u8; 32]), + let tree_nullifiers = vec![ + create_nullifier_input([1u8; 32], 1, 1, 1), + create_nullifier_input([2u8; 32], 2, 1, 1), + create_nullifier_input([3u8; 32], 3, 1, 1), ]; - tree.insert_items(nullifiers).unwrap(); + tree.insert_items(tree_nullifiers).unwrap(); let search_hashes = vec![[1u8; 32], [2u8; 32], [99u8; 32]]; let result = tree.search_item_inclusions(&search_hashes); @@ -224,9 +302,9 @@ mod tests { #[test] fn test_insert_and_get_proof_of_existing_item() { let mut tree = NullifierSparseMerkleTree::new(); - let nullifier = create_nullifier([1u8; 32]); + let tree_nullifier = create_nullifier_input([1u8; 32], 1, 1, 1); - tree.insert_item(nullifier.clone()).unwrap(); + tree.insert_item(tree_nullifier.clone()).unwrap(); let proof_result = tree.get_non_membership_proof([1u8; 32]); assert!(proof_result.is_err()); @@ -235,9 +313,12 @@ mod tests { #[test] fn test_insert_and_get_proofs_of_existing_items() { let mut tree = NullifierSparseMerkleTree::new(); - let nullifiers = vec![create_nullifier([1u8; 32]), create_nullifier([2u8; 32])]; + let tree_nullifiers = vec![ + create_nullifier_input([1u8; 32], 1, 1, 1), + create_nullifier_input([2u8; 32], 2, 1, 1), + ]; - tree.insert_items(nullifiers).unwrap(); + tree.insert_items(tree_nullifiers).unwrap(); let proof_result = tree.get_non_membership_proofs(&[[1u8; 32], [2u8; 32]]); assert!(proof_result.is_err()); diff --git a/node_core/src/chain_storage/mod.rs b/node_core/src/chain_storage/mod.rs index f8bcf4e..3500c41 100644 --- a/node_core/src/chain_storage/mod.rs +++ b/node_core/src/chain_storage/mod.rs @@ -10,12 +10,12 @@ use common::{ block::Block, merkle_tree_public::merkle_tree::{PublicTransactionMerkleTree, UTXOCommitmentsMerkleTree}, nullifier::UTXONullifier, - nullifier_sparse_merkle_tree::NullifierSparseMerkleTree, + nullifier_sparse_merkle_tree::{NullifierSparseMerkleTree, NullifierTreeInput}, utxo_commitment::UTXOCommitment, }; use k256::AffinePoint; use public_context::PublicSCContext; -use utxo::utxo_core::UTXO; +use utxo::{utxo_core::UTXO, utxo_tree::UTXOTreeInput}; use crate::ActionData; @@ -54,7 +54,7 @@ impl NodeChainStore { } pub fn dissect_insert_block(&mut self, block: Block) -> Result<()> { - for tx in &block.transactions { + for (tx_id, tx) in block.transactions.iter().enumerate() { if !tx.execution_input.is_empty() { let public_action = serde_json::from_slice::(&tx.execution_input); @@ -101,7 +101,13 @@ impl NodeChainStore { tx.nullifier_created_hashes .clone() .into_iter() - .map(|hash| UTXONullifier { utxo_hash: hash }) + .enumerate() + .map(|(idx, hash)| NullifierTreeInput { + nullifier_id: idx as u64, + tx_id: tx_id as u64, + block_id: block.block_id, + nullifier: UTXONullifier { utxo_hash: hash }, + }) .collect(), )?; @@ -109,7 +115,9 @@ impl NodeChainStore { let ephemeral_public_key_sender = serde_json::from_slice::(&tx.ephemeral_pub_key)?; - for (ciphertext, nonce, tag) in tx.encoded_data.clone() { + for (utxo_id, (ciphertext, nonce, tag)) in + tx.encoded_data.clone().into_iter().enumerate() + { let slice = nonce.as_slice(); let nonce = accounts::key_management::constants_types::Nonce::clone_from_slice(slice); @@ -125,7 +133,13 @@ impl NodeChainStore { serde_json::from_slice::(&decoded_data_curr_acc); if let Ok(utxo) = decoded_utxo_try { if &utxo.owner == acc_id { - acc.utxo_tree.insert_item(utxo)?; + let input_utxo = UTXOTreeInput { + utxo_id: utxo_id as u64, + tx_id: tx_id as u64, + block_id: block.block_id, + utxo, + }; + acc.utxo_tree.insert_item(input_utxo)?; } } } diff --git a/node_core/src/executions/de.rs b/node_core/src/executions/de.rs index f73ba85..ddc5e81 100644 --- a/node_core/src/executions/de.rs +++ b/node_core/src/executions/de.rs @@ -1,4 +1,7 @@ +use std::collections::BTreeMap; + use bincode; +use common::nullifier_sparse_merkle_tree::NullifierTreeInput; use common::{ commitment::Commitment, commitments_sparse_merkle_tree::CommitmentsSparseMerkleTree, nullifier::UTXONullifier, nullifier_sparse_merkle_tree::NullifierSparseMerkleTree, @@ -106,15 +109,25 @@ pub fn validate_nullifiers_proof( root_nullifier: [u8; 32], nullifiers_proof: &[[u8; 32]], ) -> bool { + //There is no need for storage, so ids can be default there + let id = 1; + let mut nsmt = NullifierSparseMerkleTree { curr_root: Option::Some(root_nullifier), tree: Monotree::default(), hasher: Blake3::new(), + leafs: BTreeMap::new(), }; let nullifiers: Vec<_> = nullifiers_proof .into_iter() - .map(|n_p| UTXONullifier { utxo_hash: *n_p }) + .enumerate() + .map(|(idx, n_p)| NullifierTreeInput { + nullifier_id: idx as u64, + tx_id: id, + block_id: id, + nullifier: UTXONullifier { utxo_hash: *n_p }, + }) .collect(); nsmt.insert_items(nullifiers).unwrap(); diff --git a/node_core/src/executions/private_exec.rs b/node_core/src/executions/private_exec.rs index 2e06c86..e3d55a1 100644 --- a/node_core/src/executions/private_exec.rs +++ b/node_core/src/executions/private_exec.rs @@ -1,4 +1,7 @@ +use std::collections::BTreeMap; + use bincode; +use common::nullifier_sparse_merkle_tree::NullifierTreeInput; use common::{ commitment::Commitment, commitments_sparse_merkle_tree::CommitmentsSparseMerkleTree, nullifier::UTXONullifier, nullifier_sparse_merkle_tree::NullifierSparseMerkleTree, @@ -79,15 +82,25 @@ pub fn validate_nullifiers_proof( root_nullifier: [u8; 32], nullifiers_proof: &[[u8; 32]], ) -> bool { + //There is no need for storage, so ids can be default there + let id = 1; + let mut nsmt = NullifierSparseMerkleTree { curr_root: Option::Some(root_nullifier), tree: Monotree::default(), hasher: Blake3::new(), + leafs: BTreeMap::new(), }; let nullifiers: Vec<_> = nullifiers_proof .into_iter() - .map(|n_p| UTXONullifier { utxo_hash: *n_p }) + .enumerate() + .map(|(idx, n_p)| NullifierTreeInput { + nullifier_id: idx as u64, + tx_id: id, + block_id: id, + nullifier: UTXONullifier { utxo_hash: *n_p }, + }) .collect(); nsmt.insert_items(nullifiers).unwrap(); diff --git a/node_core/src/executions/se.rs b/node_core/src/executions/se.rs index d37fee5..f9123c2 100644 --- a/node_core/src/executions/se.rs +++ b/node_core/src/executions/se.rs @@ -1,4 +1,7 @@ +use std::collections::BTreeMap; + use bincode; +use common::nullifier_sparse_merkle_tree::NullifierTreeInput; use common::{ commitment::Commitment, commitments_sparse_merkle_tree::CommitmentsSparseMerkleTree, nullifier::UTXONullifier, nullifier_sparse_merkle_tree::NullifierSparseMerkleTree, @@ -108,15 +111,25 @@ pub fn validate_nullifiers_proof( root_nullifier: [u8; 32], nullifiers_proof: &[[u8; 32]], ) -> bool { + //There is no need for storage, so ids can be default there + let id = 1; + let mut nsmt = NullifierSparseMerkleTree { curr_root: Option::Some(root_nullifier), tree: Monotree::default(), hasher: Blake3::new(), + leafs: BTreeMap::new(), }; let nullifiers: Vec<_> = nullifiers_proof .into_iter() - .map(|n_p| UTXONullifier { utxo_hash: *n_p }) + .enumerate() + .map(|(idx, n_p)| NullifierTreeInput { + nullifier_id: idx as u64, + tx_id: id, + block_id: id, + nullifier: UTXONullifier { utxo_hash: *n_p }, + }) .collect(); nsmt.insert_items(nullifiers).unwrap(); diff --git a/sc_core/src/proofs_circuits.rs b/sc_core/src/proofs_circuits.rs index 2072e55..3574ec4 100644 --- a/sc_core/src/proofs_circuits.rs +++ b/sc_core/src/proofs_circuits.rs @@ -1,4 +1,7 @@ +use std::collections::BTreeMap; + use bincode; +use common::nullifier_sparse_merkle_tree::NullifierTreeInput; use common::{ commitment::Commitment, commitments_sparse_merkle_tree::CommitmentsSparseMerkleTree, nullifier::UTXONullifier, nullifier_sparse_merkle_tree::NullifierSparseMerkleTree, @@ -81,15 +84,25 @@ pub fn validate_nullifiers_proof( root_nullifier: [u8; 32], nullifiers_proof: &[[u8; 32]], ) -> bool { + //There is no need for storage, so ids can be default there + let id = 1; + let mut nsmt = NullifierSparseMerkleTree { curr_root: Option::Some(root_nullifier), tree: Monotree::default(), hasher: Blake3::new(), + leafs: BTreeMap::new(), }; let nullifiers: Vec<_> = nullifiers_proof .into_iter() - .map(|n_p| UTXONullifier { utxo_hash: *n_p }) + .enumerate() + .map(|(idx, n_p)| NullifierTreeInput { + nullifier_id: idx as u64, + tx_id: id, + block_id: id, + nullifier: UTXONullifier { utxo_hash: *n_p }, + }) .collect(); nsmt.insert_items(nullifiers).unwrap(); diff --git a/sequencer_core/src/lib.rs b/sequencer_core/src/lib.rs index d073ffb..647af9b 100644 --- a/sequencer_core/src/lib.rs +++ b/sequencer_core/src/lib.rs @@ -5,6 +5,7 @@ use common::{ block::{Block, HashableBlockData}, merkle_tree_public::TreeHashType, nullifier::UTXONullifier, + nullifier_sparse_merkle_tree::NullifierTreeInput, transaction::{Transaction, TxKind}, utxo_commitment::UTXOCommitment, }; @@ -185,6 +186,8 @@ impl SequencerCore { fn execute_check_transaction_on_state( &mut self, tx: TransactionMempool, + tx_id: u64, + block_id: u64, ) -> Result<(), TransactionMalformationErrorKind> { let Transaction { hash, @@ -199,11 +202,16 @@ impl SequencerCore { .add_tx(UTXOCommitment { hash: *utxo_comm }); } - for nullifier in nullifier_created_hashes { + for (idx, nullifier) in nullifier_created_hashes.iter().enumerate() { self.store .nullifier_store - .insert_item(UTXONullifier { - utxo_hash: *nullifier, + .insert_item(NullifierTreeInput { + nullifier_id: idx as u64, + tx_id, + block_id, + nullifier: UTXONullifier { + utxo_hash: *nullifier, + }, }) .map_err(|err| TransactionMalformationErrorKind::FailedToInsert { tx: hash, @@ -225,12 +233,14 @@ impl SequencerCore { ///Produces new block from transactions in mempool pub fn produce_new_block_with_mempool_transactions(&mut self) -> Result { + let new_block_height = self.chain_height + 1; + let transactions = self .mempool .pop_size(self.sequencer_config.max_num_tx_in_block); - for tx in transactions.clone() { - self.execute_check_transaction_on_state(tx)?; + for (idx, tx) in transactions.clone().into_iter().enumerate() { + self.execute_check_transaction_on_state(tx, idx as u64, new_block_height)?; } let prev_block_hash = self @@ -240,7 +250,7 @@ impl SequencerCore { .hash; let hashable_data = HashableBlockData { - block_id: self.chain_height + 1, + block_id: new_block_height, prev_block_id: self.chain_height, transactions: transactions.into_iter().map(|tx_mem| tx_mem.tx).collect(), data: vec![], diff --git a/utxo/src/utxo_tree.rs b/utxo/src/utxo_tree.rs index 3c0c603..cc84dd9 100644 --- a/utxo/src/utxo_tree.rs +++ b/utxo/src/utxo_tree.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use common::merkle_tree_public::TreeHashType; use monotree::database::MemoryDB; @@ -7,11 +7,32 @@ use monotree::{Hasher, Monotree, Proof}; use crate::utxo_core::UTXO; +#[derive(Debug, Clone)] +pub struct UTXOTreeInput { + pub utxo_id: u64, + pub tx_id: u64, + pub block_id: u64, + pub utxo: UTXO, +} + +#[derive(Debug, Clone)] +pub struct TreeTxWithUTXOId { + pub id: u64, + pub utxos: BTreeMap, +} + +#[derive(Debug, Clone)] +pub struct TreeBlockWithTxId { + pub id: u64, + pub txs: BTreeMap, +} + pub struct UTXOSparseMerkleTree { pub curr_root: Option, pub tree: Monotree, pub hasher: Blake3, pub store: HashMap, + pub leafs: BTreeMap, } impl UTXOSparseMerkleTree { @@ -21,30 +42,58 @@ impl UTXOSparseMerkleTree { tree: Monotree::default(), hasher: Blake3::new(), store: HashMap::new(), + leafs: BTreeMap::new(), } } - pub fn insert_item(&mut self, utxo: UTXO) -> Result<(), monotree::Errors> { + pub fn modify_leavs_with_nullifier_input(&mut self, tree_utxo: UTXOTreeInput) { + self.leafs + .entry(tree_utxo.block_id) + .and_modify(|tree_block| { + tree_block + .txs + .entry(tree_utxo.tx_id) + .and_modify(|tree_tx| { + tree_tx.utxos.insert(tree_utxo.utxo_id, tree_utxo.utxo); + }) + .or_insert(TreeTxWithUTXOId { + id: tree_utxo.tx_id, + utxos: BTreeMap::new(), + }); + }) + .or_insert(TreeBlockWithTxId { + id: tree_utxo.block_id, + txs: BTreeMap::new(), + }); + } + + pub fn insert_item(&mut self, tree_utxo: UTXOTreeInput) -> Result<(), monotree::Errors> { let root = self.curr_root.as_ref(); - let new_root = self.tree.insert(root, &utxo.hash, &utxo.hash)?; - - self.store.insert(utxo.hash, utxo); + let new_root = self + .tree + .insert(root, &tree_utxo.utxo.hash, &tree_utxo.utxo.hash)?; self.curr_root = new_root; + self.store + .insert(tree_utxo.utxo.hash, tree_utxo.utxo.clone()); + self.modify_leavs_with_nullifier_input(tree_utxo); + Ok(()) } - pub fn insert_items(&mut self, utxos: Vec) -> Result<(), monotree::Errors> { + pub fn insert_items(&mut self, tree_utxos: Vec) -> Result<(), monotree::Errors> { let root = self.curr_root.as_ref(); - let hashes: Vec = utxos.iter().map(|item| item.hash).collect(); + let hashes: Vec = tree_utxos.iter().map(|item| item.utxo.hash).collect(); let new_root = self.tree.inserts(root, &hashes, &hashes)?; - for utxo in utxos { - self.store.insert(utxo.hash, utxo); + for tree_utxo in tree_utxos { + self.store + .insert(tree_utxo.utxo.hash, tree_utxo.utxo.clone()); + self.modify_leavs_with_nullifier_input(tree_utxo); } self.curr_root = new_root; @@ -80,17 +129,31 @@ mod tests { use super::*; use crate::utxo_core::{UTXOPayload, UTXO}; - fn sample_utxo_payload() -> UTXOPayload { + fn sample_utxo_payload(amount: u128) -> UTXOPayload { UTXOPayload { owner: AccountId::default(), asset: vec![1, 2, 3], - amount: 10, + amount, privacy_flag: false, } } - fn sample_utxo() -> anyhow::Result { - UTXO::create_utxo_from_payload(sample_utxo_payload()) + fn sample_utxo(amount: u128) -> anyhow::Result { + UTXO::create_utxo_from_payload(sample_utxo_payload(amount)) + } + + fn sample_utxo_input( + utxo_id: u64, + tx_id: u64, + block_id: u64, + amount: u128, + ) -> anyhow::Result { + sample_utxo(amount).map(|utxo| UTXOTreeInput { + utxo_id, + tx_id, + block_id, + utxo, + }) } #[test] @@ -103,7 +166,7 @@ mod tests { #[test] fn test_insert_item() { let mut smt = UTXOSparseMerkleTree::new(); - let utxo = sample_utxo().unwrap(); + let utxo = sample_utxo_input(1, 1, 1, 10).unwrap(); let result = smt.insert_item(utxo.clone()); @@ -111,7 +174,7 @@ mod tests { assert!(result.is_ok()); // Test UTXO is now stored in the tree - assert_eq!(smt.store.get(&utxo.hash).unwrap().hash, utxo.hash); + assert_eq!(smt.store.get(&utxo.utxo.hash).unwrap().hash, utxo.utxo.hash); // Test curr_root is updated assert!(smt.curr_root.is_some()); @@ -120,8 +183,8 @@ mod tests { #[test] fn test_insert_items() { let mut smt = UTXOSparseMerkleTree::new(); - let utxo1 = sample_utxo().unwrap(); - let utxo2 = sample_utxo().unwrap(); + let utxo1 = sample_utxo_input(1, 1, 1, 10).unwrap(); + let utxo2 = sample_utxo_input(2, 1, 1, 11).unwrap(); let result = smt.insert_items(vec![utxo1.clone(), utxo2.clone()]); @@ -129,8 +192,8 @@ mod tests { assert!(result.is_ok()); // Test UTXOs are now stored in the tree - assert!(smt.store.get(&utxo1.hash).is_some()); - assert!(smt.store.get(&utxo2.hash).is_some()); + assert!(smt.store.get(&utxo1.utxo.hash).is_some()); + assert!(smt.store.get(&utxo2.utxo.hash).is_some()); // Test curr_root is updated assert!(smt.curr_root.is_some()); @@ -139,20 +202,20 @@ mod tests { #[test] fn test_get_item_exists() { let mut smt = UTXOSparseMerkleTree::new(); - let utxo = sample_utxo().unwrap(); + let utxo = sample_utxo_input(1, 1, 1, 10).unwrap(); smt.insert_item(utxo.clone()).unwrap(); // Test that the UTXO can be retrieved by hash - let retrieved_utxo = smt.get_item(utxo.hash).unwrap(); + let retrieved_utxo = smt.get_item(utxo.utxo.hash).unwrap(); assert!(retrieved_utxo.is_some()); - assert_eq!(retrieved_utxo.unwrap().hash, utxo.hash); + assert_eq!(retrieved_utxo.unwrap().hash, utxo.utxo.hash); } #[test] fn test_get_item_not_exists() { let mut smt = UTXOSparseMerkleTree::new(); - let utxo = sample_utxo().unwrap(); + let utxo = sample_utxo_input(1, 1, 1, 10).unwrap(); // Insert one UTXO and try to fetch a different hash smt.insert_item(utxo).unwrap(); @@ -167,12 +230,12 @@ mod tests { #[test] fn test_get_membership_proof() { let mut smt = UTXOSparseMerkleTree::new(); - let utxo = sample_utxo().unwrap(); + let utxo = sample_utxo_input(1, 1, 1, 10).unwrap(); smt.insert_item(utxo.clone()).unwrap(); // Fetch membership proof for the inserted UTXO - let proof = smt.get_membership_proof(utxo.hash).unwrap(); + let proof = smt.get_membership_proof(utxo.utxo.hash).unwrap(); // Test proof is generated successfully assert!(proof.is_some());