sdk/crypto/smt: Implement leaf removal support

This commit is contained in:
parazyd
2024-04-03 11:15:42 +02:00
parent dcf419b0ca
commit f425397115
8 changed files with 125 additions and 37 deletions

View File

@@ -19,9 +19,11 @@
use darkfi_sdk::{
crypto::ContractId,
dark_tree::DarkLeaf,
db::{db_init, db_lookup, db_set},
error::ContractResult,
util::set_return_data,
wasm::{
db::{db_init, db_lookup, db_set},
util::set_return_data,
},
ContractCall,
};
use darkfi_serial::{deserialize, serialize};

View File

@@ -19,11 +19,11 @@
use darkfi_sdk::{
crypto::{ContractId, PublicKey},
dark_tree::DarkLeaf,
db::{db_get, db_lookup, db_set},
deploy::DeployParamsV1,
error::{ContractError, ContractResult},
msg,
pasta::pallas,
wasm::db::{db_get, db_lookup, db_set},
ContractCall,
};
use darkfi_serial::{deserialize, serialize, Encodable, WriteExt};

View File

@@ -19,10 +19,10 @@
use darkfi_sdk::{
crypto::{ContractId, PublicKey},
dark_tree::DarkLeaf,
db::{db_contains_key, db_get, db_lookup, db_set},
error::{ContractError, ContractResult},
msg,
pasta::pallas,
wasm::db::{db_contains_key, db_get, db_lookup, db_set},
ContractCall,
};
use darkfi_serial::{deserialize, serialize, Encodable, WriteExt};

View File

@@ -23,6 +23,7 @@ use darkfi_sdk::{
pasta_prelude::*,
smt::{PoseidonFp, SparseMerkleTree, StorageAdapter, EMPTY_NODES_FP, SMT_FP_DEPTH},
},
error::{ContractError, ContractResult},
wasm,
};
use darkfi_serial::{serialize, Decodable, Encodable};
@@ -42,35 +43,50 @@ pub struct SledStorage<'a> {
impl<'a> StorageAdapter for SledStorage<'a> {
type Value = pallas::Base;
fn put(&mut self, key: BigUint, value: pallas::Base) -> bool {
if self.overlay.insert(self.tree_key, &key.to_bytes_le(), &value.to_repr()).is_err() {
fn put(&mut self, key: BigUint, value: pallas::Base) -> ContractResult {
if let Err(e) = self.overlay.insert(self.tree_key, &key.to_bytes_le(), &value.to_repr()) {
error!(
target: "runtime::smt::SledStorage::put",
"[WASM] sparse_merkle_insert_batch(): inserting key {:?}, value {:?} into DB tree: {:?}",
key, value, self.tree_key
"[WASM] sparse_merkle_insert_batch(): inserting key {:?}, value {:?} into DB tree: {:?}: {}",
key, value, self.tree_key, e,
);
return false
return Err(ContractError::SmtPutFailed)
}
true
Ok(())
}
fn get(&self, key: &BigUint) -> Option<pallas::Base> {
let Ok(value) = self.overlay.get(self.tree_key, &key.to_bytes_le()) else {
error!(
target: "runtime::smt::SledStorage::get",
"[WASM] sparse_merkle_insert_batch(): fetching key {:?} from DB tree: {:?}",
key, self.tree_key
);
return None
let value = match self.overlay.get(self.tree_key, &key.to_bytes_le()) {
Ok(v) => v,
Err(e) => {
error!(
target: "runtime::smt::SledStorage::get",
"[WASM] SledStorage::get(): Fetching key {:?} from DB tree: {:?}: {}",
key, self.tree_key, e,
);
return None
}
};
let value = value?;
let mut repr = [0; 32];
repr.copy_from_slice(&value);
let value = pallas::Base::from_repr(repr);
if value.is_none().into() {
None
} else {
Some(value.unwrap())
pallas::Base::from_repr(repr).into()
}
fn del(&mut self, key: &BigUint) -> ContractResult {
if let Err(e) = self.overlay.remove(self.tree_key, &key.to_bytes_le()) {
error!(
target: "runtime::smt::SledStorage::del",
"[WASM] SledStorage::del(): Removing key {:?} from DB tree: {:?}: {}",
key, self.tree_key, e,
);
return Err(ContractError::SmtDelFailed)
}
Ok(())
}
}

View File

@@ -63,7 +63,7 @@ use std::collections::HashMap;
// Only used for the type aliases below
use pasta_curves::pallas;
use crate::error::{ContractError, ContractResult};
use crate::error::ContractResult;
use util::{FieldElement, FieldHasher};
mod empty;
@@ -92,12 +92,13 @@ pub type SmtMemoryFp = SparseMerkleTree<
pub type PathFp = Path<SMT_FP_DEPTH, pallas::Base, PoseidonFp>;
/// Pluggable storage backend for the SMT.
/// Has a minimal interface to simply put and get objects from the store.
/// Has a minimal interface to put, get, and delete objects from the store.
pub trait StorageAdapter {
type Value;
fn put(&mut self, key: BigUint, value: Self::Value) -> bool;
fn put(&mut self, key: BigUint, value: Self::Value) -> ContractResult;
fn get(&self, key: &BigUint) -> Option<Self::Value>;
fn del(&mut self, key: &BigUint) -> ContractResult;
}
/// An in-memory storage, useful for unit tests and smaller trees.
@@ -115,14 +116,19 @@ impl<F: FieldElement> MemoryStorage<F> {
impl<F: FieldElement> StorageAdapter for MemoryStorage<F> {
type Value = F;
fn put(&mut self, key: BigUint, value: F) -> bool {
fn put(&mut self, key: BigUint, value: F) -> ContractResult {
self.tree.insert(key, value);
true
Ok(())
}
fn get(&self, key: &BigUint) -> Option<F> {
self.tree.get(key).copied()
}
fn del(&mut self, key: &BigUint) -> ContractResult {
self.tree.remove(key);
Ok(())
}
}
/// The Sparse Merkle Tree struct.
@@ -212,6 +218,50 @@ impl<
Ok(())
}
pub fn remove_leaves(&mut self, leaves: Vec<(F, F)>) -> ContractResult {
if leaves.is_empty() {
return Ok(())
}
let mut dirty_idxs = Vec::new();
for (pos, _leaf) in leaves {
let idx = util::leaf_pos_to_index::<N, _>(&pos);
self.remove_node(&idx)?;
// Mark node parent as dirty
let parent_idx = util::parent(&idx).unwrap();
dirty_idxs.push(parent_idx);
}
// Depth first from the bottom of the tree
for _ in 0..N + 1 {
let mut new_dirty_idxs = Vec::new();
for idx in dirty_idxs {
let left_idx = util::left_child(&idx);
let right_idx = util::right_child(&idx);
let left = self.get_node(&left_idx);
let right = self.get_node(&right_idx);
// Recalculate the node
let node = self.hasher.hash([left, right]);
self.put_node(idx.clone(), node)?;
// Add this node's parent to the update list
let parent_idx = match util::parent(&idx) {
Some(idx) => idx,
// We are at the root node so no parents exist
None => break,
};
new_dirty_idxs.push(parent_idx);
}
dirty_idxs = new_dirty_idxs;
}
Ok(())
}
/// Returns the Merkle tree root.
pub fn root(&self) -> F {
self.get_node(&BigUint::from(0u32))
@@ -251,10 +301,11 @@ impl<
}
fn put_node(&mut self, key: BigUint, value: F) -> ContractResult {
if !self.store.put(key, value) {
return Err(ContractError::SmtPutFailed)
}
Ok(())
self.store.put(key, value)
}
fn remove_node(&mut self, key: &BigUint) -> ContractResult {
self.store.del(key)
}
}

View File

@@ -141,6 +141,11 @@ fn poseidon_smt_incl_proof() {
(Fp::from(1), Fp::random(&mut OsRng)),
(Fp::from(2), Fp::random(&mut OsRng)),
(Fp::from(3), Fp::random(&mut OsRng)),
/*
(Fp::from(1), Fp::from(111)),
(Fp::from(2), Fp::from(222)),
(Fp::from(3), Fp::from(333)),
*/
];
smt.insert_batch(leaves.clone()).unwrap();
@@ -149,4 +154,7 @@ fn poseidon_smt_incl_proof() {
let path = smt.prove_membership(&pos);
assert!(path.verify(&smt.root(), &leaf, &pos));
smt.remove_leaves(vec![(pos, leaf)]).unwrap();
assert!(!path.verify(&smt.root(), &leaf, &pos));
}

View File

@@ -21,9 +21,10 @@ use num_bigint::BigUint;
use super::{PoseidonFp, SparseMerkleTree, StorageAdapter, SMT_FP_DEPTH};
use crate::{
crypto::pasta_prelude::*,
error::ContractResult,
msg,
pasta::pallas,
wasm::db::{db_get, db_set, DbHandle},
wasm::db::{db_del, db_get, db_set, DbHandle},
};
pub type SmtWasmFp = SparseMerkleTree<
@@ -48,8 +49,8 @@ impl SmtWasmDbStorage {
impl StorageAdapter for SmtWasmDbStorage {
type Value = pallas::Base;
fn put(&mut self, key: BigUint, value: pallas::Base) -> bool {
db_set(self.db, &key.to_bytes_le(), &value.to_repr()).is_ok()
fn put(&mut self, key: BigUint, value: pallas::Base) -> ContractResult {
db_set(self.db, &key.to_bytes_le(), &value.to_repr())
}
fn get(&self, key: &BigUint) -> Option<pallas::Base> {
@@ -65,4 +66,8 @@ impl StorageAdapter for SmtWasmDbStorage {
pallas::Base::from_repr(repr).into()
}
fn del(&mut self, key: &BigUint) -> ContractResult {
db_del(self.db, &key.to_bytes_le())
}
}

View File

@@ -81,6 +81,9 @@ pub enum ContractError {
#[error("SMT: put failed with storage backend")]
SmtPutFailed,
#[error("SMT: rm failed with storage backend")]
SmtDelFailed,
#[error("Error retrieving system time")]
GetSystemTimeFailed,
@@ -119,9 +122,10 @@ pub const DB_CONTAINS_KEY_FAILED: i64 = to_builtin!(15);
pub const INVALID_FUNCTION: i64 = to_builtin!(16);
pub const DB_DEL_FAILED: i64 = to_builtin!(17);
pub const SMT_PUT_FAILED: i64 = to_builtin!(18);
pub const GET_SYSTEM_TIME_FAILED: i64 = to_builtin!(19);
pub const DATA_TOO_LARGE: i64 = to_builtin!(20);
pub const HEX_FMT_ERR: i64 = to_builtin!(21);
pub const SMT_DEL_FAILED: i64 = to_builtin!(19);
pub const GET_SYSTEM_TIME_FAILED: i64 = to_builtin!(20);
pub const DATA_TOO_LARGE: i64 = to_builtin!(21);
pub const HEX_FMT_ERR: i64 = to_builtin!(22);
impl From<ContractError> for i64 {
fn from(err: ContractError) -> Self {
@@ -143,6 +147,7 @@ impl From<ContractError> for i64 {
ContractError::InvalidFunction => INVALID_FUNCTION,
ContractError::DbDelFailed => DB_DEL_FAILED,
ContractError::SmtPutFailed => SMT_PUT_FAILED,
ContractError::SmtDelFailed => SMT_DEL_FAILED,
ContractError::GetSystemTimeFailed => GET_SYSTEM_TIME_FAILED,
ContractError::DataTooLarge => DATA_TOO_LARGE,
ContractError::HexFmtErr => HEX_FMT_ERR,
@@ -178,6 +183,7 @@ impl From<i64> for ContractError {
INVALID_FUNCTION => Self::InvalidFunction,
DB_DEL_FAILED => Self::DbDelFailed,
SMT_PUT_FAILED => Self::SmtPutFailed,
SMT_DEL_FAILED => Self::SmtDelFailed,
GET_SYSTEM_TIME_FAILED => Self::GetSystemTimeFailed,
DATA_TOO_LARGE => Self::DataTooLarge,
HEX_FMT_ERR => Self::HexFmtErr,