This commit is contained in:
Sergio Chouhy 2025-05-16 19:59:51 -03:00
parent 25ef949a97
commit f50ef5be9a
6 changed files with 30 additions and 218 deletions

View File

@ -1,4 +1,4 @@
use std::collections::HashMap;
use std::{collections::HashMap, hash::Hash};
use anyhow::Result;
use common::{merkle_tree_public::TreeHashType, nullifier::UTXONullifier, transaction::Tag};
@ -7,7 +7,6 @@ use log::info;
use serde::Serialize;
use utxo::{
utxo_core::{UTXOPayload, UTXO},
utxo_tree::UTXOSparseMerkleTree,
};
use crate::key_management::{
@ -23,7 +22,7 @@ pub struct Account {
pub key_holder: AddressKeyHolder,
pub address: AccountAddress,
pub balance: u64,
pub utxo_tree: UTXOSparseMerkleTree,
pub utxos: HashMap<TreeHashType, UTXO>,
}
///A strucure, which represents all the visible(public) information
@ -42,26 +41,26 @@ impl Account {
let key_holder = AddressKeyHolder::new_os_random();
let address = key_holder.address;
let balance = 0;
let utxo_tree = UTXOSparseMerkleTree::new();
let utxos = HashMap::new();
Self {
key_holder,
address,
balance,
utxo_tree,
utxos,
}
}
pub fn new_with_balance(balance: u64) -> Self {
let key_holder = AddressKeyHolder::new_os_random();
let address = key_holder.address;
let utxo_tree = UTXOSparseMerkleTree::new();
let utxos = HashMap::new();
Self {
key_holder,
address,
balance,
utxo_tree,
utxos,
}
}
@ -92,7 +91,7 @@ impl Account {
utxo_nullifier_map: HashMap<TreeHashType, UTXONullifier>,
) -> Result<()> {
for (hash, nullifier) in utxo_nullifier_map {
if let Some(utxo_entry) = self.utxo_tree.store.get_mut(&hash) {
if let Some(utxo_entry) = self.utxos.get_mut(&hash) {
utxo_entry.consume_utxo(nullifier)?;
}
}
@ -101,7 +100,13 @@ impl Account {
}
pub fn add_new_utxo_outputs(&mut self, utxos: Vec<UTXO>) -> Result<()> {
Ok(self.utxo_tree.insert_items(utxos)?)
for utxo in utxos {
if self.utxos.contains_key(&utxo.hash) {
return Err(anyhow::anyhow!("UTXO already exists"));
}
self.utxos.insert(utxo.hash, utxo);
}
return Ok(());
}
pub fn update_public_balance(&mut self, new_balance: u64) {
@ -123,7 +128,7 @@ impl Account {
let asset_utxo = UTXO::create_utxo_from_payload(payload_with_asset)?;
self.utxo_tree.insert_item(asset_utxo)?;
self.utxos.insert(asset_utxo.hash, asset_utxo);
Ok(())
}
@ -193,7 +198,7 @@ mod tests {
let result = account.mark_spent_utxo(utxo_nullifier_map);
assert!(result.is_ok());
assert!(account.utxo_tree.store.get(&account.address).is_none());
assert!(account.utxos.get(&account.address).is_none());
}
#[test]
@ -205,7 +210,7 @@ mod tests {
let result = account.add_new_utxo_outputs(vec![utxo1.clone(), utxo2.clone()]);
assert!(result.is_ok());
assert_eq!(account.utxo_tree.store.len(), 2);
assert_eq!(account.utxos.store.len(), 2);
}
#[test]
@ -225,6 +230,6 @@ mod tests {
let result = account.add_asset(asset, amount, false);
assert!(result.is_ok());
assert_eq!(account.utxo_tree.store.len(), 1);
assert_eq!(account.utxos.store.len(), 1);
}
}

View File

@ -125,7 +125,7 @@ impl NodeChainStore {
serde_json::from_slice::<UTXO>(&decoded_data_curr_acc);
if let Ok(utxo) = decoded_utxo_try {
if &utxo.owner == acc_id {
acc.utxo_tree.insert_item(utxo)?;
acc.utxos.insert_item(utxo)?;
}
}
}

View File

@ -1062,7 +1062,7 @@ impl NodeCore {
let acc = write_guard.acc_map.get_mut(&acc_addr).unwrap();
acc.utxo_tree.get_item(new_utxo_hash)?.unwrap().clone()
acc.utxos.get_item(new_utxo_hash)?.unwrap().clone()
};
new_utxo.log();
@ -1102,7 +1102,7 @@ impl NodeCore {
let acc = write_guard.acc_map.get_mut(&acc_addr).unwrap();
let new_utxo = acc
.utxo_tree
.utxos
.get_item(new_utxo_hash)
.unwrap()
.unwrap()
@ -1238,7 +1238,7 @@ impl NodeCore {
let acc = write_guard.acc_map.get_mut(&acc_addr_rec).unwrap();
acc.log();
acc.utxo_tree.get_item(new_utxo_hash)?.unwrap().clone()
acc.utxos.get_item(new_utxo_hash)?.unwrap().clone()
};
new_utxo.log();
info!(
@ -1278,7 +1278,7 @@ impl NodeCore {
let acc = write_guard.acc_map.get_mut(&acc_addr_rec).unwrap();
acc.log();
acc.utxo_tree.get_item(new_utxo_hash)?.unwrap().clone()
acc.utxos.get_item(new_utxo_hash)?.unwrap().clone()
};
new_utxo.log();
info!(
@ -1323,7 +1323,7 @@ impl NodeCore {
let acc = write_guard.acc_map.get_mut(&acc_addr_rec).unwrap();
acc.log();
let new_utxo = acc.utxo_tree.get_item(new_utxo_hash)?.unwrap().clone();
let new_utxo = acc.utxos.get_item(new_utxo_hash)?.unwrap().clone();
new_utxo.log();
info!(
@ -1343,7 +1343,7 @@ impl NodeCore {
let acc = write_guard.acc_map.get_mut(&acc_addr).unwrap();
acc.log();
let new_utxo = acc.utxo_tree.get_item(new_utxo_hash)?.unwrap().clone();
let new_utxo = acc.utxos.get_item(new_utxo_hash)?.unwrap().clone();
new_utxo.log();
info!(
@ -1557,7 +1557,7 @@ impl NodeCore {
let acc = write_guard.acc_map.get_mut(&acc_addr_rec).unwrap();
let new_utxo = acc
.utxo_tree
.utxos
.get_item(new_utxo_hash)
.unwrap()
.unwrap()

View File

@ -268,7 +268,7 @@ impl JsonHandler {
.ok_or(RpcError::new_internal_error(None, ACCOUNT_NOT_FOUND))?;
let utxo = acc
.utxo_tree
.utxos
.get_item(utxo_hash)
.map_err(|err| {
RpcError::new_internal_error(None, &format!("DB fetch failure {err:?}"))
@ -512,7 +512,7 @@ impl JsonHandler {
.get_mut(&acc_addr_sender)
.ok_or(RpcError::new_internal_error(None, ACCOUNT_NOT_FOUND))?;
acc.utxo_tree
acc.utxos
.get_item(utxo_hash)
.map_err(|err| {
RpcError::new_internal_error(None, &format!("DB fetch failure {err:?}"))
@ -647,7 +647,7 @@ impl JsonHandler {
.get_mut(&acc_addr_sender)
.ok_or(RpcError::new_internal_error(None, ACCOUNT_NOT_FOUND))?;
acc.utxo_tree
acc.utxos
.get_item(utxo_hash)
.map_err(|err| {
RpcError::new_internal_error(None, &format!("DB fetch failure {err:?}"))
@ -735,7 +735,7 @@ impl JsonHandler {
.get_mut(&acc_addr_sender)
.ok_or(RpcError::new_internal_error(None, ACCOUNT_NOT_FOUND))?;
acc.utxo_tree
acc.utxos
.get_item(utxo_hash)
.map_err(|err| {
RpcError::new_internal_error(None, &format!("DB fetch failure {err:?}"))

View File

@ -1,2 +1 @@
pub mod utxo_core;
pub mod utxo_tree;

View File

@ -1,192 +0,0 @@
use std::collections::HashMap;
use common::merkle_tree_public::TreeHashType;
use monotree::database::MemoryDB;
use monotree::hasher::Blake3;
use monotree::{Hasher, Monotree, Proof};
use crate::utxo_core::UTXO;
pub struct UTXOSparseMerkleTree {
pub curr_root: Option<TreeHashType>,
pub tree: Monotree<MemoryDB, Blake3>,
pub hasher: Blake3,
pub store: HashMap<TreeHashType, UTXO>,
}
impl UTXOSparseMerkleTree {
pub fn new() -> Self {
UTXOSparseMerkleTree {
curr_root: None,
tree: Monotree::default(),
hasher: Blake3::new(),
store: HashMap::new(),
}
}
pub fn insert_item(&mut self, utxo: UTXO) -> 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);
self.curr_root = new_root;
Ok(())
}
pub fn insert_items(&mut self, utxos: Vec<UTXO>) -> Result<(), monotree::Errors> {
let root = self.curr_root.as_ref();
let hashes: Vec<TreeHashType> = utxos.iter().map(|item| item.hash).collect();
let new_root = self.tree.inserts(root, &hashes, &hashes)?;
for utxo in utxos {
self.store.insert(utxo.hash, utxo);
}
self.curr_root = new_root;
Ok(())
}
pub fn get_item(&mut self, hash: TreeHashType) -> Result<Option<&UTXO>, monotree::Errors> {
let hash = self.tree.get(self.curr_root.as_ref(), &hash)?;
Ok(hash.and_then(|hash| self.store.get(&hash)))
}
pub fn get_membership_proof(
&mut self,
nullifier_hash: TreeHashType,
) -> Result<Option<Proof>, monotree::Errors> {
self.tree
.get_merkle_proof(self.curr_root.as_ref(), &nullifier_hash)
}
}
impl Default for UTXOSparseMerkleTree {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use common::AccountId;
use super::*;
use crate::utxo_core::{UTXOPayload, UTXO};
fn sample_utxo_payload() -> UTXOPayload {
UTXOPayload {
owner: AccountId::default(),
asset: vec![1, 2, 3],
amount: 10,
privacy_flag: false,
}
}
fn sample_utxo() -> anyhow::Result<UTXO> {
UTXO::create_utxo_from_payload(sample_utxo_payload())
}
#[test]
fn test_utxo_sparse_merkle_tree_new() {
let smt = UTXOSparseMerkleTree::new();
assert!(smt.curr_root.is_none());
assert_eq!(smt.store.len(), 0);
}
#[test]
fn test_insert_item() {
let mut smt = UTXOSparseMerkleTree::new();
let utxo = sample_utxo().unwrap();
let result = smt.insert_item(utxo.clone());
// Test insertion is successful
assert!(result.is_ok());
// Test UTXO is now stored in the tree
assert_eq!(smt.store.get(&utxo.hash).unwrap().hash, utxo.hash);
// Test curr_root is updated
assert!(smt.curr_root.is_some());
}
#[test]
fn test_insert_items() {
let mut smt = UTXOSparseMerkleTree::new();
let utxo1 = sample_utxo().unwrap();
let utxo2 = sample_utxo().unwrap();
let result = smt.insert_items(vec![utxo1.clone(), utxo2.clone()]);
// Test insertion of multiple items is successful
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());
// Test curr_root is updated
assert!(smt.curr_root.is_some());
}
#[test]
fn test_get_item_exists() {
let mut smt = UTXOSparseMerkleTree::new();
let utxo = sample_utxo().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();
assert!(retrieved_utxo.is_some());
assert_eq!(retrieved_utxo.unwrap().hash, utxo.hash);
}
#[test]
fn test_get_item_not_exists() {
let mut smt = UTXOSparseMerkleTree::new();
let utxo = sample_utxo().unwrap();
// Insert one UTXO and try to fetch a different hash
smt.insert_item(utxo).unwrap();
let non_existent_hash = TreeHashType::default();
let result = smt.get_item(non_existent_hash).unwrap();
// Test that retrieval for a non-existent UTXO returns None
assert!(result.is_none());
}
#[test]
fn test_get_membership_proof() {
let mut smt = UTXOSparseMerkleTree::new();
let utxo = sample_utxo().unwrap();
smt.insert_item(utxo.clone()).unwrap();
// Fetch membership proof for the inserted UTXO
let proof = smt.get_membership_proof(utxo.hash).unwrap();
// Test proof is generated successfully
assert!(proof.is_some());
}
#[test]
fn test_get_membership_proof_not_exists() {
let mut smt = UTXOSparseMerkleTree::new();
// Try fetching proof for a non-existent UTXO hash
let non_existent_hash = TreeHashType::default();
let proof = smt.get_membership_proof(non_existent_hash).unwrap();
// Test no proof is generated for a non-existent UTXO
assert!(proof.is_none());
}
}