From 1e585c0c26f7323fb2715989a0c42b3f70a40dc5 Mon Sep 17 00:00:00 2001 From: AurelienFT Date: Thu, 28 Mar 2024 11:34:41 +0100 Subject: [PATCH] Remove misleading function and add a constructor for triekey --- src/trie/merkle_tree.rs | 56 +++++++++++++++++------------------------ src/trie/mod.rs | 2 +- src/trie/path.rs | 1 + src/trie/trie_db.rs | 14 +++++++++-- 4 files changed, 37 insertions(+), 36 deletions(-) diff --git a/src/trie/merkle_tree.rs b/src/trie/merkle_tree.rs index 8ad7025..c10d0d9 100644 --- a/src/trie/merkle_tree.rs +++ b/src/trie/merkle_tree.rs @@ -22,6 +22,7 @@ use crate::{error::BonsaiStorageError, id::Id, BonsaiDatabase, KeyValueDB}; use super::{ merkle_node::{BinaryNode, Direction, EdgeNode, Node, NodeHandle, NodeId}, path::Path, + trie_db::TrieKeyType, TrieKey, }; @@ -275,15 +276,6 @@ pub(crate) enum InsertOrRemove { Remove, } -/// Note to developers : Must be used to build a key for the database. -// Allow because needed for no-std -#[allow(clippy::ptr_arg)] -pub fn build_db_key(identifier: &Vec, key: &[u8]) -> Vec { - let mut db_key = identifier.clone(); - db_key.extend_from_slice(key); - db_key -} - impl MerkleTree { /// Less visible initialization for `MerkleTree` as the main entry points should be /// [`MerkleTree::::load`] for persistent trees and [`MerkleTree::empty`] for @@ -294,12 +286,12 @@ impl MerkleTree { identifier: Vec, ) -> Result> { let nodes_mapping: HashMap = HashMap::new(); - let root_node = db.get(&TrieKey::Trie(build_db_key(&identifier, &[])))?; + let root_node = db.get(&TrieKey::new(&identifier, TrieKeyType::Trie, &[]))?; let node = if let Some(root_node) = root_node { Node::decode(&mut root_node.as_slice())? } else { db.insert( - &TrieKey::Trie(build_db_key(&identifier, &[])), + &TrieKey::new(&identifier, TrieKeyType::Trie, &[]), &Node::Unresolved(Felt::ZERO).encode(), None, )?; @@ -358,7 +350,7 @@ impl MerkleTree { self.commit_subtree::(&mut updates, self.root_handle, Path(BitVec::new()))?; for (key, value) in mem::take(&mut self.cache_leaf_modified) { updates.insert( - TrieKey::Flat(build_db_key(&self.identifier, &key)), + TrieKey::new(&self.identifier, TrieKeyType::Flat, &key), match value { InsertOrRemove::Insert(value) => InsertOrRemove::Insert(value.encode()), InsertOrRemove::Remove => InsertOrRemove::Remove, @@ -404,7 +396,7 @@ impl MerkleTree { Unresolved(hash) => { if path.0.is_empty() { updates.insert( - TrieKey::Trie(build_db_key(&self.identifier, &[])), + TrieKey::new(&self.identifier, TrieKeyType::Trie, &[]), InsertOrRemove::Insert(Node::Unresolved(hash).encode()), ); Ok(hash) @@ -421,13 +413,9 @@ impl MerkleTree { binary.hash = Some(hash); binary.left = NodeHandle::Hash(left_hash); binary.right = NodeHandle::Hash(right_hash); - let key_bytes = if path.0.is_empty() { - vec![] - } else { - [&[path.0.len() as u8], path.0.as_raw_slice()].concat() - }; + let key_bytes: Vec = path.into(); updates.insert( - TrieKey::Trie(build_db_key(&self.identifier, &key_bytes)), + TrieKey::new(&self.identifier, TrieKeyType::Trie, &key_bytes), InsertOrRemove::Insert(Node::Binary(binary).encode()), ); Ok(hash) @@ -450,13 +438,9 @@ impl MerkleTree { let hash = H::hash(&child_hash, &felt_path) + length; edge.hash = Some(hash); edge.child = NodeHandle::Hash(child_hash); - let key_bytes = if path.0.is_empty() { - vec![] - } else { - [&[path.0.len() as u8], path.0.as_raw_slice()].concat() - }; + let key_bytes: Vec = path.into(); updates.insert( - TrieKey::Trie(build_db_key(&self.identifier, &key_bytes)), + TrieKey::new(&self.identifier, TrieKeyType::Trie, &key_bytes), InsertOrRemove::Insert(Node::Edge(edge).encode()), ); Ok(hash) @@ -485,9 +469,11 @@ impl MerkleTree { return Ok(()); } } - if let Some(value_db) = - db.get(&TrieKey::Flat(build_db_key(&self.identifier, &key_bytes)))? - { + if let Some(value_db) = db.get(&TrieKey::new( + &self.identifier, + TrieKeyType::Flat, + &key_bytes, + ))? { if value == Felt::decode(&mut value_db.as_slice()).unwrap() { return Ok(()); } @@ -671,7 +657,11 @@ impl MerkleTree { // Then we are done. let key_bytes = bitslice_to_bytes(key); if db - .get(&TrieKey::Flat(build_db_key(&self.identifier, &key_bytes)))? + .get(&TrieKey::new( + &self.identifier, + TrieKeyType::Flat, + &key_bytes, + ))? .is_none() && !self.cache_leaf_modified.contains_key(&key_bytes) { @@ -702,7 +692,7 @@ impl MerkleTree { last_binary_path = new_path; let path: Vec = (&last_binary_path).into(); self.death_row - .push(TrieKey::Trie(build_db_key(&self.identifier, &path))); + .push(TrieKey::new(&self.identifier, TrieKeyType::Trie, &path)); } } !node.is_binary() @@ -792,7 +782,7 @@ impl MerkleTree { InsertOrRemove::Insert(value) => return Ok(Some(*value)), } } - db.get(&TrieKey::Flat(build_db_key(&self.identifier, &key))) + db.get(&TrieKey::new(&self.identifier, TrieKeyType::Flat, &key)) .map(|r| r.map(|opt| Felt::decode(&mut opt.as_slice()).unwrap())) } @@ -808,7 +798,7 @@ impl MerkleTree { InsertOrRemove::Insert(_) => return Ok(true), } } - db.contains(&TrieKey::Flat(build_db_key(&self.identifier, &key))) + db.contains(&TrieKey::new(&self.identifier, TrieKeyType::Flat, &key)) } /// Returns the list of nodes along the path. @@ -1124,7 +1114,7 @@ impl MerkleTree { path: &Path, ) -> Result, BonsaiStorageError> { let path: Vec = path.into(); - db.get(&TrieKey::Trie(build_db_key(&self.identifier, &path)))? + db.get(&TrieKey::new(&self.identifier, TrieKeyType::Trie, &path))? .map(|node| { Node::decode(&mut node.as_slice()).map_err(|err| { BonsaiStorageError::Trie(format!("Couldn't decode node: {}", err)) diff --git a/src/trie/mod.rs b/src/trie/mod.rs index 0d06bb8..dc6e304 100644 --- a/src/trie/mod.rs +++ b/src/trie/mod.rs @@ -3,4 +3,4 @@ pub mod merkle_tree; mod path; mod trie_db; -pub use trie_db::TrieKey; +pub(crate) use trie_db::TrieKey; diff --git a/src/trie/path.rs b/src/trie/path.rs index 1c5a8e8..b5f8cc4 100644 --- a/src/trie/path.rs +++ b/src/trie/path.rs @@ -104,6 +104,7 @@ impl Path { } } +/// Convert Path to Vec can be used, for example, to create keys for the database impl From for Vec { fn from(path: Path) -> Self { let key = if path.0.is_empty() { diff --git a/src/trie/trie_db.rs b/src/trie/trie_db.rs index b7debf2..3dfafc4 100644 --- a/src/trie/trie_db.rs +++ b/src/trie/trie_db.rs @@ -4,13 +4,14 @@ use crate::bonsai_database::DatabaseKey; use alloc::vec::Vec; /// Key in the database of the different elements that are used in the storage of the trie data. +/// Use `new` function to create a new key. #[derive(Debug, Clone, Hash, PartialEq, Eq)] -pub enum TrieKey { +pub(crate) enum TrieKey { Trie(Vec), Flat(Vec), } -enum TrieKeyType { +pub(crate) enum TrieKeyType { Trie = 0, Flat = 1, } @@ -34,6 +35,15 @@ impl From<&TrieKey> for u8 { } impl TrieKey { + pub fn new(identifier: &[u8], key_type: TrieKeyType, key: &[u8]) -> Self { + let mut final_key = identifier.to_owned(); + final_key.extend_from_slice(key); + match key_type { + TrieKeyType::Trie => TrieKey::Trie(final_key), + TrieKeyType::Flat => TrieKey::Flat(final_key), + } + } + pub fn from_variant_and_bytes(variant: u8, bytes: Vec) -> Self { match variant { x if x == TrieKeyType::Trie as u8 => TrieKey::Trie(bytes),