Skip to content

Commit

Permalink
refactor: tree traversal implemented using a new iterator abstraction
Browse files Browse the repository at this point in the history
  • Loading branch information
@cchudant
cchudant committed Sep 25, 2024
1 parent 755034f commit d698670
Show file tree
Hide file tree
Showing 19 changed files with 2,042 additions and 1,708 deletions.
21 changes: 17 additions & 4 deletions benches/storage.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
use std::hint::black_box;

use bitvec::vec::BitVec;
use bonsai_trie::{
databases::HashMapDb,
id::{BasicId, BasicIdBuilder},
Expand All @@ -10,7 +9,7 @@ use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
use rand::{prelude::*, thread_rng};
use starknet_types_core::{
felt::Felt,
hash::{Pedersen, StarkHash},
hash::{Pedersen, Poseidon, StarkHash},
};

mod flamegraph;
Expand Down Expand Up @@ -244,7 +243,7 @@ fn multiple_contracts(c: &mut Criterion) {
});
}

fn hash(c: &mut Criterion) {
fn pedersen_hash(c: &mut Criterion) {
c.bench_function("pedersen hash", move |b| {
let felt0 =
Felt::from_hex("0x100bd6fbfced88ded1b34bd1a55b747ce3a9fde9a914bca75571e4496b56443")
Expand All @@ -258,9 +257,23 @@ fn hash(c: &mut Criterion) {
});
}

fn poseidon_hash(c: &mut Criterion) {
c.bench_function("poseidon hash", move |b| {
let felt0 =
Felt::from_hex("0x100bd6fbfced88ded1b34bd1a55b747ce3a9fde9a914bca75571e4496b56443")
.unwrap();
let felt1 =
Felt::from_hex("0x00a038cda302fedbc4f6117648c6d3faca3cda924cb9c517b46232c6316b152f")
.unwrap();
b.iter(|| {
black_box(Poseidon::hash(&felt0, &felt1));
})
});
}

criterion_group! {
name = benches;
config = Criterion::default(); // .with_profiler(flamegraph::FlamegraphProfiler::new(100));
targets = storage, one_update, five_updates, hash, drop_storage, storage_with_insert, multiple_contracts
targets = storage, one_update, five_updates, pedersen_hash, poseidon_hash, drop_storage, storage_with_insert, multiple_contracts
}
criterion_main!(benches);
6 changes: 2 additions & 4 deletions src/key_value_db.rs
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
use crate::{
changes::key_new_value, format, trie::merkle_tree::bytes_to_bitvec, BTreeSet, ByteVec,
changes::key_new_value, format, trie::merkle_tree::bytes_to_bitvec, BTreeSet, BitVec, ByteVec,
Change as ExternChange, ToString,
};
use bitvec::{order::Msb0, vec::BitVec};
use hashbrown::HashMap;
use log::trace;
use parity_scale_codec::Decode;
Expand Down Expand Up @@ -92,8 +91,7 @@ where
pub(crate) fn get_changes(
&self,
id: ID,
) -> Result<HashMap<BitVec<u8, Msb0>, ExternChange>, BonsaiStorageError<DB::DatabaseError>>
{
) -> Result<HashMap<BitVec, ExternChange>, BonsaiStorageError<DB::DatabaseError>> {
if self.changes_store.id_queue.contains(&id) {
let mut leaf_changes = HashMap::new();
let changes = ChangeBatch::deserialize(
Expand Down
122 changes: 62 additions & 60 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,12 +111,30 @@ pub(crate) use std::{
vec::Vec,
};

use crate::trie::merkle_tree::MerkleTree;

pub type ByteVec = smallvec::SmallVec<[u8; 32]>;
pub type BitVec = bitvec::vec::BitVec<u8, bitvec::order::Msb0>;
pub type BitSlice = bitvec::slice::BitSlice<u8, bitvec::order::Msb0>;

mod changes;
mod key_value_db;
mod trie;

mod bonsai_database;
/// All databases already implemented in this crate.
pub mod databases;
mod error;
/// Definition and basic implementation of an CommitID
pub mod id;

pub use bonsai_database::{BonsaiDatabase, BonsaiPersistentDatabase, DBError, DatabaseKey};
pub use error::BonsaiStorageError;
pub use trie::merkle_tree::Membership;

#[cfg(test)]
mod tests;

pub(crate) trait EncodeExt: parity_scale_codec::Encode {
fn encode_sbytevec(&self) -> ByteVec {
fn encode_bytevec(&self) -> ByteVec {
struct Out(ByteVec);
impl parity_scale_codec::Output for Out {
#[inline]
Expand All @@ -132,32 +150,10 @@ pub(crate) trait EncodeExt: parity_scale_codec::Encode {
}
impl<T: parity_scale_codec::Encode> EncodeExt for T {}

use bitvec::{order::Msb0, slice::BitSlice, vec::BitVec};
use changes::ChangeBatch;
use key_value_db::KeyValueDB;
use starknet_types_core::{
felt::Felt,
hash::{Pedersen, StarkHash},
};

mod changes;
mod key_value_db;
mod trie;

mod bonsai_database;
/// All databases already implemented in this crate.
pub mod databases;
mod error;
/// Definition and basic implementation of an CommitID
pub mod id;

pub use bonsai_database::{BonsaiDatabase, BonsaiPersistentDatabase, DBError, DatabaseKey};
pub use error::BonsaiStorageError;
use trie::merkle_tree::{bytes_to_bitvec, MerkleTrees};
pub use trie::merkle_tree::{Membership, ProofNode};

#[cfg(test)]
mod tests;
use starknet_types_core::{felt::Felt, hash::StarkHash};
use trie::{merkle_tree::bytes_to_bitvec, trees::MerkleTrees};

/// Structure that contains the configuration for the BonsaiStorage.
/// A default implementation is provided with coherent values.
Expand Down Expand Up @@ -264,7 +260,7 @@ where
pub fn insert(
&mut self,
identifier: &[u8],
key: &BitSlice<u8, Msb0>,
key: &BitSlice,
value: &Felt,
) -> Result<(), BonsaiStorageError<DB::DatabaseError>> {
self.tries.set(identifier, key, *value)?;
Expand All @@ -276,7 +272,7 @@ where
pub fn remove(
&mut self,
identifier: &[u8],
key: &BitSlice<u8, Msb0>,
key: &BitSlice,
) -> Result<(), BonsaiStorageError<DB::DatabaseError>> {
self.tries.set(identifier, key, Felt::ZERO)?;
Ok(())
Expand All @@ -286,7 +282,7 @@ where
pub fn get(
&self,
identifier: &[u8],
key: &BitSlice<u8, Msb0>,
key: &BitSlice,
) -> Result<Option<Felt>, BonsaiStorageError<DB::DatabaseError>> {
self.tries.get(identifier, key)
}
Expand All @@ -298,7 +294,7 @@ where
pub fn get_at(
&self,
identifier: &[u8],
key: &BitSlice<u8, Msb0>,
key: &BitSlice,
id: ChangeID,
) -> Result<Option<Felt>, BonsaiStorageError<DB::DatabaseError>> {
self.tries.get_at(identifier, key, id)
Expand All @@ -308,7 +304,7 @@ where
pub fn contains(
&self,
identifier: &[u8],
key: &BitSlice<u8, Msb0>,
key: &BitSlice,
) -> Result<bool, BonsaiStorageError<DB::DatabaseError>> {
self.tries.contains(identifier, key)
}
Expand Down Expand Up @@ -402,7 +398,7 @@ where
pub fn get_changes(
&self,
id: ChangeID,
) -> Result<HashMap<BitVec<u8, Msb0>, Change>, BonsaiStorageError<DB::DatabaseError>> {
) -> Result<HashMap<BitVec, Change>, BonsaiStorageError<DB::DatabaseError>> {
self.tries.db_ref().get_changes(id)
}

Expand All @@ -411,6 +407,11 @@ where
self.tries.db_ref().db.dump_database();
}

#[cfg(test)]
pub fn dump(&self) {
self.tries.dump();
}

/// Get trie root hash at the latest commit
pub fn root_hash(
&self,
Expand All @@ -430,24 +431,25 @@ where
Ok(())
}

/// Generates a merkle-proof for a given `key`.
///
/// Returns vector of [`TrieNode`] which form a chain from the root to the key,
/// if it exists, or down to the node which proves that the key does not exist.
///
/// The nodes are returned in order, root first.
///
/// Verification is performed by confirming that:
/// 1. the chain follows the path of `key`, and
/// 2. the hashes are correct, and
/// 3. the root hash matches the known root
pub fn get_proof(
&self,
identifier: &[u8],
key: &BitSlice<u8, Msb0>,
) -> Result<Vec<ProofNode>, BonsaiStorageError<DB::DatabaseError>> {
self.tries.get_proof(identifier, key)
}
// /// Generates a merkle-proof for a given `key`.
// ///
// /// Returns vector of [`TrieNode`] which form a chain from the root to the key,
// /// if it exists, or down to the node which proves that the key does not exist.
// ///
// /// The nodes are returned in order, root first.
// ///
// /// Verification is performed by confirming that:
// /// 1. the chain follows the path of `key`, and
// /// 2. the hashes are correct, and
// /// 3. the root hash matches the known root
// pub fn get_proof(
// &self,
// identifier: &[u8],
// key: &BitSlice,
// ) -> Result<Vec<ProofNode>, BonsaiStorageError<DB::DatabaseError>> {
// todo!()
// // self.tries.get_proof(identifier, key)
// }

/// Get all the keys in a specific trie.
pub fn get_keys(
Expand All @@ -471,15 +473,15 @@ where
self.tries.db_ref().get_latest_id()
}

/// Verifies a merkle-proof for a given `key` and `value`.
pub fn verify_proof(
root: Felt,
key: &BitSlice<u8, Msb0>,
value: Felt,
proofs: &[ProofNode],
) -> Option<Membership> {
MerkleTree::<Pedersen>::verify_proof(root, key, value, proofs)
}
// /// Verifies a merkle-proof for a given `key` and `value`.
// pub fn verify_proof(
// root: Felt,
// key: &BitSlice,
// value: Felt,
// proofs: &[ProofNode],
// ) -> Option<Membership> {
// MerkleTree::<Pedersen>::verify_proof(root, key, value, proofs)
// }
}

impl<ChangeID, DB, H> BonsaiStorage<ChangeID, DB, H>
Expand Down
10 changes: 4 additions & 6 deletions src/tests/madara_comparison.rs
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,9 @@
#![cfg(all(feature = "std", feature = "rocksdb"))]
use bitvec::{bits, order::Msb0, vec::BitVec};
use bitvec::{bits, order::Msb0};
use starknet_types_core::{felt::Felt, hash::Pedersen};

use crate::{
databases::{create_rocks_db, RocksDB, RocksDBConfig},
id::BasicIdBuilder,
BonsaiStorage, BonsaiStorageConfig,
databases::{create_rocks_db, RocksDB, RocksDBConfig}, id::BasicIdBuilder, BitVec, BonsaiStorage, BonsaiStorageConfig
};

#[test]
Expand All @@ -17,7 +15,7 @@ fn trie_height_251() {
let mut bonsai_storage: BonsaiStorage<_, _, Pedersen> =
BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap();
for i in 0..251 {
let mut key: BitVec<u8, Msb0> = bits![u8, Msb0; 0; 251].to_bitvec();
let mut key: BitVec = bits![u8, Msb0; 0; 251].to_bitvec();
key.set(i, true);
let value = Felt::from_hex("0x01").unwrap();
bonsai_storage
Expand All @@ -33,7 +31,7 @@ fn trie_height_251() {
// #[test]// fn test_height_251() {
// let mut tree = super::merkle_patricia_tree::merkle_tree::MerkleTree::<PedersenHasher>::empty();
// for i in 0..251 {
// let mut key: BitVec<u8, Msb0> = bits![u8, Msb0; 0; 251].to_bitvec();
// let mut key: BitVec = bits![u8, Msb0; 0; 251].to_bitvec();
// key.set(i, true);
// let value = Felt::from_hex_be("0x01").unwrap();
// tree.set(key.as_bitslice(), value);
Expand Down
11 changes: 4 additions & 7 deletions src/tests/merge.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,30 +2,27 @@
#![cfg(all(feature = "std", feature = "rocksdb"))]

use crate::{
databases::{create_rocks_db, RocksDB, RocksDBConfig, RocksDBTransaction},
id::{BasicId, BasicIdBuilder},
BonsaiStorage, BonsaiStorageConfig,
databases::{create_rocks_db, RocksDB, RocksDBConfig, RocksDBTransaction}, id::{BasicId, BasicIdBuilder}, BitVec, BonsaiStorage, BonsaiStorageConfig
};
use bitvec::vec::BitVec;
use once_cell::sync::Lazy;
use rocksdb::OptimisticTransactionDB;
use starknet_types_core::{felt::Felt, hash::Pedersen};

static PAIR1: Lazy<(BitVec<u8, bitvec::prelude::Msb0>, Felt)> = Lazy::new(|| {
static PAIR1: Lazy<(BitVec, Felt)> = Lazy::new(|| {
(
BitVec::from_vec(vec![1, 2, 2]),
Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(),
)
});

static PAIR2: Lazy<(BitVec<u8, bitvec::prelude::Msb0>, Felt)> = Lazy::new(|| {
static PAIR2: Lazy<(BitVec, Felt)> = Lazy::new(|| {
(
BitVec::from_vec(vec![1, 2, 3]),
Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(),
)
});

static PAIR3: Lazy<(BitVec<u8, bitvec::prelude::Msb0>, Felt)> = Lazy::new(|| {
static PAIR3: Lazy<(BitVec, Felt)> = Lazy::new(|| {
(
BitVec::from_vec(vec![1, 2, 4]),
Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(),
Expand Down
Loading

0 comments on commit d698670

Please sign in to comment.