From 9568867a728c54b7cdd80b87bbd7e7167fdc9743 Mon Sep 17 00:00:00 2001 From: Trantorian Date: Thu, 4 Apr 2024 02:03:11 +0000 Subject: [PATCH 1/3] test(insert): added tests for insertions into `BonsaiStorage` without skiping first 5 bites --- Cargo.toml | 1 + src/trie/merkle_tree.rs | 520 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 521 insertions(+) diff --git a/Cargo.toml b/Cargo.toml index 1b23a4e..a5dd114 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,3 +43,4 @@ pathfinder-storage = { git = "https://github.com/massalabs/pathfinder.git", pack rand = "0.8.5" tempfile = "3.8.0" rstest = "0.18.2" +test-log = "0.2.15" diff --git a/src/trie/merkle_tree.rs b/src/trie/merkle_tree.rs index 2abae46..b03c1f0 100644 --- a/src/trie/merkle_tree.rs +++ b/src/trie/merkle_tree.rs @@ -1312,6 +1312,526 @@ pub(crate) fn bytes_to_bitvec(bytes: &[u8]) -> BitVec { #[cfg(test)] #[cfg(all(test, feature = "std"))] mod tests { + use std::collections::HashMap; + + use bitvec::{order::Msb0, slice::BitSlice}; + use starknet_types_core::{felt::Felt, hash::Pedersen}; + + use crate::{databases::HashMapDb, id::BasicId, BonsaiStorage, BonsaiStorageConfig}; + + #[test_log::test] + /// The whole point of this test is to experiment with inserting keys into the database without + /// truncating the first 5 bits to which which parts of the lib are causing errors with this. + /// Testing is done over the contract storage of the first few blocks, with a single commit + /// between each block. + fn test_node_decode_error() { + let db = HashMapDb::::default(); + let mut bonsai = + BonsaiStorage::::new(db, BonsaiStorageConfig::default()).unwrap(); + + let block_0 = vec![ + ( + str_to_felt_bytes( + "0x031c887d82502ceb218c06ebb46198da3f7b92864a8223746bc836dda3e34b52", + ), + vec![ + ( + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000005", + ), + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000065", + ), + ), + ( + str_to_felt_bytes( + "0x00cfc2e2866fd08bfb4ac73b70e0c136e326ae18fc797a2c090c8811c695577e", + ), + str_to_felt_bytes( + "0x05f1dd5a5aef88e0498eeca4e7b2ea0fa7110608c11531278742f0b5499af4b3", + ), + ), + ( + str_to_felt_bytes( + "0x05aee31408163292105d875070f98cb48275b8c87e80380b78d30647e05854d5", + ), + str_to_felt_bytes( + "0x00000000000000000000000000000000000000000000000000000000000007c7", + ), + ), + ( + str_to_felt_bytes( + "0x05fac6815fddf6af1ca5e592359862ede14f171e1544fd9e792288164097c35d", + ), + str_to_felt_bytes( + "0x00299e2f4b5a873e95e65eb03d31e532ea2cde43b498b50cd3161145db5542a5", + ), + ), + ( + str_to_felt_bytes( + "0x05fac6815fddf6af1ca5e592359862ede14f171e1544fd9e792288164097c35e", + ), + str_to_felt_bytes( + "0x03d6897cf23da3bf4fd35cc7a43ccaf7c5eaf8f7c5b9031ac9b09a929204175f", + ), + ), + ], + ), + ( + str_to_felt_bytes( + "0x06ee3440b08a9c805305449ec7f7003f27e9f7e287b83610952ec36bdc5a6bae", + ), + vec![ + ( + str_to_felt_bytes( + "0x01e2cd4b3588e8f6f9c4e89fb0e293bf92018c96d7a93ee367d29a284223b6ff", + ), + str_to_felt_bytes( + "0x071d1e9d188c784a0bde95c1d508877a0d93e9102b37213d1e13f3ebc54a7751", + ), + ), + ( + str_to_felt_bytes( + "0x0449908c349e90f81ab13042b1e49dc251eb6e3e51092d9a40f86859f7f415b0", + ), + str_to_felt_bytes( + "0x06cb6104279e754967a721b52bcf5be525fdc11fa6db6ef5c3a4db832acf7804", + ), + ), + ( + str_to_felt_bytes( + "0x048cba68d4e86764105adcdcf641ab67b581a55a4f367203647549c8bf1feea2", + ), + str_to_felt_bytes( + "0x0362d24a3b030998ac75e838955dfee19ec5b6eceb235b9bfbeccf51b6304d0b", + ), + ), + ( + str_to_felt_bytes( + "0x05bdaf1d47b176bfcd1114809af85a46b9c4376e87e361d86536f0288a284b65", + ), + str_to_felt_bytes( + "0x028dff6722aa73281b2cf84cac09950b71fa90512db294d2042119abdd9f4b87", + ), + ), + ( + str_to_felt_bytes( + "0x05bdaf1d47b176bfcd1114809af85a46b9c4376e87e361d86536f0288a284b66", + ), + str_to_felt_bytes( + "0x057a8f8a019ccab5bfc6ff86c96b1392257abb8d5d110c01d326b94247af161c", + ), + ), + ( + str_to_felt_bytes( + "0x05f750dc13ed239fa6fc43ff6e10ae9125a33bd05ec034fc3bb4dd168df3505f", + ), + str_to_felt_bytes( + "0x00000000000000000000000000000000000000000000000000000000000007e5", + ), + ), + ], + ), + ( + str_to_felt_bytes( + "0x0735596016a37ee972c42adef6a3cf628c19bb3794369c65d2c82ba034aecf2c", + ), + vec![ + ( + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000005", + ), + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000064", + ), + ), + ( + str_to_felt_bytes( + "0x002f50710449a06a9fa789b3c029a63bd0b1f722f46505828a9f815cf91b31d8", + ), + str_to_felt_bytes( + "0x02a222e62eabe91abdb6838fa8b267ffe81a6eb575f61e96ec9aa4460c0925a2", + ), + ), + ], + ), + ( + str_to_felt_bytes( + "0x020cfa74ee3564b4cd5435cdace0f9c4d43b939620e4a0bb5076105df0a626c6", + ), + vec![ + ( + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000005", + ), + str_to_felt_bytes( + "0x000000000000000000000000000000000000000000000000000000000000022b", + ), + ), + ( + str_to_felt_bytes( + "0x0313ad57fdf765addc71329abf8d74ac2bce6d46da8c2b9b82255a5076620300", + ), + str_to_felt_bytes( + "0x04e7e989d58a17cd279eca440c5eaa829efb6f9967aaad89022acbe644c39b36", + ), + ), + ( + str_to_felt_bytes( + "0x0313ad57fdf765addc71329abf8d74ac2bce6d46da8c2b9b82255a5076620301", + ), + str_to_felt_bytes( + "0x0453ae0c9610197b18b13645c44d3d0a407083d96562e8752aab3fab616cecb0", + ), + ), + ( + str_to_felt_bytes( + "0x05aee31408163292105d875070f98cb48275b8c87e80380b78d30647e05854d5", + ), + str_to_felt_bytes( + "0x00000000000000000000000000000000000000000000000000000000000007e5", + ), + ), + ( + str_to_felt_bytes( + "0x06cf6c2f36d36b08e591e4489e92ca882bb67b9c39a3afccf011972a8de467f0", + ), + str_to_felt_bytes( + "0x07ab344d88124307c07b56f6c59c12f4543e9c96398727854a322dea82c73240", + ), + ), + ], + ), + ( + str_to_felt_bytes( + "0x031c887d82502ceb218c06ebb46198da3f7b92864a8223746bc836dda3e34b52", + ), + vec![ + ( + str_to_felt_bytes( + "0x00df28e613c065616a2e79ca72f9c1908e17b8c913972a9993da77588dc9cae9", + ), + str_to_felt_bytes( + "0x01432126ac23c7028200e443169c2286f99cdb5a7bf22e607bcd724efa059040", + ), + ), + ( + str_to_felt_bytes( + "0x05f750dc13ed239fa6fc43ff6e10ae9125a33bd05ec034fc3bb4dd168df3505f", + ), + str_to_felt_bytes( + "0x00000000000000000000000000000000000000000000000000000000000007c7", + ), + ), + ], + ), + ]; + + for (contract_address, storage) in block_0.iter() { + bonsai.init_tree(contract_address).unwrap(); + + for (key, value) in storage { + let key = BitSlice::::from_slice(key); + let value = &Felt::from_bytes_be_slice(value); + + bonsai.insert(contract_address, key, value).unwrap(); + } + } + bonsai.commit(BasicId::new(0)).unwrap(); + assert!(bonsai + .get_transactional_state(BasicId::new(0), BonsaiStorageConfig::default()) + .is_ok()); + + // aggregates all changes made so far, keeping only the latest storage + // in case of storage updates + let mut storage_map = HashMap::, HashMap, Vec>>::new(); + for (contract_address, storage) in block_0.iter() { + let map = storage_map + .entry(contract_address.to_vec()) + .or_insert(HashMap::new()); + + for (key, value) in storage { + map.insert(key.to_vec(), value.to_vec()); + } + } + + log::info!("checking after block 0"); + for (contract_address, storage) in storage_map { + log::info!( + "contract address: {:#064x}", + Felt::from_bytes_be_slice(&contract_address) + ); + assert!(bonsai.init_tree(&contract_address).is_ok()); + + for (key, value) in storage { + let key1 = BitSlice::::from_slice(&key); + let value1 = Felt::from_bytes_be_slice(&value); + + match bonsai.get(&contract_address, &key1) { + Ok(Some(value2)) => assert_eq!(value1, value2), + _ => { + log::info!("Missing key {:#064x}", Felt::from_bytes_be_slice(&key)); + panic!("Values do not match insertions!") + } + } + } + } + + let block_1 = vec![ + ( + str_to_felt_bytes( + "0x06538fdd3aa353af8a87f5fe77d1f533ea82815076e30a86d65b72d3eb4f0b80", + ), + vec![ + ( + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000005", + ), + str_to_felt_bytes( + "0x000000000000000000000000000000000000000000000000000000000000022b", + ), + ), + ( + str_to_felt_bytes( + "0x01aed933fd362faecd8ea54ee749092bd21f89901b7d1872312584ac5b636c6d", + ), + str_to_felt_bytes( + "0x00000000000000000000000000000000000000000000000000000000000007e5", + ), + ), + ( + str_to_felt_bytes( + "0x010212fa2be788e5d943714d6a9eac5e07d8b4b48ead96b8d0a0cbe7a6dc3832", + ), + str_to_felt_bytes( + "0x008a81230a7e3ffa40abe541786a9b69fbb601434cec9536d5d5b2ee4df90383", + ), + ), + ( + str_to_felt_bytes( + "0x00ffda4b5cf0dce9bc9b0d035210590c73375fdbb70cd94ec6949378bffc410c", + ), + str_to_felt_bytes( + "0x02b36318931915f71777f7e59246ecab3189db48408952cefda72f4b7977be51", + ), + ), + ( + str_to_felt_bytes( + "0x00ffda4b5cf0dce9bc9b0d035210590c73375fdbb70cd94ec6949378bffc410d", + ), + str_to_felt_bytes( + "0x07e928dcf189b05e4a3dae0bc2cb98e447f1843f7debbbf574151eb67cda8797", + ), + ), + ], + ), + ( + str_to_felt_bytes( + "0x0327d34747122d7a40f4670265b098757270a449ec80c4871450fffdab7c2fa8", + ), + vec![ + ( + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000005", + ), + str_to_felt_bytes( + "0x0000000000000000000000000000000000000000000000000000000000000065", + ), + ), + ( + str_to_felt_bytes( + "0x01aed933fd362faecd8ea54ee749092bd21f89901b7d1872312584ac5b636c6d", + ), + str_to_felt_bytes( + "0x00000000000000000000000000000000000000000000000000000000000007c7", + ), + ), + ( + str_to_felt_bytes( + "0x04184fa5a6d40f47a127b046ed6facfa3e6bc3437b393da65cc74afe47ca6c6e", + ), + str_to_felt_bytes( + "0x001ef78e458502cd457745885204a4ae89f3880ec24db2d8ca97979dce15fedc", + ), + ), + ( + str_to_felt_bytes( + "0x05591c8c3c8d154a30869b463421cd5933770a0241e1a6e8ebcbd91bdd69bec4", + ), + str_to_felt_bytes( + "0x026b5943d4a0c420607cee8030a8cdd859bf2814a06633d165820960a42c6aed", + ), + ), + ( + str_to_felt_bytes( + "0x05591c8c3c8d154a30869b463421cd5933770a0241e1a6e8ebcbd91bdd69bec5", + ), + str_to_felt_bytes( + "0x01518eec76afd5397cefd14eda48d01ad59981f9ce9e70c233ca67acd8754008", + ), + ), + ], + ), + ]; + + for (contract_address, storage) in block_1.iter() { + bonsai.init_tree(contract_address).unwrap(); + + for (key, value) in storage { + let key = BitSlice::::from_slice(key); + let value = &Felt::from_bytes_be_slice(value); + + bonsai.insert(contract_address, key, value).unwrap(); + } + } + bonsai.commit(BasicId::new(1)).unwrap(); + assert!(bonsai + .get_transactional_state(BasicId::new(1), BonsaiStorageConfig::default()) + .is_ok()); + + let mut storage_map = HashMap::, HashMap, Vec>>::new(); + for (contract_address, storage) in block_0.iter().chain(block_1.iter()) { + let map = storage_map + .entry(contract_address.to_vec()) + .or_insert(HashMap::new()); + + for (key, value) in storage { + map.insert(key.to_vec(), value.to_vec()); + } + } + + log::info!("checking after block 1"); + for (contract_address, storage) in storage_map { + log::info!( + "contract address: {:#064x}", + Felt::from_bytes_be_slice(&contract_address) + ); + assert!(bonsai.init_tree(&contract_address).is_ok()); + + for (key, value) in storage { + let key1 = BitSlice::::from_slice(&key); + let value1 = Felt::from_bytes_be_slice(&value); + + match bonsai.get(&contract_address, &key1) { + Ok(Some(value2)) => assert_eq!(value1, value2), + _ => { + log::info!("Missing key {:#064x}", Felt::from_bytes_be_slice(&key)); + panic!("Values do not match insertions!") + } + } + } + } + + let block_2 = vec![ + ( + str_to_felt_bytes( + "0x001fb4457f3fe8a976bdb9c04dd21549beeeb87d3867b10effe0c4bd4064a8e4", + ), + vec![( + str_to_felt_bytes( + "0x056c060e7902b3d4ec5a327f1c6e083497e586937db00af37fe803025955678f", + ), + str_to_felt_bytes( + "0x075495b43f53bd4b9c9179db113626af7b335be5744d68c6552e3d36a16a747c", + ), + )], + ), + ( + str_to_felt_bytes( + "0x05790719f16afe1450b67a92461db7d0e36298d6a5f8bab4f7fd282050e02f4f", + ), + vec![( + str_to_felt_bytes( + "0x0772c29fae85f8321bb38c9c3f6edb0957379abedc75c17f32bcef4e9657911a", + ), + str_to_felt_bytes( + "0x06d4ca0f72b553f5338a95625782a939a49b98f82f449c20f49b42ec60ed891c", + ), + )], + ), + ( + str_to_felt_bytes( + "0x057b973bf2eb26ebb28af5d6184b4a044b24a8dcbf724feb95782c4d1aef1ca9", + ), + vec![( + str_to_felt_bytes( + "0x04f2c206f3f2f1380beeb9fe4302900701e1cb48b9b33cbe1a84a175d7ce8b50", + ), + str_to_felt_bytes( + "0x02a614ae71faa2bcdacc5fd66965429c57c4520e38ebc6344f7cf2e78b21bd2f", + ), + )], + ), + ( + str_to_felt_bytes( + "0x02d6c9569dea5f18628f1ef7c15978ee3093d2d3eec3b893aac08004e678ead3", + ), + vec![( + str_to_felt_bytes( + "0x07f93985c1baa5bd9b2200dd2151821bd90abb87186d0be295d7d4b9bc8ca41f", + ), + str_to_felt_bytes( + "0x0127cd00a078199381403a33d315061123ce246c8e5f19aa7f66391a9d3bf7c6", + ), + )], + ), + ]; + + for (contract_address, storage) in block_2.iter() { + bonsai.init_tree(contract_address).unwrap(); + + for (key, value) in storage { + let key = BitSlice::::from_slice(key); + let value = &Felt::from_bytes_be_slice(value); + + bonsai.insert(contract_address, key, value).unwrap(); + } + } + bonsai.commit(BasicId::new(2)).unwrap(); + + let mut storage_map = HashMap::, HashMap, Vec>>::new(); + for (contract_address, storage) in + block_0.iter().chain(block_1.iter()).chain(block_2.iter()) + { + let map = storage_map + .entry(contract_address.to_vec()) + .or_insert(HashMap::new()); + + for (key, value) in storage { + map.insert(key.to_vec(), value.to_vec()); + } + } + + log::info!("checking after block 2"); + for (contract_address, storage) in storage_map { + log::info!( + "contract address: {:#064x}", + Felt::from_bytes_be_slice(&contract_address) + ); + assert!(bonsai.init_tree(&contract_address).is_ok()); + + for (key, value) in storage { + let key1 = BitSlice::::from_slice(&key); + let value1 = Felt::from_bytes_be_slice(&value); + + match bonsai.get(&contract_address, &key1) { + Ok(Some(value2)) => assert_eq!(value1, value2), + _ => { + log::info!("Missing key {:#064x}", Felt::from_bytes_be_slice(&key)); + panic!("Values do not match insertions!") + } + } + } + } + + assert!(bonsai + .get_transactional_state(BasicId::new(2), BonsaiStorageConfig::default()) + .is_ok()) + } + + fn str_to_felt_bytes(hex: &str) -> [u8; 32] { + Felt::from_hex(hex).unwrap().to_bytes_be() + } // use crate::{ // databases::{create_rocks_db, RocksDB, RocksDBConfig}, // id::BasicId, From cdc12703ff945ad485ca7d0b6679e2cd5eede516 Mon Sep 17 00:00:00 2001 From: Trantorian Date: Wed, 24 Apr 2024 12:16:38 +0000 Subject: [PATCH 2/3] feat(storage): added and updated functions related to storage access Updated `get_keys` to be easier to use. Added `get_key_value_pairs` to retrieve all key-values for a given identifier. Added `get_latest_id` to retrieve the latest id committed to a `BonsaiStorage`. --- Cargo.toml | 14 +- README.md | 7 + benches/flamegraph.rs | 39 +++ benches/storage.rs | 157 +++++++++ ensure_no_std/Cargo.lock | 32 +- src/changes.rs | 2 + src/databases/mod.rs | 2 +- src/key_value_db.rs | 5 + src/lib.rs | 65 +++- src/tests/merge.rs | 576 +++++++++++++++++++++++++++++++ src/tests/mod.rs | 1 + src/tests/transactional_state.rs | 167 +++++++++ src/trie/merkle_tree.rs | 514 ++++++++++++++++----------- 13 files changed, 1350 insertions(+), 231 deletions(-) create mode 100644 benches/flamegraph.rs create mode 100644 benches/storage.rs create mode 100644 src/tests/merge.rs diff --git a/Cargo.toml b/Cargo.toml index a5dd114..0dfc007 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,8 @@ version = "0.1.0" default = ["std", "rocksdb"] rocksdb = ["dep:rocksdb"] std = ["parity-scale-codec/std", "bitvec/std", "starknet-types-core/std"] +# internal +bench = [] [dependencies] bitvec = { version = "1", default-features = false, features = ["alloc"] } @@ -25,7 +27,7 @@ serde = { version = "1.0.195", default-features = false, features = [ "derive", "alloc", ] } -starknet-types-core = { git = "https://github.com/starknet-io/types-rs", branch = "main", default-features = false, features = [ +starknet-types-core = { version = "0.1", default-features = false, features = [ "hash", "parity-scale-codec", ] } @@ -36,6 +38,9 @@ rocksdb = { optional = true, version = "0.21.0", features = [ ] } [dev-dependencies] +env_logger = "0.11.3" +once_cell = "1.19.0" +pprof = { version = "0.3", features = ["flamegraph"] } pathfinder-common = { git = "https://github.com/massalabs/pathfinder.git", package = "pathfinder-common", rev = "b7b6d76a76ab0e10f92e5f84ce099b5f727cb4db" } pathfinder-crypto = { git = "https://github.com/massalabs/pathfinder.git", package = "pathfinder-crypto", rev = "b7b6d76a76ab0e10f92e5f84ce099b5f727cb4db" } pathfinder-merkle-tree = { git = "https://github.com/massalabs/pathfinder.git", package = "pathfinder-merkle-tree", rev = "b7b6d76a76ab0e10f92e5f84ce099b5f727cb4db" } @@ -44,3 +49,10 @@ rand = "0.8.5" tempfile = "3.8.0" rstest = "0.18.2" test-log = "0.2.15" +indexmap = "2.2.6" +criterion = "0.5.1" + +[[bench]] +name = "storage" +required-features = ["bench"] +harness = false diff --git a/README.md b/README.md index d8b1d20..4e18fa8 100644 --- a/README.md +++ b/README.md @@ -153,6 +153,13 @@ fn main() { } ``` +## Build and run benchmarks + +This crate uses `rayon` to parallelize hash computations. As such, results will vary depending on the number of cores of your cpu. +``` +cargo bench +``` + ## Acknowledgements - Shout out to [Danno Ferrin](https://github.com/shemnon) and [Karim Taam](https://github.com/matkt) for their work on Bonsai. This project is heavily inspired by their work. diff --git a/benches/flamegraph.rs b/benches/flamegraph.rs new file mode 100644 index 0000000..4fc346b --- /dev/null +++ b/benches/flamegraph.rs @@ -0,0 +1,39 @@ +use criterion::profiler::Profiler; +use pprof::ProfilerGuard; +use std::{fs::File, os::raw::c_int, path::Path}; + +pub struct FlamegraphProfiler<'a> { + frequency: c_int, + active_profiler: Option>, +} + +impl<'a> FlamegraphProfiler<'a> { + #[allow(dead_code)] + pub fn new(frequency: c_int) -> Self { + FlamegraphProfiler { + frequency, + active_profiler: None, + } + } +} + +impl<'a> Profiler for FlamegraphProfiler<'a> { + fn start_profiling(&mut self, _benchmark_id: &str, _benchmark_dir: &Path) { + self.active_profiler = Some(ProfilerGuard::new(self.frequency).unwrap()); + } + + fn stop_profiling(&mut self, _benchmark_id: &str, benchmark_dir: &Path) { + std::fs::create_dir_all(benchmark_dir).unwrap(); + let flamegraph_path = benchmark_dir.join("flamegraph.svg"); + let flamegraph_file = File::create(&flamegraph_path) + .expect("File system error while creating flamegraph.svg"); + if let Some(profiler) = self.active_profiler.take() { + profiler + .report() + .build() + .unwrap() + .flamegraph(flamegraph_file) + .expect("Error writing flamegraph"); + } + } +} diff --git a/benches/storage.rs b/benches/storage.rs new file mode 100644 index 0000000..373e92f --- /dev/null +++ b/benches/storage.rs @@ -0,0 +1,157 @@ +use std::hint::black_box; + +use bitvec::vec::BitVec; +use bonsai_trie::{ + databases::HashMapDb, + id::{BasicId, BasicIdBuilder}, + BonsaiStorage, BonsaiStorageConfig, +}; +use criterion::{criterion_group, criterion_main, Criterion}; +use rand::{prelude::*, thread_rng}; +use starknet_types_core::{ + felt::Felt, + hash::{Pedersen, StarkHash}, +}; + +mod flamegraph; + +fn storage(c: &mut Criterion) { + c.bench_function("storage commit", move |b| { + let mut bonsai_storage: BonsaiStorage = BonsaiStorage::new( + HashMapDb::::default(), + BonsaiStorageConfig::default(), + ) + .unwrap(); + let mut rng = thread_rng(); + + let felt = Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(); + for _ in 0..1000 { + let bitvec = BitVec::from_vec(vec![ + rng.gen(), + rng.gen(), + rng.gen(), + rng.gen(), + rng.gen(), + rng.gen(), + ]); + bonsai_storage.insert(&[], &bitvec, &felt).unwrap(); + } + + let mut id_builder = BasicIdBuilder::new(); + b.iter_batched( + || bonsai_storage.clone(), + |mut bonsai_storage| { + bonsai_storage.commit(id_builder.new_id()).unwrap(); + }, + criterion::BatchSize::LargeInput, + ); + }); +} + +fn one_update(c: &mut Criterion) { + c.bench_function("one update", move |b| { + let mut bonsai_storage: BonsaiStorage = BonsaiStorage::new( + HashMapDb::::default(), + BonsaiStorageConfig::default(), + ) + .unwrap(); + let mut rng = thread_rng(); + + let felt = Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(); + for _ in 0..1000 { + let bitvec = BitVec::from_vec(vec![ + rng.gen(), + rng.gen(), + rng.gen(), + rng.gen(), + rng.gen(), + rng.gen(), + ]); + bonsai_storage.insert(&[], &bitvec, &felt).unwrap(); + } + + let mut id_builder = BasicIdBuilder::new(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + + b.iter_batched( + || bonsai_storage.clone(), + |mut bonsai_storage| { + let bitvec = BitVec::from_vec(vec![0, 1, 2, 3, 4, 5]); + bonsai_storage.insert(&[], &bitvec, &felt).unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + }, + criterion::BatchSize::LargeInput, + ); + }); +} + +fn five_updates(c: &mut Criterion) { + c.bench_function("five updates", move |b| { + let mut bonsai_storage: BonsaiStorage = BonsaiStorage::new( + HashMapDb::::default(), + BonsaiStorageConfig::default(), + ) + .unwrap(); + let mut rng = thread_rng(); + + let felt = Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(); + for _ in 0..1000 { + let bitvec = BitVec::from_vec(vec![ + rng.gen(), + rng.gen(), + rng.gen(), + rng.gen(), + rng.gen(), + rng.gen(), + ]); + bonsai_storage.insert(&[], &bitvec, &felt).unwrap(); + } + + let mut id_builder = BasicIdBuilder::new(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + + b.iter_batched( + || bonsai_storage.clone(), + |mut bonsai_storage| { + bonsai_storage + .insert(&[], &BitVec::from_vec(vec![0, 1, 2, 3, 4, 5]), &felt) + .unwrap(); + bonsai_storage + .insert(&[], &BitVec::from_vec(vec![0, 2, 2, 5, 4, 5]), &felt) + .unwrap(); + bonsai_storage + .insert(&[], &BitVec::from_vec(vec![0, 1, 2, 3, 3, 5]), &felt) + .unwrap(); + bonsai_storage + .insert(&[], &BitVec::from_vec(vec![0, 1, 1, 3, 99, 3]), &felt) + .unwrap(); + bonsai_storage + .insert(&[], &BitVec::from_vec(vec![0, 1, 2, 3, 4, 6]), &felt) + .unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + }, + criterion::BatchSize::LargeInput, + ); + }); +} + +fn hash(c: &mut Criterion) { + c.bench_function("pedersen hash", move |b| { + let felt0 = + Felt::from_hex("0x100bd6fbfced88ded1b34bd1a55b747ce3a9fde9a914bca75571e4496b56443") + .unwrap(); + let felt1 = + Felt::from_hex("0x00a038cda302fedbc4f6117648c6d3faca3cda924cb9c517b46232c6316b152f") + .unwrap(); + b.iter(|| { + black_box(Pedersen::hash(&felt0, &felt1)); + }) + }); +} + +criterion_group! { + name = benches; + config = Criterion::default(); // .with_profiler(flamegraph::FlamegraphProfiler::new(100)); + targets = storage, one_update, five_updates, hash +} +criterion_main!(benches); diff --git a/ensure_no_std/Cargo.lock b/ensure_no_std/Cargo.lock index 30a9d11..689c82e 100644 --- a/ensure_no_std/Cargo.lock +++ b/ensure_no_std/Cargo.lock @@ -229,9 +229,9 @@ dependencies = [ [[package]] name = "lambdaworks-crypto" -version = "0.5.0" +version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d4c222d5b2fdc0faf702d3ab361d14589b097f40eac9dc550e27083483edc65" +checksum = "458fee521f12d0aa97a2e06eaf134398a5d2ae7b2074af77eb402b0d93138c47" dependencies = [ "lambdaworks-math", "sha2", @@ -240,18 +240,9 @@ dependencies = [ [[package]] name = "lambdaworks-math" -version = "0.5.0" +version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ee7dcab3968c71896b8ee4dc829147acc918cffe897af6265b1894527fe3add" - -[[package]] -name = "lazy_static" -version = "1.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" -dependencies = [ - "spin", -] +checksum = "6c74ce6f0d9cb672330b6ca59e85a6c3607a3329e0372ab0d3fe38c2d38e50f9" [[package]] name = "libc" @@ -300,9 +291,9 @@ dependencies = [ [[package]] name = "num-traits" -version = "0.2.17" +version = "0.2.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c" +checksum = "da0df0e5185db44f69b44f26786fe401b6c293d1907744beaa7fa62b2e5a517a" dependencies = [ "autocfg", ] @@ -438,22 +429,15 @@ version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6ecd384b10a64542d77071bd64bd7b231f4ed5940fba55e98c3de13824cf3d7" -[[package]] -name = "spin" -version = "0.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d" - [[package]] name = "starknet-types-core" -version = "0.0.7" +version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6d791c671fecde494f435170a01c6fcb2949d0dd61be0b31b7c410b041609f96" +checksum = "1051b4f4af0bb9b546388a404873ee1e6b9787b9d5b0b3319ecbfadf315ef276" dependencies = [ "bitvec", "lambdaworks-crypto", "lambdaworks-math", - "lazy_static", "num-bigint", "num-integer", "num-traits", diff --git a/src/changes.rs b/src/changes.rs index 3d9a19c..3018736 100644 --- a/src/changes.rs +++ b/src/changes.rs @@ -15,6 +15,7 @@ pub struct Change { } #[derive(Debug, Default)] +#[cfg_attr(feature = "bench", derive(Clone))] pub struct ChangeBatch(pub(crate) HashMap); const KEY_SEPARATOR: u8 = 0x00; @@ -115,6 +116,7 @@ impl ChangeBatch { } } +#[cfg_attr(feature = "bench", derive(Clone))] pub struct ChangeStore where ID: Id, diff --git a/src/databases/mod.rs b/src/databases/mod.rs index e73625d..6b54e2f 100644 --- a/src/databases/mod.rs +++ b/src/databases/mod.rs @@ -6,4 +6,4 @@ pub use hashmap_db::HashMapDb; mod rocks_db; #[cfg(feature = "rocksdb")] -pub use rocks_db::{create_rocks_db, RocksDB, RocksDBBatch, RocksDBConfig}; +pub use rocks_db::{create_rocks_db, RocksDB, RocksDBBatch, RocksDBConfig, RocksDBTransaction}; diff --git a/src/key_value_db.rs b/src/key_value_db.rs index 84ce41f..2b4d099 100644 --- a/src/key_value_db.rs +++ b/src/key_value_db.rs @@ -18,6 +18,7 @@ use crate::{ }; /// Crate Trie <= KeyValueDB => BonsaiDatabase +#[cfg_attr(feature = "bench", derive(Clone))] pub struct KeyValueDB where DB: BonsaiDatabase, @@ -170,6 +171,10 @@ where Ok(self.db.get(&key.into())?) } + pub(crate) fn get_latest_id(&self) -> Option { + self.changes_store.id_queue.back().cloned() + } + pub(crate) fn contains( &self, key: &TrieKey, diff --git a/src/lib.rs b/src/lib.rs index 24d4f31..0272a72 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -87,7 +87,7 @@ #[cfg(not(feature = "std"))] extern crate alloc; -use crate::trie::merkle_tree::MerkleTree; +use crate::trie::merkle_tree::{bytes_to_bitvec, MerkleTree}; #[cfg(not(feature = "std"))] use alloc::{format, vec::Vec}; use bitvec::{order::Msb0, slice::BitSlice, vec::BitVec}; @@ -168,6 +168,20 @@ where tries: MerkleTrees, } +#[cfg(feature = "bench")] +impl Clone for BonsaiStorage +where + DB: BonsaiDatabase + Clone, + ChangeID: id::Id, + H: StarkHash + Send + Sync, +{ + fn clone(&self) -> Self { + Self { + tries: self.tries.clone(), + } + } +} + /// Trie root hash type. pub type BonsaiTrieHash = Felt; @@ -397,6 +411,20 @@ where self.tries.get_keys(identifier) } + /// Get all the key-value pairs in a specific trie. + #[allow(clippy::type_complexity)] + pub fn get_key_value_pairs( + &self, + identifier: &[u8], + ) -> Result, Vec)>, BonsaiStorageError> { + self.tries.get_key_value_pairs(identifier) + } + + /// Get the id from the latest commit, or `None` if no commit has taken place yet. + pub fn get_latest_id(&self) -> Option { + self.tries.db_ref().get_latest_id() + } + /// Verifies a merkle-proof for a given `key` and `value`. pub fn verify_proof( root: Felt, @@ -460,9 +488,38 @@ where &mut self, transactional_bonsai_storage: BonsaiStorage, ) -> Result<(), BonsaiStorageError<>::DatabaseError>> + where + ::DatabaseError: core::fmt::Debug, { - self.tries - .db_mut() - .merge(transactional_bonsai_storage.tries.db()) + // memorize changes + let MerkleTrees { db, trees, .. } = transactional_bonsai_storage.tries; + + self.tries.db_mut().merge(db)?; + + // apply changes + for (identifier, tree) in trees { + for (k, op) in tree.cache_leaf_modified() { + match op { + crate::trie::merkle_tree::InsertOrRemove::Insert(v) => { + self.insert(&identifier, &bytes_to_bitvec(k), v) + .map_err(|e| { + BonsaiStorageError::Merge(format!( + "While merging insert({:?} {}) faced error: {:?}", + k, v, e + )) + })?; + } + crate::trie::merkle_tree::InsertOrRemove::Remove => { + self.remove(&identifier, &bytes_to_bitvec(k)).map_err(|e| { + BonsaiStorageError::Merge(format!( + "While merging remove({:?}) faced error: {:?}", + k, e + )) + })?; + } + } + } + } + Ok(()) } } diff --git a/src/tests/merge.rs b/src/tests/merge.rs new file mode 100644 index 0000000..83181a2 --- /dev/null +++ b/src/tests/merge.rs @@ -0,0 +1,576 @@ +#![cfg(feature = "std")] +use crate::{ + databases::{create_rocks_db, RocksDB, RocksDBConfig, RocksDBTransaction}, + id::{BasicId, BasicIdBuilder}, + BonsaiStorage, BonsaiStorageConfig, +}; +use bitvec::vec::BitVec; +use rocksdb::OptimisticTransactionDB; +use starknet_types_core::{felt::Felt, hash::Pedersen}; + +use once_cell::sync::Lazy; + +static PAIR1: Lazy<(BitVec, Felt)> = Lazy::new(|| { + ( + BitVec::from_vec(vec![1, 2, 2]), + Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), + ) +}); + +static PAIR2: Lazy<(BitVec, Felt)> = Lazy::new(|| { + ( + BitVec::from_vec(vec![1, 2, 3]), + Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(), + ) +}); + +static PAIR3: Lazy<(BitVec, Felt)> = Lazy::new(|| { + ( + BitVec::from_vec(vec![1, 2, 4]), + Felt::from_hex("0x66342762FD54D033c195fec3ce2568b62052e").unwrap(), + ) +}); + +/// Initializes a test environment for the BonsaiStorage data structure. +/// +/// # Arguments +/// +/// * `db` - An instance of the `OptimisticTransactionDB` struct. +/// +/// # Returns +/// +/// A tuple containing the following elements: +/// * `identifier` - A vector of bytes. +/// * `bonsai_storage` - An instance of `BonsaiStorage` with a `RocksDB` +/// backend. +/// * `bonsai_at_txn` - An instance of `BonsaiStorage` representing the +/// transactional state of `bonsai_storage` at `start_id`. +/// * `id_builder` - An instance of `BasicIdBuilder`. +/// * `start_id` - A `BasicId` representing the commit ID of the changes made in +/// `bonsai_storage`. +fn init_test<'db>( + db: &'db OptimisticTransactionDB, +) -> ( + Vec, + BonsaiStorage, Pedersen>, + BonsaiStorage, Pedersen>, + BasicIdBuilder, + BasicId, +) { + let identifier = vec![]; + + let config = BonsaiStorageConfig::default(); + let mut bonsai_storage = + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config) + .expect("Failed to create BonsaiStorage"); + + let mut id_builder = BasicIdBuilder::new(); + + bonsai_storage + .insert(&identifier, &PAIR1.0, &PAIR1.1) + .expect("Failed to insert key-value pair"); + + let start_id = id_builder.new_id(); + bonsai_storage + .commit(start_id) + .expect("Failed to commit changes"); + + let bonsai_at_txn = bonsai_storage + .get_transactional_state(start_id, BonsaiStorageConfig::default()) + .expect("Failed to get transactional state") + .expect("Transactional state not found"); + + ( + identifier, + bonsai_storage, + bonsai_at_txn, + id_builder, + start_id, + ) +} + +#[test] +fn merge_before_simple() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, _) = init_test(&db); + + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + bonsai_storage.merge(bonsai_at_txn).unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + + assert_eq!( + bonsai_storage.get(&identifier, &PAIR2.0).unwrap(), + Some(PAIR2.1) + ); +} + +#[test] +fn merge_before_simple_remove() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, start_id) = + init_test(&db); + + bonsai_at_txn.remove(&identifier, &PAIR1.0).unwrap(); + bonsai_storage.merge(bonsai_at_txn).unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + + assert_eq!( + bonsai_storage.contains(&identifier, &PAIR1.0).unwrap(), + false + ); + + bonsai_storage.revert_to(start_id).unwrap(); + + assert_eq!( + bonsai_storage.get(&identifier, &PAIR1.0).unwrap(), + Some(PAIR1.1) + ); +} + +#[test] +fn merge_tx_commit_simple_remove() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, start_id) = + init_test(&db); + + bonsai_at_txn.remove(&identifier, &PAIR1.0).unwrap(); + bonsai_at_txn + .transactional_commit(id_builder.new_id()) + .unwrap(); + + bonsai_storage.merge(bonsai_at_txn).unwrap(); + + assert_eq!( + bonsai_storage.contains(&identifier, &PAIR1.0).unwrap(), + false + ); + + bonsai_storage.revert_to(start_id).unwrap(); + + assert_eq!( + bonsai_storage.get(&identifier, &PAIR1.0).unwrap(), + Some(PAIR1.1) + ); +} + +#[test] +fn merge_before_simple_revert_to() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, start_id) = + init_test(&db); + + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + bonsai_storage.merge(bonsai_at_txn).unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + bonsai_storage.revert_to(start_id).unwrap(); + + assert_eq!( + bonsai_storage.get(&identifier, &PAIR2.0).unwrap().is_none(), + true + ); +} + +#[test] +fn merge_transactional_commit_in_txn_before() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, start_id) = + init_test(&db); + + let id2 = id_builder.new_id(); + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + bonsai_at_txn.transactional_commit(id2).unwrap(); + + let id3 = id_builder.new_id(); + bonsai_at_txn + .insert(&identifier, &PAIR3.0, &PAIR3.1) + .unwrap(); + bonsai_storage.merge(bonsai_at_txn).unwrap(); + bonsai_storage.commit(id3).unwrap(); + + assert_eq!( + bonsai_storage.get(&identifier, &PAIR2.0).unwrap(), + Some(PAIR2.1) + ); + assert_eq!( + bonsai_storage.get(&identifier, &PAIR3.0).unwrap(), + Some(PAIR3.1) + ); + bonsai_storage.revert_to(id2).unwrap(); + assert_eq!( + bonsai_storage.get(&identifier, &PAIR2.0).unwrap(), + Some(PAIR2.1) + ); + assert_eq!( + bonsai_storage.get(&identifier, &PAIR3.0).unwrap().is_none(), + true + ); + bonsai_storage.revert_to(start_id).unwrap(); + assert_eq!( + bonsai_storage.get(&identifier, &PAIR2.0).unwrap().is_none(), + true + ); + assert_eq!( + bonsai_storage.get(&identifier, &PAIR3.0).unwrap().is_none(), + true + ); +} + +#[test] +fn merge_transactional_commit_in_txn_before_existing_key() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, start_id) = + init_test(&db); + + bonsai_at_txn.remove(&identifier, &PAIR1.0).unwrap(); + + let id2 = id_builder.new_id(); + bonsai_at_txn.transactional_commit(id2).unwrap(); + bonsai_storage.merge(bonsai_at_txn).unwrap(); + bonsai_storage.revert_to(id2).unwrap(); + + assert_eq!( + bonsai_storage.get(&identifier, &PAIR1.0).unwrap().is_none(), + true + ); + + bonsai_storage.revert_to(start_id).unwrap(); + + assert_eq!( + bonsai_storage.get(&identifier, &PAIR1.0.clone()).unwrap(), + Some(PAIR1.1) + ); +} + +#[test] +fn merge_get_uncommitted() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, _, _) = init_test(&db); + + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + bonsai_storage.merge(bonsai_at_txn).unwrap(); + + assert_eq!( + bonsai_storage.get(&identifier, &PAIR2.0).unwrap(), + Some(PAIR2.1) + ); +} + +#[test] +fn merge_conflict_commited_vs_commited() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, _) = init_test(&db); + + // insert a key in the transactional state + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + let id = id_builder.new_id(); + bonsai_at_txn.transactional_commit(id).unwrap(); + + // insert same key with a different value in the bonsai_storage + bonsai_storage + .insert(&identifier, &PAIR2.0, &PAIR3.1) + .unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + + match bonsai_storage.merge(bonsai_at_txn) { + Ok(_) => panic!("Expected merge conflict error"), + Err(err) => assert_eq!( + err.to_string(), + "Merge error: Transaction created_at BasicId(0) is lower than the last recorded id" + ), + } +} + +#[test] +fn merge_conflict_commited_vs_commited_change_order() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, _) = init_test(&db); + + // insert same key with a different value in the bonsai_storage + bonsai_storage + .insert(&identifier, &PAIR2.0, &PAIR3.1) + .unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + + // insert a key in the transactional state + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + bonsai_at_txn + .transactional_commit(id_builder.new_id()) + .unwrap(); + + match bonsai_storage.merge(bonsai_at_txn) { + Ok(_) => panic!("Expected merge conflict error"), + Err(err) => assert_eq!( + err.to_string(), + "Merge error: Transaction created_at BasicId(0) is lower than the last recorded id" + ), + } +} + +#[test] +fn merge_conflict_commited_vs_commited_and_noncommited() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, _) = init_test(&db); + + // insert a key in the transactional state + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + let id = id_builder.new_id(); + bonsai_at_txn.transactional_commit(id).unwrap(); + bonsai_at_txn + .insert(&identifier, &PAIR3.0, &PAIR3.1) + .unwrap(); + + // insert same key with a different value in the bonsai_storage + bonsai_storage + .insert(&identifier, &PAIR2.0, &PAIR3.1) + .unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + + match bonsai_storage.merge(bonsai_at_txn) { + Ok(_) => panic!("Expected merge conflict error"), + Err(err) => assert_eq!( + err.to_string(), + "Merge error: Transaction created_at BasicId(0) is lower than the last recorded id" + ), + } +} + +#[test] +fn merge_conflict_commited_and_noncommited_vs_commited() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, _) = init_test(&db); + + // insert a key in the transactional state + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + let id = id_builder.new_id(); + bonsai_at_txn.transactional_commit(id).unwrap(); + + // insert same key with a different value in the bonsai_storage + bonsai_storage + .insert(&identifier, &PAIR2.0, &PAIR3.1) + .unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + bonsai_storage.remove(&identifier, &PAIR2.0).unwrap(); + // .insert(&identifier, &PAIR3.0, &PAIR3.1) + // .unwrap(); + + match bonsai_storage.merge(bonsai_at_txn) { + Ok(_) => panic!("Expected merge conflict error"), + Err(err) => assert_eq!( + err.to_string(), + "Merge error: Transaction created_at BasicId(0) is lower than the last recorded id" + ), + } +} + +#[test] +fn merge_conflict_commited_and_noncommited_vs_noncommited() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, _) = init_test(&db); + + // insert a key in the transactional state + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + + // insert same key with a different value in the bonsai_storage + bonsai_storage + .insert(&identifier, &PAIR2.0, &PAIR3.1) + .unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + bonsai_storage.remove(&identifier, &PAIR2.0).unwrap(); + // .insert(&identifier, &PAIR3.0, &PAIR3.1) + // .unwrap(); + + match bonsai_storage.merge(bonsai_at_txn) { + Ok(_) => panic!("Expected merge conflict error"), + Err(err) => assert_eq!( + err.to_string(), + "Merge error: Transaction created_at BasicId(0) is lower than the last recorded id" + ), + } +} + +#[test] +fn merge_conflict_commited_vs_noncommited() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, _) = init_test(&db); + + // insert a key in the transactional state + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + + // insert same key with a different value in the bonsai_storage + bonsai_storage + .insert(&identifier, &PAIR2.0, &PAIR3.1) + .unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + + match bonsai_storage.merge(bonsai_at_txn) { + Ok(_) => panic!("Expected merge conflict error"), + Err(err) => assert_eq!( + err.to_string(), + "Merge error: Transaction created_at BasicId(0) is lower than the last recorded id" + ), + } +} + +#[test] +fn merge_conflict_noncommited_vs_commited() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, _) = init_test(&db); + + // insert same key with a different value in the bonsai_storage + bonsai_storage + .insert(&identifier, &PAIR2.0, &PAIR3.1) + .unwrap(); + + // insert a key in the transactional state + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + bonsai_at_txn + .transactional_commit(id_builder.new_id()) + .unwrap(); + + bonsai_storage.merge(bonsai_at_txn).unwrap(); + + // check that changes in the transactional state overwrite the ones in the + // storage + let get = bonsai_storage.get(&identifier, &PAIR2.0).unwrap(); + assert_eq!(get, Some(PAIR2.1)); +} + +#[test] +fn merge_conflict_noncommited_vs_noncommited() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, _, _) = init_test(&db); + + // insert same key with a different value in the bonsai_storage + bonsai_storage + .insert(&identifier, &PAIR2.0, &PAIR3.1) + .unwrap(); + + // insert a key in the transactional state + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + + bonsai_storage.merge(bonsai_at_txn).unwrap(); + + // check that changes in the transactional state overwrite the ones in the + // storage + let get = bonsai_storage.get(&identifier, &PAIR2.0).unwrap(); + assert_eq!(get, Some(PAIR2.1)); +} + +#[test] +fn merge_conflict_noncommited_vs_commited_noncommited() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, _) = init_test(&db); + + // insert same key with a different value in the bonsai_storage + bonsai_storage + .insert(&identifier, &PAIR2.0, &PAIR3.1) + .unwrap(); + + // insert a key in the transactional state + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + bonsai_at_txn + .transactional_commit(id_builder.new_id()) + .unwrap(); + bonsai_at_txn + .insert(&identifier, &PAIR3.0, &PAIR3.1) + .unwrap(); + + bonsai_storage.merge(bonsai_at_txn).unwrap(); + + // change in the transactional state overwrites any noncommited changes in + // the storage + let get = bonsai_storage.get(&identifier, &PAIR2.0).unwrap(); + assert_eq!(get, Some(PAIR2.1)); + let get = bonsai_storage.get(&identifier, &PAIR3.0).unwrap(); + assert_eq!(get, Some(PAIR2.1)); +} + +#[test] +fn merge_conflict_commited_noncommited_vs_commited_noncommited() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, _) = init_test(&db); + + // insert same key with a different value in the bonsai_storage + bonsai_storage + .insert(&identifier, &PAIR2.0, &PAIR3.1) + .unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + bonsai_storage + .insert(&identifier, &PAIR3.0, &PAIR2.1) + .unwrap(); + + // insert a key in the transactional state + bonsai_at_txn + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + bonsai_at_txn + .transactional_commit(id_builder.new_id()) + .unwrap(); + bonsai_at_txn + .insert(&identifier, &PAIR3.0, &PAIR3.1) + .unwrap(); + + match bonsai_storage.merge(bonsai_at_txn) { + Ok(_) => { + panic!("Expected merge conflict error") + } + Err(err) => assert_eq!( + err.to_string(), + "Merge error: Transaction created_at BasicId(0) is lower than the last recorded id" + ), + } +} + +#[test] +fn merge_nonconflict_commited_vs_commited() { + let db = create_rocks_db(tempfile::tempdir().unwrap().path()).unwrap(); + let (identifier, mut bonsai_storage, mut bonsai_at_txn, mut id_builder, _) = init_test(&db); + + bonsai_storage + .insert(&identifier, &PAIR2.0, &PAIR2.1) + .unwrap(); + bonsai_storage.commit(id_builder.new_id()).unwrap(); + + bonsai_at_txn + .insert(&identifier, &PAIR3.0, &PAIR3.1) + .unwrap(); + bonsai_at_txn + .transactional_commit(id_builder.new_id()) + .unwrap(); + + match bonsai_storage.merge(bonsai_at_txn) { + Ok(_) => { + panic!("Expected merge conflict error") + } + Err(err) => assert_eq!( + err.to_string(), + "Merge error: Transaction created_at BasicId(0) is lower than the last recorded id" + ), + } +} diff --git a/src/tests/mod.rs b/src/tests/mod.rs index 0292186..81223fc 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -1,4 +1,5 @@ mod madara_comparison; +mod merge; mod proof; mod simple; mod transactional_state; diff --git a/src/tests/transactional_state.rs b/src/tests/transactional_state.rs index 7271e8e..1a4d38a 100644 --- a/src/tests/transactional_state.rs +++ b/src/tests/transactional_state.rs @@ -5,6 +5,7 @@ use crate::{ BonsaiStorage, BonsaiStorageConfig, }; use bitvec::vec::BitVec; +use log::LevelFilter; use starknet_types_core::{felt::Felt, hash::Pedersen}; #[test] @@ -200,6 +201,172 @@ fn merge() { ); } +#[test] +fn merge_with_uncommitted_insert() { + let identifier = vec![]; + let tempdir = tempfile::tempdir().unwrap(); + let db = create_rocks_db(tempdir.path()).unwrap(); + let config = BonsaiStorageConfig::default(); + let mut bonsai_storage = + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + let mut id_builder = BasicIdBuilder::new(); + + let pair1 = ( + BitVec::from_vec(vec![1, 2, 2]), + Felt::from_hex("0x66342762FDD5D033c195fec3ce2568b62052e").unwrap(), + ); + let pair2 = ( + BitVec::from_vec(vec![1, 2, 3]), + Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), + ); + + let id1 = id_builder.new_id(); + bonsai_storage + .insert(&identifier, &pair1.0, &pair1.1) + .unwrap(); + bonsai_storage.commit(id1).unwrap(); + + let mut bonsai_at_txn: BonsaiStorage<_, _, Pedersen> = bonsai_storage + .get_transactional_state(id1, BonsaiStorageConfig::default()) + .unwrap() + .unwrap(); + bonsai_at_txn + .insert(&identifier, &pair2.0, &pair2.1) + .unwrap(); + + bonsai_storage.merge(bonsai_at_txn).unwrap(); + + // commit after merge + let revert_id = id_builder.new_id(); + bonsai_storage.commit(revert_id).unwrap(); + + // overwrite pair2 + bonsai_storage + .insert( + &identifier, + &pair2.0, + &Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052F").unwrap(), + ) + .unwrap(); + + // revert to commit + bonsai_storage.revert_to(revert_id).unwrap(); + + assert_eq!( + bonsai_storage.get(&identifier, &pair2.0).unwrap(), + Some(pair2.1) + ); +} + +#[test] +fn merge_with_uncommitted_remove() { + let _ = env_logger::builder().is_test(true).try_init(); + log::set_max_level(LevelFilter::Trace); + + let identifier = vec![]; + let tempdir = tempfile::tempdir().unwrap(); + let db = create_rocks_db(tempdir.path()).unwrap(); + let config = BonsaiStorageConfig::default(); + let mut bonsai_storage = + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + let mut id_builder = BasicIdBuilder::new(); + + let pair1 = ( + BitVec::from_vec(vec![1, 2, 2]), + Felt::from_hex("0x66342762FDD5D033c195fec3ce2568b62052e").unwrap(), + ); + let pair2 = ( + BitVec::from_vec(vec![1, 2, 3]), + Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), + ); + + let id1 = id_builder.new_id(); + bonsai_storage + .insert(&identifier, &pair1.0, &pair1.1) + .unwrap(); + bonsai_storage.commit(id1).unwrap(); + + let mut bonsai_at_txn: BonsaiStorage<_, _, Pedersen> = bonsai_storage + .get_transactional_state(id1, BonsaiStorageConfig::default()) + .unwrap() + .unwrap(); + bonsai_at_txn + .insert(&identifier, &pair2.0, &pair2.1) + .unwrap(); + bonsai_at_txn + .transactional_commit(id_builder.new_id()) + .unwrap(); + + // remove pair2 but don't commit in transational state + bonsai_at_txn.remove(&identifier, &pair2.0).unwrap(); + assert_eq!( + bonsai_at_txn.contains(&identifier, &pair2.0).unwrap(), + false + ); + + let merge = bonsai_storage.merge(bonsai_at_txn); + match merge { + Ok(_) => println!("merge succeeded"), + Err(e) => { + println!("merge failed"); + panic!("{}", e); + } + }; + + // commit after merge + bonsai_storage.commit(id_builder.new_id()).unwrap(); + + assert_eq!( + bonsai_storage.contains(&identifier, &pair2.0).unwrap(), + false + ); +} + +#[test] +fn transactional_state_after_uncommitted() { + let _ = env_logger::builder().is_test(true).try_init(); + log::set_max_level(LevelFilter::Trace); + + let identifier = vec![]; + let tempdir = tempfile::tempdir().unwrap(); + let db = create_rocks_db(tempdir.path()).unwrap(); + let config = BonsaiStorageConfig::default(); + let mut bonsai_storage = + BonsaiStorage::new(RocksDB::new(&db, RocksDBConfig::default()), config).unwrap(); + let mut id_builder = BasicIdBuilder::new(); + + let pair1 = ( + BitVec::from_vec(vec![1, 2, 2]), + Felt::from_hex("0x66342762FDD5D033c195fec3ce2568b62052e").unwrap(), + ); + let pair2 = ( + BitVec::from_vec(vec![1, 2, 3]), + Felt::from_hex("0x66342762FDD54D033c195fec3ce2568b62052e").unwrap(), + ); + + let id1 = id_builder.new_id(); + bonsai_storage + .insert(&identifier, &pair1.0, &pair1.1) + .unwrap(); + bonsai_storage.commit(id1).unwrap(); + + // make a change to original tree but don't commit it + bonsai_storage + .insert(&identifier, &pair2.0, &pair2.1) + .unwrap(); + + // create a transactional state after the uncommitted change + let bonsai_at_txn: BonsaiStorage<_, _, Pedersen> = bonsai_storage + .get_transactional_state(id1, BonsaiStorageConfig::default()) + .unwrap() + .unwrap(); + + // uncommitted changes, done after the transactional state was created, + // are not included in the transactional state + let contains = bonsai_at_txn.contains(&identifier, &pair2.0).unwrap(); + assert_eq!(contains, false); +} + #[test] fn merge_override() { let identifier = vec![]; diff --git a/src/trie/merkle_tree.rs b/src/trie/merkle_tree.rs index b03c1f0..122d506 100644 --- a/src/trie/merkle_tree.rs +++ b/src/trie/merkle_tree.rs @@ -1,5 +1,5 @@ #[cfg(not(feature = "std"))] -use alloc::{format, string::ToString, vec::Vec}; +use alloc::{format, string::ToString, vec, vec::Vec}; use bitvec::{ prelude::{BitSlice, BitVec, Msb0}, view::BitView, @@ -68,15 +68,25 @@ impl ProofNode { pub(crate) struct MerkleTrees { pub db: KeyValueDB, - _hasher: PhantomData, pub trees: HashMap, MerkleTree>, } +#[cfg(feature = "bench")] +impl Clone + for MerkleTrees +{ + fn clone(&self) -> Self { + Self { + db: self.db.clone(), + trees: self.trees.clone(), + } + } +} + impl MerkleTrees { pub(crate) fn new(db: KeyValueDB) -> Self { Self { db, - _hasher: PhantomData, trees: HashMap::new(), } } @@ -150,10 +160,6 @@ impl MerkleTrees KeyValueDB { - self.db - } - pub(crate) fn root_hash( &self, identifier: &[u8], @@ -177,40 +183,61 @@ impl MerkleTrees identifier.len() { + Some(key[identifier.len() + 1..].to_vec()) + } else { + None + } + }) + .collect() + }) + .map_err(|e| e.into()) + } + + #[allow(clippy::type_complexity)] + pub(crate) fn get_key_value_pairs( + &self, + identifier: &[u8], + ) -> Result, Vec)>, BonsaiStorageError> { + self.db + .db + .get_by_prefix(&crate::DatabaseKey::Flat(identifier)) + .map(|key_value_pairs| { + key_value_pairs + .into_iter() + // FIXME: this does not filter out keys values correctly for `HashMapDb` due + // to branches and leafs not being differenciated + .filter_map(|(key, value)| { + if key.len() > identifier.len() { + Some((key[identifier.len() + 1..].to_vec(), value)) + } else { + None + } + }) .collect() }) .map_err(|e| e.into()) } pub(crate) fn commit(&mut self) -> Result<(), BonsaiStorageError> { - #[allow(clippy::type_complexity)] #[cfg(not(feature = "std"))] - let db_changes: Vec< - Result< - HashMap>>, - BonsaiStorageError, - >, - > = self + let db_changes = self .trees .iter_mut() .map(|(_, tree)| tree.get_updates::()) - .collect(); - #[allow(clippy::type_complexity)] + .collect::, BonsaiStorageError>>()?; #[cfg(feature = "std")] - let db_changes: Vec< - Result< - HashMap>>, - BonsaiStorageError, - >, - > = self + let db_changes = self .trees .par_iter_mut() .map(|(_, tree)| tree.get_updates::()) - .collect(); + .collect::, BonsaiStorageError>>()?; + let mut batch = self.db.create_batch(); for changes in db_changes { - let changes = changes?; for (key, value) in changes { match value { InsertOrRemove::Insert(value) => { @@ -270,11 +297,32 @@ pub struct MerkleTree { _hasher: PhantomData, } -#[derive(Debug, PartialEq, Eq)] +// NB: #[derive(Clone)] does not work because it expands to an impl block which forces H: Clone, which Pedersen/Poseidon aren't. +#[cfg(feature = "bench")] +impl Clone for MerkleTree { + fn clone(&self) -> Self { + Self { + root_handle: self.root_handle.clone(), + root_hash: self.root_hash.clone(), + identifier: self.identifier.clone(), + storage_nodes: self.storage_nodes.clone(), + latest_node_id: self.latest_node_id.clone(), + death_row: self.death_row.clone(), + cache_leaf_modified: self.cache_leaf_modified.clone(), + _hasher: PhantomData, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] pub(crate) enum InsertOrRemove { Insert(T), Remove, } +enum NodeOrFelt<'a> { + Node(&'a Node), + Felt(Felt), +} impl MerkleTree { /// Less visible initialization for `MerkleTree` as the main entry points should be @@ -315,6 +363,10 @@ impl MerkleTree { self.root_hash } + pub fn cache_leaf_modified(&self) -> &HashMap, InsertOrRemove> { + &self.cache_leaf_modified + } + /// Remove all the modifications that have been done since the last commit. pub fn reset_to_last_commit( &mut self, @@ -340,22 +392,30 @@ impl MerkleTree { #[allow(clippy::type_complexity)] pub(crate) fn get_updates( &mut self, - ) -> Result>>, BonsaiStorageError> + ) -> Result>)>, BonsaiStorageError> { - let mut updates = HashMap::new(); + let mut updates = vec![]; for node_key in mem::take(&mut self.death_row) { - updates.insert(node_key, InsertOrRemove::Remove); + updates.push((node_key, InsertOrRemove::Remove)); } - let root_hash = - self.commit_subtree::(&mut updates, self.root_handle, Path(BitVec::new()))?; + + let mut hashes = vec![]; + self.compute_root_hash::(&mut hashes)?; + let root_hash = self.commit_subtree::( + &mut updates, + self.root_handle, + Path(BitVec::new()), + &mut hashes.drain(..), + )?; + for (key, value) in mem::take(&mut self.cache_leaf_modified) { - updates.insert( + updates.push(( TrieKey::new(&self.identifier, TrieKeyType::Flat, &key), match value { InsertOrRemove::Insert(value) => InsertOrRemove::Insert(value.encode()), InsertOrRemove::Remove => InsertOrRemove::Remove, }, - ); + )); } self.latest_node_id.reset(); self.root_hash = root_hash; @@ -363,24 +423,147 @@ impl MerkleTree { Ok(updates) } + fn get_node_or_felt( + &self, + node_handle: &NodeHandle, + ) -> Result> { + let node_id = match node_handle { + NodeHandle::Hash(hash) => return Ok(NodeOrFelt::Felt(*hash)), + NodeHandle::InMemory(root_id) => root_id, + }; + let node = self + .storage_nodes + .0 + .get(node_id) + .ok_or(BonsaiStorageError::Trie( + "Couldn't fetch node in the temporary storage".to_string(), + ))?; + Ok(NodeOrFelt::Node(node)) + } + + fn compute_root_hash( + &self, + hashes: &mut Vec, + ) -> Result> { + match self.get_node_or_felt::(&self.root_handle)? { + NodeOrFelt::Felt(felt) => Ok(felt), + NodeOrFelt::Node(node) => self.compute_hashes::(node, Path(BitVec::new()), hashes), + } + } + + /// Compute the hashes of all of the updated nodes in the merkle tree. This step + /// is separate from [`commit_subtree`] as it is done in parallel using rayon. + /// Computed hashes are pushed to the `hashes` vector, depth first. + fn compute_hashes( + &self, + node: &Node, + path: Path, + hashes: &mut Vec, + ) -> Result> { + use Node::*; + + match node { + Unresolved(hash) => Ok(*hash), + Binary(binary) => { + // we check if we have one or two changed children + + let left_path = path.new_with_direction(Direction::Left); + let node_left = self.get_node_or_felt::(&binary.left)?; + let right_path = path.new_with_direction(Direction::Right); + let node_right = self.get_node_or_felt::(&binary.right)?; + + let (left_hash, right_hash) = match (node_left, node_right) { + #[cfg(feature = "std")] + (NodeOrFelt::Node(left), NodeOrFelt::Node(right)) => { + // two children: use rayon + let (left, right) = rayon::join( + || self.compute_hashes::(left, left_path, hashes), + || { + let mut hashes = vec![]; + let felt = + self.compute_hashes::(right, right_path, &mut hashes)?; + Ok::<_, BonsaiStorageError>((felt, hashes)) + }, + ); + let (left_hash, (right_hash, hashes2)) = (left?, right?); + hashes.extend(hashes2); + + (left_hash, right_hash) + } + (left, right) => { + let left_hash = match left { + NodeOrFelt::Felt(felt) => felt, + NodeOrFelt::Node(node) => { + self.compute_hashes::(node, left_path, hashes)? + } + }; + let right_hash = match right { + NodeOrFelt::Felt(felt) => felt, + NodeOrFelt::Node(node) => { + self.compute_hashes::(node, right_path, hashes)? + } + }; + (left_hash, right_hash) + } + }; + + let hash = H::hash(&left_hash, &right_hash); + hashes.push(hash); + Ok(hash) + } + + Edge(edge) => { + let mut child_path = path.clone(); + child_path.0.extend(&edge.path.0); + let child_hash = match self.get_node_or_felt::(&edge.child)? { + NodeOrFelt::Felt(felt) => felt, + NodeOrFelt::Node(node) => { + self.compute_hashes::(node, child_path, hashes)? + } + }; + + let mut bytes = [0u8; 32]; + bytes.view_bits_mut::()[256 - edge.path.0.len()..] + .copy_from_bitslice(&edge.path.0); + + let felt_path = Felt::from_bytes_be(&bytes); + let mut length = [0; 32]; + // Safe as len() is guaranteed to be <= 251 + length[31] = edge.path.0.len() as u8; + + let length = Felt::from_bytes_be(&length); + let hash = H::hash(&child_hash, &felt_path) + length; + hashes.push(hash); + Ok(hash) + } + } + } + /// Persists any changes in this subtree to storage. /// /// This necessitates recursively calculating the hash of, and /// in turn persisting, any changed child nodes. This is necessary /// as the parent node's hash relies on its children hashes. + /// Hash computation is done in parallel with [`compute_hashes`] beforehand. /// /// In effect, the entire tree gets persisted. /// /// # Arguments /// - /// * `node` - The top node from the subtree to commit. + /// * `node_handle` - The top node from the subtree to commit. + /// * `hashes` - The precomputed hashes for the subtree as returned by [`compute_hashes`]. + /// The order is depth first, left to right. + /// + /// # Panics + /// + /// Panics if the precomputed `hashes` do not match the length of the modified subtree. fn commit_subtree( &mut self, - updates: &mut HashMap>>, + updates: &mut Vec<(TrieKey, InsertOrRemove>)>, node_handle: NodeHandle, path: Path, + hashes: &mut impl Iterator, ) -> Result> { - use Node::*; let node_id = match node_handle { NodeHandle::Hash(hash) => return Ok(hash), NodeHandle::InMemory(root_id) => root_id, @@ -393,56 +576,48 @@ impl MerkleTree { .ok_or(BonsaiStorageError::Trie( "Couldn't fetch node in the temporary storage".to_string(), ))? { - Unresolved(hash) => { + Node::Unresolved(hash) => { if path.0.is_empty() { - updates.insert( + updates.push(( TrieKey::new(&self.identifier, TrieKeyType::Trie, &[]), InsertOrRemove::Insert(Node::Unresolved(hash).encode()), - ); + )); Ok(hash) } else { Ok(hash) } } - Binary(mut binary) => { + Node::Binary(mut binary) => { let left_path = path.new_with_direction(Direction::Left); - let left_hash = self.commit_subtree::(updates, binary.left, left_path)?; + let left_hash = + self.commit_subtree::(updates, binary.left, left_path, hashes)?; let right_path = path.new_with_direction(Direction::Right); - let right_hash = self.commit_subtree::(updates, binary.right, right_path)?; - let hash = H::hash(&left_hash, &right_hash); + let right_hash = + self.commit_subtree::(updates, binary.right, right_path, hashes)?; + let hash = hashes.next().expect("mismatched hash state"); binary.hash = Some(hash); binary.left = NodeHandle::Hash(left_hash); binary.right = NodeHandle::Hash(right_hash); let key_bytes: Vec = path.into(); - updates.insert( + updates.push(( TrieKey::new(&self.identifier, TrieKeyType::Trie, &key_bytes), InsertOrRemove::Insert(Node::Binary(binary).encode()), - ); + )); Ok(hash) } - - Edge(mut edge) => { + Node::Edge(mut edge) => { let mut child_path = path.clone(); child_path.0.extend(&edge.path.0); - let child_hash = self.commit_subtree::(updates, edge.child, child_path)?; - let mut bytes = [0u8; 32]; - bytes.view_bits_mut::()[256 - edge.path.0.len()..] - .copy_from_bitslice(&edge.path.0); - - let felt_path = Felt::from_bytes_be(&bytes); - let mut length = [0; 32]; - // Safe as len() is guaranteed to be <= 251 - length[31] = edge.path.0.len() as u8; - - let length = Felt::from_bytes_be(&length); - let hash = H::hash(&child_hash, &felt_path) + length; + let child_hash = + self.commit_subtree::(updates, edge.child, child_path, hashes)?; + let hash = hashes.next().expect("mismatched hash state"); edge.hash = Some(hash); edge.child = NodeHandle::Hash(child_hash); let key_bytes: Vec = path.into(); - updates.insert( + updates.push(( TrieKey::new(&self.identifier, TrieKeyType::Trie, &key_bytes), InsertOrRemove::Insert(Node::Edge(edge).encode()), - ); + )); Ok(hash) } } @@ -1312,20 +1487,23 @@ pub(crate) fn bytes_to_bitvec(bytes: &[u8]) -> BitVec { #[cfg(test)] #[cfg(all(test, feature = "std"))] mod tests { - use std::collections::HashMap; - - use bitvec::{order::Msb0, slice::BitSlice}; + use bitvec::{order::Msb0, vec::BitVec, view::BitView}; + use indexmap::IndexMap; use starknet_types_core::{felt::Felt, hash::Pedersen}; - use crate::{databases::HashMapDb, id::BasicId, BonsaiStorage, BonsaiStorageConfig}; + use crate::{ + databases::{create_rocks_db, RocksDB, RocksDBConfig}, + id::BasicId, + BonsaiStorage, BonsaiStorageConfig, + }; #[test_log::test] - /// The whole point of this test is to experiment with inserting keys into the database without - /// truncating the first 5 bits to which which parts of the lib are causing errors with this. - /// Testing is done over the contract storage of the first few blocks, with a single commit - /// between each block. - fn test_node_decode_error() { - let db = HashMapDb::::default(); + // The whole point of this test is to make sure it is possible to reconstruct the original + // keys from the data present in the db. + fn test_key_retrieval() { + let tempdir = tempfile::tempdir().unwrap(); + let rocksdb = create_rocks_db(tempdir.path()).unwrap(); + let db = RocksDB::new(&rocksdb, RocksDBConfig::default()); let mut bonsai = BonsaiStorage::::new(db, BonsaiStorageConfig::default()).unwrap(); @@ -1527,56 +1705,6 @@ mod tests { ), ]; - for (contract_address, storage) in block_0.iter() { - bonsai.init_tree(contract_address).unwrap(); - - for (key, value) in storage { - let key = BitSlice::::from_slice(key); - let value = &Felt::from_bytes_be_slice(value); - - bonsai.insert(contract_address, key, value).unwrap(); - } - } - bonsai.commit(BasicId::new(0)).unwrap(); - assert!(bonsai - .get_transactional_state(BasicId::new(0), BonsaiStorageConfig::default()) - .is_ok()); - - // aggregates all changes made so far, keeping only the latest storage - // in case of storage updates - let mut storage_map = HashMap::, HashMap, Vec>>::new(); - for (contract_address, storage) in block_0.iter() { - let map = storage_map - .entry(contract_address.to_vec()) - .or_insert(HashMap::new()); - - for (key, value) in storage { - map.insert(key.to_vec(), value.to_vec()); - } - } - - log::info!("checking after block 0"); - for (contract_address, storage) in storage_map { - log::info!( - "contract address: {:#064x}", - Felt::from_bytes_be_slice(&contract_address) - ); - assert!(bonsai.init_tree(&contract_address).is_ok()); - - for (key, value) in storage { - let key1 = BitSlice::::from_slice(&key); - let value1 = Felt::from_bytes_be_slice(&value); - - match bonsai.get(&contract_address, &key1) { - Ok(Some(value2)) => assert_eq!(value1, value2), - _ => { - log::info!("Missing key {:#064x}", Felt::from_bytes_be_slice(&key)); - panic!("Values do not match insertions!") - } - } - } - } - let block_1 = vec![ ( str_to_felt_bytes( @@ -1674,54 +1802,6 @@ mod tests { ), ]; - for (contract_address, storage) in block_1.iter() { - bonsai.init_tree(contract_address).unwrap(); - - for (key, value) in storage { - let key = BitSlice::::from_slice(key); - let value = &Felt::from_bytes_be_slice(value); - - bonsai.insert(contract_address, key, value).unwrap(); - } - } - bonsai.commit(BasicId::new(1)).unwrap(); - assert!(bonsai - .get_transactional_state(BasicId::new(1), BonsaiStorageConfig::default()) - .is_ok()); - - let mut storage_map = HashMap::, HashMap, Vec>>::new(); - for (contract_address, storage) in block_0.iter().chain(block_1.iter()) { - let map = storage_map - .entry(contract_address.to_vec()) - .or_insert(HashMap::new()); - - for (key, value) in storage { - map.insert(key.to_vec(), value.to_vec()); - } - } - - log::info!("checking after block 1"); - for (contract_address, storage) in storage_map { - log::info!( - "contract address: {:#064x}", - Felt::from_bytes_be_slice(&contract_address) - ); - assert!(bonsai.init_tree(&contract_address).is_ok()); - - for (key, value) in storage { - let key1 = BitSlice::::from_slice(&key); - let value1 = Felt::from_bytes_be_slice(&value); - - match bonsai.get(&contract_address, &key1) { - Ok(Some(value2)) => assert_eq!(value1, value2), - _ => { - log::info!("Missing key {:#064x}", Felt::from_bytes_be_slice(&key)); - panic!("Values do not match insertions!") - } - } - } - } - let block_2 = vec![ ( str_to_felt_bytes( @@ -1777,61 +1857,93 @@ mod tests { ), ]; - for (contract_address, storage) in block_2.iter() { - bonsai.init_tree(contract_address).unwrap(); + let blocks = block_0.iter().chain(block_1.iter()).chain(block_2.iter()); - for (key, value) in storage { - let key = BitSlice::::from_slice(key); - let value = &Felt::from_bytes_be_slice(value); + // Inserts all storage updates into the bonsai + for (contract_address, storage) in blocks.clone() { + log::info!( + "contract address (write): {:#064x}", + Felt::from_bytes_be_slice(contract_address) + ); + assert!(bonsai.init_tree(contract_address).is_ok()); + + for (k, v) in storage { + // truncate only keeps the first 251 bits in a key + // so there should be no error during insertion + let ktrunc = &truncate(k); + let kfelt0 = Felt::from_bytes_be_slice(k); + let kfelt1 = Felt::from_bytes_be_slice(ktrunc.as_raw_slice()); - bonsai.insert(contract_address, key, value).unwrap(); + // quick sanity check to make sure truncating a key does not remove any data + assert_eq!(kfelt0, kfelt1); + + let v = &Felt::from_bytes_be_slice(v); + assert!(bonsai.insert(contract_address, ktrunc, v).is_ok()); } } - bonsai.commit(BasicId::new(2)).unwrap(); + assert!(bonsai.commit(BasicId::new(0)).is_ok()); - let mut storage_map = HashMap::, HashMap, Vec>>::new(); - for (contract_address, storage) in - block_0.iter().chain(block_1.iter()).chain(block_2.iter()) - { + // aggreates all storage changes to their latest state + // (replacements are takent into account) + let mut storage_map = IndexMap::, IndexMap>::new(); + for (contract_address, storage) in blocks.clone() { let map = storage_map .entry(contract_address.to_vec()) - .or_insert(HashMap::new()); + .or_insert(IndexMap::new()); - for (key, value) in storage { - map.insert(key.to_vec(), value.to_vec()); + for (k, v) in storage { + let k = Felt::from_bytes_be_slice(k); + let v = Felt::from_bytes_be_slice(v); + map.insert(k, v); } } - log::info!("checking after block 2"); - for (contract_address, storage) in storage_map { + // checks for each contract if the original key can be reconstructed + // from the data stored in the db + for (contract_address, storage) in storage_map.iter() { log::info!( - "contract address: {:#064x}", - Felt::from_bytes_be_slice(&contract_address) + "contract address (read): {:#064x}", + Felt::from_bytes_be_slice(contract_address) ); - assert!(bonsai.init_tree(&contract_address).is_ok()); - for (key, value) in storage { - let key1 = BitSlice::::from_slice(&key); - let value1 = Felt::from_bytes_be_slice(&value); + let keys = bonsai.get_keys(contract_address).unwrap(); + log::debug!("{keys:?}"); + for k in keys { + // if all has gone well, the db should contain the first 251 bits of the key, + // which should represent the entirety of the data + let k = Felt::from_bytes_be_slice(&k); + log::info!("looking for key: {k:#064x}"); - match bonsai.get(&contract_address, &key1) { - Ok(Some(value2)) => assert_eq!(value1, value2), - _ => { - log::info!("Missing key {:#064x}", Felt::from_bytes_be_slice(&key)); - panic!("Values do not match insertions!") - } - } + assert!(storage.contains_key(&k)); } } - assert!(bonsai - .get_transactional_state(BasicId::new(2), BonsaiStorageConfig::default()) - .is_ok()) + // makes sure retrieving key-value pairs works for each contract + for (contract_address, storage) in storage_map.iter() { + log::info!( + "contract address (read): {:#064x}", + Felt::from_bytes_be_slice(contract_address) + ); + + let kv = bonsai.get_key_value_pairs(contract_address).unwrap(); + log::debug!("{kv:?}"); + for (k, v) in kv { + let k = Felt::from_bytes_be_slice(&k); + let v = Felt::from_bytes_be_slice(&v); + log::info!("checking for key-value pair:({k:#064x}, {v:#064x})"); + + assert_eq!(*storage.get(&k).unwrap(), v); + } + } } fn str_to_felt_bytes(hex: &str) -> [u8; 32] { Felt::from_hex(hex).unwrap().to_bytes_be() } + + fn truncate(key: &[u8]) -> BitVec { + key.view_bits()[5..].to_owned() + } // use crate::{ // databases::{create_rocks_db, RocksDB, RocksDBConfig}, // id::BasicId, From 320de334c7dcb0ba82f8c78046276230cbe0755e Mon Sep 17 00:00:00 2001 From: Trantorian Date: Wed, 24 Apr 2024 12:39:17 +0000 Subject: [PATCH 3/3] feat(storage): added `get_at` to query storage at a specifi commit --- src/changes.rs | 42 +++++++++++++++++++++++------------------ src/key_value_db.rs | 34 ++++++++++++++++++++++++++++++++- src/lib.rs | 13 +++++++++++++ src/trie/merkle_tree.rs | 25 ++++++++++++++++++++++++ 4 files changed, 95 insertions(+), 19 deletions(-) diff --git a/src/changes.rs b/src/changes.rs index 3018736..5de582b 100644 --- a/src/changes.rs +++ b/src/changes.rs @@ -39,11 +39,9 @@ impl ChangeBatch { } pub fn serialize(&self, id: &ID) -> Vec<(Vec, &[u8])> { - let id = id.to_bytes(); self.0 .iter() .flat_map(|(change_key, change)| { - let key_slice = change_key.as_slice(); let mut changes = Vec::new(); if let Some(old_value) = &change.old_value { @@ -52,26 +50,12 @@ impl ChangeBatch { return changes; } } - let key = [ - id.as_slice(), - &[KEY_SEPARATOR], - key_slice, - &[change_key.into()], - &[OLD_VALUE], - ] - .concat(); + let key = key_old_value(id, change_key); changes.push((key, old_value.as_slice())); } if let Some(new_value) = &change.new_value { - let key = [ - id.as_slice(), - &[KEY_SEPARATOR], - key_slice, - &[change_key.into()], - &[NEW_VALUE], - ] - .concat(); + let key = key_new_value(id, change_key); changes.push((key, new_value.as_slice())); } changes @@ -116,6 +100,28 @@ impl ChangeBatch { } } +pub fn key_old_value(id: &ID, key: &TrieKey) -> Vec { + [ + id.to_bytes().as_slice(), + &[KEY_SEPARATOR], + key.as_slice(), + &[key.into()], + &[OLD_VALUE], + ] + .concat() +} + +pub fn key_new_value(id: &ID, key: &TrieKey) -> Vec { + [ + id.to_bytes().as_slice(), + &[KEY_SEPARATOR], + key.as_slice(), + &[key.into()], + &[NEW_VALUE], + ] + .concat() +} + #[cfg_attr(feature = "bench", derive(Clone))] pub struct ChangeStore where diff --git a/src/key_value_db.rs b/src/key_value_db.rs index 2b4d099..c5e5f71 100644 --- a/src/key_value_db.rs +++ b/src/key_value_db.rs @@ -1,4 +1,4 @@ -use crate::{trie::merkle_tree::bytes_to_bitvec, Change as ExternChange}; +use crate::{changes::key_new_value, trie::merkle_tree::bytes_to_bitvec, Change as ExternChange}; #[cfg(not(feature = "std"))] use alloc::{collections::BTreeSet, format, string::ToString, vec::Vec}; use bitvec::{order::Msb0, vec::BitVec}; @@ -171,6 +171,38 @@ where Ok(self.db.get(&key.into())?) } + pub(crate) fn get_at( + &self, + key: &TrieKey, + id: ID, + ) -> Result>, BonsaiStorageError> { + trace!("Getting from KeyValueDB: {:?} at ID: {:?}", key, id); + + // makes sure given id exists + let Ok(id_position) = self.changes_store.id_queue.binary_search(&id) else { + return Err(BonsaiStorageError::Transaction(format!( + "invalid id {:?}", + id + ))); + }; + + // looking for the first storage insertion with given key + let iter = self + .changes_store + .id_queue + .iter() + .take(id_position + 1) + .rev(); + for id in iter { + let key = key_new_value(id, key); + if let Some(value) = self.db.get(&DatabaseKey::TrieLog(&key))? { + return Ok(Some(value)); + } + } + + Ok(None) + } + pub(crate) fn get_latest_id(&self) -> Option { self.changes_store.id_queue.back().cloned() } diff --git a/src/lib.rs b/src/lib.rs index 0272a72..b51edc8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -259,6 +259,19 @@ where self.tries.get(identifier, key) } + /// Gets a value in a trie at a given commit ID. + /// + /// Note that this is much faster that calling `revert_to1 + /// as it only reverts storage for a single key. + pub fn get_at( + &self, + identifier: &[u8], + key: &BitSlice, + id: ChangeID, + ) -> Result, BonsaiStorageError> { + self.tries.get_at(identifier, key, id) + } + /// Checks if the key exists in the trie. pub fn contains( &self, diff --git a/src/trie/merkle_tree.rs b/src/trie/merkle_tree.rs index 122d506..77452d0 100644 --- a/src/trie/merkle_tree.rs +++ b/src/trie/merkle_tree.rs @@ -130,6 +130,20 @@ impl MerkleTrees, + id: CommitID, + ) -> Result, BonsaiStorageError> { + let tree = self.trees.get(identifier); + if let Some(tree) = tree { + tree.get_at(&self.db, key, id) + } else { + Ok(None) + } + } + pub(crate) fn contains( &self, identifier: &[u8], @@ -961,6 +975,17 @@ impl MerkleTree { .map(|r| r.map(|opt| Felt::decode(&mut opt.as_slice()).unwrap())) } + pub fn get_at( + &self, + db: &KeyValueDB, + key: &BitSlice, + id: ID, + ) -> Result, BonsaiStorageError> { + let key = bitslice_to_bytes(key); + db.get_at(&TrieKey::new(&self.identifier, TrieKeyType::Flat, &key), id) + .map(|r| r.map(|opt| Felt::decode(&mut opt.as_slice()).unwrap())) + } + pub fn contains( &self, db: &KeyValueDB,