diff --git a/grovedb/src/batch/mod.rs b/grovedb/src/batch/mod.rs index 22f859d5..5d463f91 100644 --- a/grovedb/src/batch/mod.rs +++ b/grovedb/src/batch/mod.rs @@ -80,7 +80,7 @@ use grovedb_merk::{ value_hash, NULL_HASH, }, CryptoHash, Error as MerkError, Merk, MerkType, RootHashKeyAndSum, - TreeFeatureType::{BasicMerk, SummedMerk}, + TreeFeatureType::{BasicMerkNode, SummedMerkNode}, }; use grovedb_path::SubtreePath; use grovedb_storage::{ @@ -766,8 +766,12 @@ where if recursions_allowed == 1 { let referenced_element_value_hash_opt = cost_return_on_error!( &mut cost, - merk.get_value_hash(key.as_ref(), true) - .map_err(|e| Error::CorruptedData(e.to_string())) + merk.get_value_hash( + key.as_ref(), + true, + Some(Element::value_defined_cost_for_serialized_value) + ) + .map_err(|e| Error::CorruptedData(e.to_string())) ); let referenced_element_value_hash = cost_return_on_error!( @@ -806,8 +810,12 @@ where // change in the batch. let referenced_element = cost_return_on_error!( &mut cost, - merk.get(key.as_ref(), true) - .map_err(|e| Error::CorruptedData(e.to_string())) + merk.get( + key.as_ref(), + true, + Some(Element::value_defined_cost_for_serialized_value) + ) + .map_err(|e| Error::CorruptedData(e.to_string())) ); let referenced_element = cost_return_on_error_no_add!( @@ -1130,14 +1138,18 @@ where } else { let value = cost_return_on_error!( &mut cost, - merk.get(key_info.as_slice(), true) - .map(|result_value| result_value - .map_err(Error::MerkError) - .and_then(|maybe_value| maybe_value.ok_or( - Error::InvalidInput( - "trying to refresh a non existing reference", - ) - ))) + merk.get( + key_info.as_slice(), + true, + Some(Element::value_defined_cost_for_serialized_value) + ) + .map( + |result_value| result_value.map_err(Error::MerkError).and_then( + |maybe_value| maybe_value.ok_or(Error::InvalidInput( + "trying to refresh a non existing reference", + )) + ) + ) ); cost_return_on_error_no_add!( &cost, @@ -1155,9 +1167,9 @@ where }; let merk_feature_type = if is_sum_tree { - SummedMerk(0) + SummedMerkNode(0) } else { - BasicMerk + BasicMerkNode }; let path_reference = cost_return_on_error!( @@ -1276,7 +1288,7 @@ where } cost_return_on_error!( &mut cost, - merk.apply_unchecked::<_, Vec, _, _, _>( + merk.apply_unchecked::<_, Vec, _, _, _, _>( &batch_operations, &[], Some(batch_apply_options.as_merk_options()), @@ -1284,6 +1296,7 @@ where Element::specialized_costs_for_key_value(key, value, is_sum_tree) .map_err(|e| MerkError::ClientCorruptionError(e.to_string())) }, + Some(&Element::value_defined_cost_for_serialized_value), &mut |storage_costs, old_value, new_value| { // todo: change the flags without full deserialization let old_element = Element::deserialize(old_value.as_slice()) @@ -1779,13 +1792,16 @@ impl GroveDb { ); let is_sum_tree = element.is_sum_tree(); if let Element::Tree(root_key, _) | Element::SumTree(root_key, ..) = element { - Merk::open_layered_with_root_key(storage, root_key, is_sum_tree) - .map_err(|_| { - Error::CorruptedData( - "cannot open a subtree with given root key".to_owned(), - ) - }) - .add_cost(cost) + Merk::open_layered_with_root_key( + storage, + root_key, + is_sum_tree, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .map_err(|_| { + Error::CorruptedData("cannot open a subtree with given root key".to_owned()) + }) + .add_cost(cost) } else { Err(Error::CorruptedPath( "cannot open a subtree as parent exists but is not a tree", @@ -1797,9 +1813,13 @@ impl GroveDb { if new_merk { Ok(Merk::open_empty(storage, MerkType::BaseMerk, false)).wrap_with_cost(cost) } else { - Merk::open_base(storage, false) - .map_err(|_| Error::CorruptedData("cannot open a the root subtree".to_owned())) - .add_cost(cost) + Merk::open_base( + storage, + false, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .map_err(|_| Error::CorruptedData("cannot open a the root subtree".to_owned())) + .add_cost(cost) } } } @@ -1835,11 +1855,16 @@ impl GroveDb { ); let is_sum_tree = element.is_sum_tree(); if let Element::Tree(root_key, _) | Element::SumTree(root_key, ..) = element { - Merk::open_layered_with_root_key(storage, root_key, is_sum_tree) - .map_err(|_| { - Error::CorruptedData("cannot open a subtree with given root key".to_owned()) - }) - .add_cost(local_cost) + Merk::open_layered_with_root_key( + storage, + root_key, + is_sum_tree, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .map_err(|_| { + Error::CorruptedData("cannot open a subtree with given root key".to_owned()) + }) + .add_cost(local_cost) } else { Err(Error::CorruptedData( "cannot open a subtree as parent exists but is not a tree".to_owned(), @@ -1847,9 +1872,13 @@ impl GroveDb { .wrap_with_cost(local_cost) } } else { - Merk::open_base(storage, false) - .map_err(|_| Error::CorruptedData("cannot open a subtree".to_owned())) - .add_cost(local_cost) + Merk::open_base( + storage, + false, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .map_err(|_| Error::CorruptedData("cannot open a subtree".to_owned())) + .add_cost(local_cost) } } diff --git a/grovedb/src/element/delete.rs b/grovedb/src/element/delete.rs index 9766c546..92087bc4 100644 --- a/grovedb/src/element/delete.rs +++ b/grovedb/src/element/delete.rs @@ -57,10 +57,16 @@ impl Element { }; let batch = [(key, op)]; let uses_sum_nodes = merk.is_sum_tree; - merk.apply_with_specialized_costs::<_, Vec>(&batch, &[], merk_options, &|key, value| { - Self::specialized_costs_for_key_value(key, value, uses_sum_nodes) - .map_err(|e| MerkError::ClientCorruptionError(e.to_string())) - }) + merk.apply_with_specialized_costs::<_, Vec>( + &batch, + &[], + merk_options, + &|key, value| { + Self::specialized_costs_for_key_value(key, value, uses_sum_nodes) + .map_err(|e| MerkError::ClientCorruptionError(e.to_string())) + }, + Some(&Element::value_defined_cost_for_serialized_value), + ) .map_err(|e| Error::CorruptedData(e.to_string())) } @@ -97,6 +103,7 @@ impl Element { Self::specialized_costs_for_key_value(key, value, uses_sum_nodes) .map_err(|e| MerkError::ClientCorruptionError(e.to_string())) }, + Some(&Element::value_defined_cost_for_serialized_value), &mut |_costs, _old_value, _value| Ok((false, None)), sectioned_removal, ) diff --git a/grovedb/src/element/exists.rs b/grovedb/src/element/exists.rs index d0bf5ecb..c3bf61fa 100644 --- a/grovedb/src/element/exists.rs +++ b/grovedb/src/element/exists.rs @@ -48,7 +48,10 @@ impl Element { merk: &mut Merk, key: K, ) -> CostResult { - merk.exists(key.as_ref()) - .map_err(|e| Error::CorruptedData(e.to_string())) + merk.exists( + key.as_ref(), + Some(&Element::value_defined_cost_for_serialized_value), + ) + .map_err(|e| Error::CorruptedData(e.to_string())) } } diff --git a/grovedb/src/element/get.rs b/grovedb/src/element/get.rs index 57cd4012..b6f75b10 100644 --- a/grovedb/src/element/get.rs +++ b/grovedb/src/element/get.rs @@ -78,8 +78,12 @@ impl Element { let value_opt = cost_return_on_error!( &mut cost, - merk.get(key.as_ref(), allow_cache) - .map_err(|e| Error::CorruptedData(e.to_string())) + merk.get( + key.as_ref(), + allow_cache, + Some(&Element::value_defined_cost_for_serialized_value) + ) + .map_err(|e| Error::CorruptedData(e.to_string())) ); let element = cost_return_on_error_no_add!( &cost, @@ -230,8 +234,12 @@ impl Element { let value_hash = cost_return_on_error!( &mut cost, - merk.get_value_hash(key.as_ref(), allow_cache) - .map_err(|e| Error::CorruptedData(e.to_string())) + merk.get_value_hash( + key.as_ref(), + allow_cache, + Some(&Element::value_defined_cost_for_serialized_value) + ) + .map_err(|e| Error::CorruptedData(e.to_string())) ); Ok(value_hash).wrap_with_cost(cost) @@ -253,7 +261,13 @@ mod tests { let ctx = storage .get_storage_context(SubtreePath::empty(), Some(&batch)) .unwrap(); - let mut merk = Merk::open_base(ctx, false).unwrap().unwrap(); + let mut merk = Merk::open_base( + ctx, + false, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .unwrap() + .unwrap(); Element::empty_tree() .insert(&mut merk, b"mykey", None) .unwrap() @@ -271,7 +285,13 @@ mod tests { let ctx = storage .get_storage_context(SubtreePath::empty(), None) .unwrap(); - let mut merk = Merk::open_base(ctx, false).unwrap().unwrap(); + let mut merk = Merk::open_base( + ctx, + false, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .unwrap() + .unwrap(); assert_eq!( Element::get(&merk, b"another-key", true) diff --git a/grovedb/src/element/helpers.rs b/grovedb/src/element/helpers.rs index 394ec3da..85ce8f4e 100644 --- a/grovedb/src/element/helpers.rs +++ b/grovedb/src/element/helpers.rs @@ -29,11 +29,15 @@ //! Helpers //! Implements helper functions in Element +use grovedb_merk::tree::kv::{ + ValueDefinedCostType, + ValueDefinedCostType::{LayeredValueDefinedCost, SpecializedValueDefinedCost}, +}; #[cfg(feature = "full")] use grovedb_merk::{ tree::{kv::KV, TreeNode}, TreeFeatureType, - TreeFeatureType::{BasicMerk, SummedMerk}, + TreeFeatureType::{BasicMerkNode, SummedMerkNode}, }; #[cfg(feature = "full")] use integer_encoding::VarInt; @@ -114,8 +118,8 @@ impl Element { /// Get the tree feature type pub fn get_feature_type(&self, parent_is_sum_tree: bool) -> Result { match parent_is_sum_tree { - true => Ok(SummedMerk(self.sum_value_or_default())), - false => Ok(BasicMerk), + true => Ok(SummedMerkNode(self.sum_value_or_default())), + false => Ok(BasicMerkNode), } } @@ -307,13 +311,44 @@ impl Element { )), } } + + #[cfg(feature = "full")] + /// Get the value defined cost for a serialized value + pub fn value_defined_cost(&self) -> Option { + let Some(value_cost) = self.get_specialized_cost().ok() else { + return None; + }; + + let cost = value_cost + + self.get_flags().as_ref().map_or(0, |flags| { + let flags_len = flags.len() as u32; + flags_len + flags_len.required_space() as u32 + }); + match self { + Element::Tree(..) => Some(LayeredValueDefinedCost(cost)), + Element::SumTree(..) => Some(LayeredValueDefinedCost(cost)), + Element::SumItem(..) => Some(SpecializedValueDefinedCost(cost)), + _ => None, + } + } + + #[cfg(feature = "full")] + /// Get the value defined cost for a serialized value + pub fn value_defined_cost_for_serialized_value(value: &[u8]) -> Option { + let element = Element::deserialize(value).ok()?; + element.value_defined_cost() + } } #[cfg(feature = "full")] /// Decode from bytes pub fn raw_decode(bytes: &[u8]) -> Result { - let tree = - TreeNode::decode_raw(bytes, vec![]).map_err(|e| Error::CorruptedData(e.to_string()))?; + let tree = TreeNode::decode_raw( + bytes, + vec![], + Some(Element::value_defined_cost_for_serialized_value), + ) + .map_err(|e| Error::CorruptedData(e.to_string()))?; let element: Element = Element::deserialize(tree.value_as_slice())?; Ok(element) } diff --git a/grovedb/src/element/insert.rs b/grovedb/src/element/insert.rs index e500b41a..2ba7f92d 100644 --- a/grovedb/src/element/insert.rs +++ b/grovedb/src/element/insert.rs @@ -92,6 +92,7 @@ impl Element { Self::specialized_costs_for_key_value(key, value, uses_sum_nodes) .map_err(|e| MerkError::ClientCorruptionError(e.to_string())) }, + Some(&Element::value_defined_cost_for_serialized_value), ) .map_err(|e| Error::CorruptedData(e.to_string())) } @@ -287,6 +288,7 @@ impl Element { Self::specialized_costs_for_key_value(key, value, uses_sum_nodes) .map_err(|e| MerkError::ClientCorruptionError(e.to_string())) }, + Some(&Element::value_defined_cost_for_serialized_value), ) .map_err(|e| Error::CorruptedData(e.to_string())) } @@ -356,6 +358,7 @@ impl Element { Self::specialized_costs_for_key_value(key, value, uses_sum_nodes) .map_err(|e| MerkError::ClientCorruptionError(e.to_string())) }, + Some(&Element::value_defined_cost_for_serialized_value), ) .map_err(|e| Error::CorruptedData(e.to_string())) } diff --git a/grovedb/src/estimated_costs/average_case_costs.rs b/grovedb/src/estimated_costs/average_case_costs.rs index e93e1b8f..d93b6451 100644 --- a/grovedb/src/estimated_costs/average_case_costs.rs +++ b/grovedb/src/estimated_costs/average_case_costs.rs @@ -461,7 +461,7 @@ mod test { use grovedb_costs::OperationCost; use grovedb_merk::{ estimated_costs::average_case_costs::add_average_case_get_merk_node, - test_utils::make_batch_seq, Merk, + test_utils::make_batch_seq, tree::kv::ValueDefinedCostType, Merk, }; use grovedb_storage::{ rocksdb_storage::RocksDbStorage, worst_case_costs::WorstKeyLength, Storage, StorageBatch, @@ -487,6 +487,7 @@ mod test { .get_storage_context(EMPTY_PATH, Some(&batch)) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .expect("cannot open merk"); @@ -505,6 +506,7 @@ mod test { let merk = Merk::open_base( storage.get_storage_context(EMPTY_PATH, None).unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .expect("cannot open merk"); @@ -514,7 +516,11 @@ mod test { // 2. Left link exists // 3. Right link exists // Based on merk's avl rotation algorithm node is key 8 satisfies this - let node_result = merk.get(&8_u64.to_be_bytes(), true); + let node_result = merk.get( + &8_u64.to_be_bytes(), + true, + None::<&fn(&[u8]) -> Option>, + ); // By tweaking the max element size, we can adapt the average case function to // this scenario. make_batch_seq creates values that are 60 bytes in size diff --git a/grovedb/src/estimated_costs/worst_case_costs.rs b/grovedb/src/estimated_costs/worst_case_costs.rs index d72b86e8..5a72a405 100644 --- a/grovedb/src/estimated_costs/worst_case_costs.rs +++ b/grovedb/src/estimated_costs/worst_case_costs.rs @@ -411,6 +411,7 @@ mod test { use grovedb_merk::{ estimated_costs::worst_case_costs::add_worst_case_get_merk_node, test_utils::{empty_path_merk, empty_path_merk_read_only, make_batch_seq}, + tree::kv::ValueDefinedCostType, }; use grovedb_storage::{ rocksdb_storage::{test_utils::TempStorage, RocksDbStorage}, @@ -451,7 +452,11 @@ mod test { // 2. Left link exists // 3. Right link exists // Based on merk's avl rotation algorithm node is key 8 satisfies this - let node_result = merk.get(&8_u64.to_be_bytes(), true); + let node_result = merk.get( + &8_u64.to_be_bytes(), + true, + None::<&fn(&[u8]) -> Option>, + ); // By tweaking the max element size, we can adapt the worst case function to // this scenario. make_batch_seq creates values that are 60 bytes in size diff --git a/grovedb/src/lib.rs b/grovedb/src/lib.rs index 87ca2443..30ca2e91 100644 --- a/grovedb/src/lib.rs +++ b/grovedb/src/lib.rs @@ -193,6 +193,7 @@ pub use grovedb_merk::estimated_costs::{ pub use grovedb_merk::proofs::query::query_item::QueryItem; #[cfg(any(feature = "full", feature = "verify"))] pub use grovedb_merk::proofs::Query; +use grovedb_merk::tree::kv::ValueDefinedCostType; #[cfg(feature = "full")] use grovedb_merk::{ self, @@ -289,11 +290,16 @@ impl GroveDb { ); let is_sum_tree = element.is_sum_tree(); if let Element::Tree(root_key, _) | Element::SumTree(root_key, ..) = element { - Merk::open_layered_with_root_key(storage, root_key, is_sum_tree) - .map_err(|_| { - Error::CorruptedData("cannot open a subtree with given root key".to_owned()) - }) - .add_cost(cost) + Merk::open_layered_with_root_key( + storage, + root_key, + is_sum_tree, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .map_err(|_| { + Error::CorruptedData("cannot open a subtree with given root key".to_owned()) + }) + .add_cost(cost) } else { Err(Error::CorruptedPath( "cannot open a subtree as parent exists but is not a tree", @@ -301,9 +307,13 @@ impl GroveDb { .wrap_with_cost(cost) } } else { - Merk::open_base(storage, false) - .map_err(|_| Error::CorruptedData("cannot open a the root subtree".to_owned())) - .add_cost(cost) + Merk::open_base( + storage, + false, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .map_err(|_| Error::CorruptedData("cannot open a the root subtree".to_owned())) + .add_cost(cost) } } @@ -340,20 +350,29 @@ impl GroveDb { .unwrap()?; let is_sum_tree = element.is_sum_tree(); if let Element::Tree(root_key, _) | Element::SumTree(root_key, ..) = element { - Merk::open_layered_with_root_key(storage, root_key, is_sum_tree) - .map_err(|_| { - Error::CorruptedData("cannot open a subtree with given root key".to_owned()) - }) - .unwrap() + Merk::open_layered_with_root_key( + storage, + root_key, + is_sum_tree, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .map_err(|_| { + Error::CorruptedData("cannot open a subtree with given root key".to_owned()) + }) + .unwrap() } else { Err(Error::CorruptedPath( "cannot open a subtree as parent exists but is not a tree", )) } } else { - Merk::open_base(storage, false) - .map_err(|_| Error::CorruptedData("cannot open a the root subtree".to_owned())) - .unwrap() + Merk::open_base( + storage, + false, + None::<&fn(&[u8]) -> Option>, + ) + .map_err(|_| Error::CorruptedData("cannot open a the root subtree".to_owned())) + .unwrap() } } @@ -391,11 +410,16 @@ impl GroveDb { ); let is_sum_tree = element.is_sum_tree(); if let Element::Tree(root_key, _) | Element::SumTree(root_key, ..) = element { - Merk::open_layered_with_root_key(storage, root_key, is_sum_tree) - .map_err(|_| { - Error::CorruptedData("cannot open a subtree with given root key".to_owned()) - }) - .add_cost(cost) + Merk::open_layered_with_root_key( + storage, + root_key, + is_sum_tree, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .map_err(|_| { + Error::CorruptedData("cannot open a subtree with given root key".to_owned()) + }) + .add_cost(cost) } else { Err(Error::CorruptedPath( "cannot open a subtree as parent exists but is not a tree", @@ -403,9 +427,13 @@ impl GroveDb { .wrap_with_cost(cost) } } else { - Merk::open_base(storage, false) - .map_err(|_| Error::CorruptedData("cannot open a the root subtree".to_owned())) - .add_cost(cost) + Merk::open_base( + storage, + false, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .map_err(|_| Error::CorruptedData("cannot open a the root subtree".to_owned())) + .add_cost(cost) } } @@ -675,7 +703,11 @@ impl GroveDb { key: K, ) -> CostResult { subtree - .get(key.as_ref(), true) + .get( + key.as_ref(), + true, + Some(&Element::value_defined_cost_for_serialized_value), + ) .map_err(|_| { Error::InvalidPath("can't find subtree in parent during propagation".to_owned()) }) @@ -831,7 +863,11 @@ impl GroveDb { let element = raw_decode(&element_value).unwrap(); if element.is_tree() { let (kv_value, element_value_hash) = merk - .get_value_and_value_hash(&key, true) + .get_value_and_value_hash( + &key, + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .unwrap() .unwrap(); diff --git a/grovedb/src/operations/delete/mod.rs b/grovedb/src/operations/delete/mod.rs index fb2ce5ce..59092697 100644 --- a/grovedb/src/operations/delete/mod.rs +++ b/grovedb/src/operations/delete/mod.rs @@ -495,7 +495,8 @@ impl GroveDb { Merk::open_layered_with_root_key( storage, subtree_to_delete_from.root_key(), - element.is_sum_tree() + element.is_sum_tree(), + Some(&Element::value_defined_cost_for_serialized_value), ) .map_err(|_| { Error::CorruptedData("cannot open a subtree with given root key".to_owned()) diff --git a/grovedb/src/operations/insert/mod.rs b/grovedb/src/operations/insert/mod.rs index 9b83ff87..d691ea34 100644 --- a/grovedb/src/operations/insert/mod.rs +++ b/grovedb/src/operations/insert/mod.rs @@ -218,7 +218,11 @@ impl GroveDb { let maybe_element_bytes = cost_return_on_error!( &mut cost, subtree_to_insert_into - .get(key, true) + .get( + key, + true, + Some(&Element::value_defined_cost_for_serialized_value) + ) .map_err(|e| Error::CorruptedData(e.to_string())) ); if let Some(element_bytes) = maybe_element_bytes { @@ -353,7 +357,11 @@ impl GroveDb { let maybe_element_bytes = cost_return_on_error!( &mut cost, subtree_to_insert_into - .get(key, true) + .get( + key, + true, + Some(&Element::value_defined_cost_for_serialized_value) + ) .map_err(|e| Error::CorruptedData(e.to_string())) ); if let Some(element_bytes) = maybe_element_bytes { @@ -867,6 +875,11 @@ mod tests { // Child Heights 2 // Total 37 + 85 + 48 = 170 + + // replaced bytes + // 133 for key1 (higher node/same merk level) + // ? + assert_eq!( cost, OperationCost { diff --git a/grovedb/src/replication.rs b/grovedb/src/replication.rs index 86c1c3f0..898f5ff1 100644 --- a/grovedb/src/replication.rs +++ b/grovedb/src/replication.rs @@ -181,6 +181,7 @@ impl<'db> Restorer<'db> { .get_immediate_storage_context(SubtreePath::empty(), tx) .unwrap(), false, + Some(&Element::value_defined_cost_for_serialized_value), ) .unwrap() .map_err(|e| RestorerError(e.to_string()))?, diff --git a/grovedb/src/tests/mod.rs b/grovedb/src/tests/mod.rs index e078237d..922f5e9b 100644 --- a/grovedb/src/tests/mod.rs +++ b/grovedb/src/tests/mod.rs @@ -2365,9 +2365,13 @@ fn test_find_subtrees() { fn test_root_subtree_has_root_key() { let db = make_test_grovedb(); let storage = db.db.get_storage_context(EMPTY_PATH, None).unwrap(); - let root_merk = Merk::open_base(storage, false) - .unwrap() - .expect("expected to get root merk"); + let root_merk = Merk::open_base( + storage, + false, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .unwrap() + .expect("expected to get root merk"); let (_, root_key, _) = root_merk .root_hash_key_and_sum() .unwrap() @@ -2457,10 +2461,14 @@ fn test_get_subtree() { .db .get_storage_context([TEST_LEAF, b"key1", b"key2"].as_ref().into(), None) .unwrap(); - let subtree = - Merk::open_layered_with_root_key(subtree_storage, Some(b"key3".to_vec()), false) - .unwrap() - .expect("cannot open merk"); + let subtree = Merk::open_layered_with_root_key( + subtree_storage, + Some(b"key3".to_vec()), + false, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .unwrap() + .expect("cannot open merk"); let result_element = Element::get(&subtree, b"key3", true).unwrap().unwrap(); assert_eq!(result_element, Element::new_item(b"ayy".to_vec())); } @@ -2497,9 +2505,14 @@ fn test_get_subtree() { &transaction, ) .unwrap(); - let subtree = Merk::open_layered_with_root_key(subtree_storage, Some(b"key4".to_vec()), false) - .unwrap() - .expect("cannot open merk"); + let subtree = Merk::open_layered_with_root_key( + subtree_storage, + Some(b"key4".to_vec()), + false, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .unwrap() + .expect("cannot open merk"); let result_element = Element::get(&subtree, b"key4", true).unwrap().unwrap(); assert_eq!(result_element, Element::new_item(b"ayy".to_vec())); @@ -2509,9 +2522,14 @@ fn test_get_subtree() { .db .get_storage_context([TEST_LEAF, b"key1", b"key2"].as_ref().into(), None) .unwrap(); - let subtree = Merk::open_layered_with_root_key(subtree_storage, Some(b"key3".to_vec()), false) - .unwrap() - .expect("cannot open merk"); + let subtree = Merk::open_layered_with_root_key( + subtree_storage, + Some(b"key3".to_vec()), + false, + Some(&Element::value_defined_cost_for_serialized_value), + ) + .unwrap() + .expect("cannot open merk"); let result_element = Element::get(&subtree, b"key3", true).unwrap().unwrap(); assert_eq!(result_element, Element::new_item(b"ayy".to_vec())); } diff --git a/grovedb/src/tests/sum_tree_tests.rs b/grovedb/src/tests/sum_tree_tests.rs index 3bc6896e..6c4a7589 100644 --- a/grovedb/src/tests/sum_tree_tests.rs +++ b/grovedb/src/tests/sum_tree_tests.rs @@ -30,7 +30,8 @@ use grovedb_merk::{ proofs::Query, - TreeFeatureType::{BasicMerk, SummedMerk}, + tree::kv::ValueDefinedCostType, + TreeFeatureType::{BasicMerkNode, SummedMerkNode}, }; use grovedb_storage::StorageBatch; @@ -266,28 +267,44 @@ fn test_homogenous_node_type_in_sum_trees_and_regular_trees() { .unwrap() .expect("should open tree"); assert!(matches!( - merk.get_feature_type(b"item1", true) - .unwrap() - .expect("node should exist"), - Some(SummedMerk(30)) + merk.get_feature_type( + b"item1", + true, + None::<&fn(&[u8]) -> Option> + ) + .unwrap() + .expect("node should exist"), + Some(SummedMerkNode(30)) )); assert!(matches!( - merk.get_feature_type(b"item2", true) - .unwrap() - .expect("node should exist"), - Some(SummedMerk(10)) + merk.get_feature_type( + b"item2", + true, + None::<&fn(&[u8]) -> Option> + ) + .unwrap() + .expect("node should exist"), + Some(SummedMerkNode(10)) )); assert!(matches!( - merk.get_feature_type(b"item3", true) - .unwrap() - .expect("node should exist"), - Some(SummedMerk(0)) + merk.get_feature_type( + b"item3", + true, + None::<&fn(&[u8]) -> Option> + ) + .unwrap() + .expect("node should exist"), + Some(SummedMerkNode(0)) )); assert!(matches!( - merk.get_feature_type(b"item4", true) - .unwrap() - .expect("node should exist"), - Some(SummedMerk(0)) + merk.get_feature_type( + b"item4", + true, + None::<&fn(&[u8]) -> Option> + ) + .unwrap() + .expect("node should exist"), + Some(SummedMerkNode(0)) )); assert_eq!(merk.sum().expect("expected to get sum"), Some(40)); @@ -326,16 +343,24 @@ fn test_homogenous_node_type_in_sum_trees_and_regular_trees() { .unwrap() .expect("should open tree"); assert!(matches!( - merk.get_feature_type(b"item1", true) - .unwrap() - .expect("node should exist"), - Some(BasicMerk) + merk.get_feature_type( + b"item1", + true, + Some(&Element::value_defined_cost_for_serialized_value) + ) + .unwrap() + .expect("node should exist"), + Some(BasicMerkNode) )); assert!(matches!( - merk.get_feature_type(b"item2", true) - .unwrap() - .expect("node should exist"), - Some(BasicMerk) + merk.get_feature_type( + b"item2", + true, + Some(&Element::value_defined_cost_for_serialized_value) + ) + .unwrap() + .expect("node should exist"), + Some(BasicMerkNode) )); assert_eq!(merk.sum().expect("expected to get sum"), None); } @@ -582,10 +607,14 @@ fn test_sum_tree_propagation() { .expect("should open tree"); assert!(matches!( test_leaf_merk - .get_feature_type(b"key", true) + .get_feature_type( + b"key", + true, + Some(&Element::value_defined_cost_for_serialized_value) + ) .unwrap() .expect("node should exist"), - Some(BasicMerk) + Some(BasicMerkNode) )); let parent_sum_tree = db @@ -594,12 +623,16 @@ fn test_sum_tree_propagation() { .expect("should open tree"); assert!(matches!( parent_sum_tree - .get_feature_type(b"tree2", true) + .get_feature_type( + b"tree2", + true, + Some(&Element::value_defined_cost_for_serialized_value) + ) .unwrap() .expect("node should exist"), - Some(SummedMerk(15)) /* 15 because the child sum tree has one sum item of - * value 5 and - * another of value 10 */ + Some(SummedMerkNode(15)) /* 15 because the child sum tree has one sum item of + * value 5 and + * another of value 10 */ )); let child_sum_tree = db @@ -611,33 +644,49 @@ fn test_sum_tree_propagation() { .expect("should open tree"); assert!(matches!( child_sum_tree - .get_feature_type(b"item1", true) + .get_feature_type( + b"item1", + true, + None::<&fn(&[u8]) -> Option> + ) .unwrap() .expect("node should exist"), - Some(SummedMerk(0)) + Some(SummedMerkNode(0)) )); assert!(matches!( child_sum_tree - .get_feature_type(b"sumitem1", true) + .get_feature_type( + b"sumitem1", + true, + None::<&fn(&[u8]) -> Option> + ) .unwrap() .expect("node should exist"), - Some(SummedMerk(5)) + Some(SummedMerkNode(5)) )); assert!(matches!( child_sum_tree - .get_feature_type(b"sumitem2", true) + .get_feature_type( + b"sumitem2", + true, + None::<&fn(&[u8]) -> Option> + ) .unwrap() .expect("node should exist"), - Some(SummedMerk(10)) + Some(SummedMerkNode(10)) )); // TODO: should references take the sum of the referenced element?? assert!(matches!( child_sum_tree - .get_feature_type(b"item2", true) + .get_feature_type( + b"item2", + true, + None::<&fn(&[u8]) -> Option> + ) .unwrap() .expect("node should exist"), - Some(SummedMerk(0)) + Some(SummedMerkNode(0)) )); } @@ -673,17 +722,25 @@ fn test_sum_tree_with_batches() { assert!(matches!( sum_tree - .get_feature_type(b"a", true) + .get_feature_type( + b"a", + true, + Some(&Element::value_defined_cost_for_serialized_value) + ) .unwrap() .expect("node should exist"), - Some(SummedMerk(0)) + Some(SummedMerkNode(0)) )); assert!(matches!( sum_tree - .get_feature_type(b"b", true) + .get_feature_type( + b"b", + true, + Some(&Element::value_defined_cost_for_serialized_value) + ) .unwrap() .expect("node should exist"), - Some(SummedMerk(10)) + Some(SummedMerkNode(10)) )); // Create new batch to use existing tree @@ -703,10 +760,14 @@ fn test_sum_tree_with_batches() { .expect("should open tree"); assert!(matches!( sum_tree - .get_feature_type(b"c", true) + .get_feature_type( + b"c", + true, + None::<&fn(&[u8]) -> Option> + ) .unwrap() .expect("node should exist"), - Some(SummedMerk(10)) + Some(SummedMerkNode(10)) )); assert_eq!(sum_tree.sum().expect("expected to get sum"), Some(20)); diff --git a/grovedb/src/tests/tree_hashes_tests.rs b/grovedb/src/tests/tree_hashes_tests.rs index b67d1cf9..d2418132 100644 --- a/grovedb/src/tests/tree_hashes_tests.rs +++ b/grovedb/src/tests/tree_hashes_tests.rs @@ -28,7 +28,9 @@ //! Tree hashes tests -use grovedb_merk::tree::{combine_hash, kv_digest_to_kv_hash, node_hash, value_hash, NULL_HASH}; +use grovedb_merk::tree::{ + combine_hash, kv::ValueDefinedCostType, kv_digest_to_kv_hash, node_hash, value_hash, NULL_HASH, +}; use grovedb_storage::StorageBatch; use crate::{ @@ -58,19 +60,31 @@ fn test_node_hashes_when_inserting_item() { .expect("should open merk"); let (elem_value, elem_value_hash) = test_leaf_merk - .get_value_and_value_hash(b"key1", true) + .get_value_and_value_hash( + b"key1", + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get value hash") .expect("value hash should be some"); let elem_kv_hash = test_leaf_merk - .get_kv_hash(b"key1", true) + .get_kv_hash( + b"key1", + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get value hash") .expect("value hash should be some"); let elem_node_hash = test_leaf_merk - .get_hash(b"key1", true) + .get_hash( + b"key1", + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get value hash") .expect("value hash should be some"); @@ -110,19 +124,31 @@ fn test_tree_hashes_when_inserting_empty_tree() { .expect("should open merk"); let (elem_value, elem_value_hash) = test_leaf_merk - .get_value_and_value_hash(b"key1", true) + .get_value_and_value_hash( + b"key1", + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get value hash") .expect("value hash should be some"); let elem_kv_hash = test_leaf_merk - .get_kv_hash(b"key1", true) + .get_kv_hash( + b"key1", + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get value hash") .expect("value hash should be some"); let elem_node_hash = test_leaf_merk - .get_hash(b"key1", true) + .get_hash( + b"key1", + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get value hash") .expect("value hash should be some"); @@ -187,7 +213,11 @@ fn test_tree_hashes_when_inserting_empty_trees_twice_under_each_other() { // Let's first verify that the lowest nodes hashes are as we expect let (elem_value, elem_value_hash) = middle_merk_key1 - .get_value_and_value_hash(b"key2", true) + .get_value_and_value_hash( + b"key2", + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get value hash") .expect("value hash should be some"); @@ -210,7 +240,11 @@ fn test_tree_hashes_when_inserting_empty_trees_twice_under_each_other() { assert_eq!(elem_value_hash, combined_value_hash_key2); let elem_kv_hash = middle_merk_key1 - .get_kv_hash(b"key2", true) + .get_kv_hash( + b"key2", + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get kv hash") .expect("value hash should be some"); @@ -220,7 +254,11 @@ fn test_tree_hashes_when_inserting_empty_trees_twice_under_each_other() { assert_eq!(elem_kv_hash, kv_hash_key2); let elem_node_hash = middle_merk_key1 - .get_hash(b"key2", true) + .get_hash( + b"key2", + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get kv hash") .expect("value hash should be some"); @@ -238,7 +276,11 @@ fn test_tree_hashes_when_inserting_empty_trees_twice_under_each_other() { assert_eq!(root_hash, node_hash_key2); let (middle_elem_value_key1, middle_elem_value_hash_key1) = under_top_merk - .get_value_and_value_hash(b"key1", true) + .get_value_and_value_hash( + b"key1", + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get value hash") .expect("value hash should be some"); @@ -276,7 +318,11 @@ fn test_tree_hashes_when_inserting_empty_trees_twice_under_each_other() { ); let middle_elem_kv_hash_key1 = under_top_merk - .get_kv_hash(b"key1", true) + .get_kv_hash( + b"key1", + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get value hash") .expect("value hash should be some"); @@ -290,7 +336,11 @@ fn test_tree_hashes_when_inserting_empty_trees_twice_under_each_other() { ); let middle_elem_node_hash_key1 = under_top_merk - .get_hash(b"key1", true) + .get_hash( + b"key1", + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get value hash") .expect("value hash should be some"); diff --git a/grovedb/src/util.rs b/grovedb/src/util.rs index da478128..cbedc9af 100644 --- a/grovedb/src/util.rs +++ b/grovedb/src/util.rs @@ -198,8 +198,11 @@ macro_rules! merk_optional_tx { { let $subtree = cost_return_on_error!( &mut $cost, - ::grovedb_merk::Merk::open_base(storage.unwrap_add_cost(&mut $cost), false) - .map(|merk_res| + ::grovedb_merk::Merk::open_base( + storage.unwrap_add_cost(&mut $cost), + false, + Some(&Element::value_defined_cost_for_serialized_value) + ).map(|merk_res| merk_res .map_err(|_| crate::Error::CorruptedData( "cannot open a subtree".to_owned() @@ -226,7 +229,8 @@ macro_rules! merk_optional_tx { ::grovedb_merk::Merk::open_layered_with_root_key( storage, root_key, - is_sum_tree + is_sum_tree, + Some(&Element::value_defined_cost_for_serialized_value), ).map(|merk_res| merk_res .map_err(|_| crate::Error::CorruptedData( @@ -271,7 +275,8 @@ macro_rules! merk_optional_tx_path_not_empty { ::grovedb_merk::Merk::open_layered_with_root_key( storage, root_key, - is_sum_tree + is_sum_tree, + Some(&Element::value_defined_cost_for_serialized_value), ).map(|merk_res| merk_res .map_err(|_| crate::Error::CorruptedData( @@ -308,8 +313,11 @@ macro_rules! root_merk_optional_tx { { let $subtree = cost_return_on_error!( &mut $cost, - ::grovedb_merk::Merk::open_base(storage.unwrap_add_cost(&mut $cost), false) - .map(|merk_res| + ::grovedb_merk::Merk::open_base( + storage.unwrap_add_cost(&mut $cost), + false, + Some(&Element::value_defined_cost_for_serialized_value) + ).map(|merk_res| merk_res .map_err(|_| crate::Error::CorruptedData( "cannot open a subtree".to_owned() diff --git a/merk/src/merk/apply.rs b/merk/src/merk/apply.rs index 52b966a2..4418dfc1 100644 --- a/merk/src/merk/apply.rs +++ b/merk/src/merk/apply.rs @@ -31,15 +31,15 @@ where /// # Example /// ``` /// # let mut store = grovedb_merk::test_utils::TempMerk::new(); - /// # store.apply::<_, Vec<_>>(&[(vec![4,5,6], Op::Put(vec![0], BasicMerk))], &[], None) + /// # store.apply::<_, Vec<_>>(&[(vec![4,5,6], Op::Put(vec![0], BasicMerkNode))], &[], None) /// .unwrap().expect(""); /// /// use grovedb_merk::Op; - /// use grovedb_merk::TreeFeatureType::BasicMerk; + /// use grovedb_merk::TreeFeatureType::BasicMerkNode; /// /// let batch = &[ /// // puts value [4,5,6] to key[1,2,3] - /// (vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerk)), + /// (vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerkNode)), /// // deletes key [4,5,6] /// (vec![4, 5, 6], Op::Delete), /// ]; @@ -67,6 +67,7 @@ where use_sum_nodes, )) }, + None::<&fn(&[u8]) -> Option>, &mut |_costs, _old_value, _value| Ok((false, None)), &mut |_a, key_bytes_to_remove, value_bytes_to_remove| { Ok(( @@ -87,15 +88,15 @@ where /// # Example /// ``` /// # let mut store = grovedb_merk::test_utils::TempMerk::new(); - /// # store.apply::<_, Vec<_>>(&[(vec![4,5,6], Op::Put(vec![0], BasicMerk))], &[], None) + /// # store.apply::<_, Vec<_>>(&[(vec![4,5,6], Op::Put(vec![0], BasicMerkNode))], &[], None) /// .unwrap().expect(""); /// /// use grovedb_merk::Op; - /// use grovedb_merk::TreeFeatureType::BasicMerk; + /// use grovedb_merk::TreeFeatureType::BasicMerkNode; /// /// let batch = &[ /// // puts value [4,5,6] to key[1,2,3] - /// (vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerk)), + /// (vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerkNode)), /// // deletes key [4,5,6] /// (vec![4, 5, 6], Op::Delete), /// ]; @@ -107,6 +108,7 @@ where aux: &AuxMerkBatch, options: Option, old_specialized_cost: &impl Fn(&Vec, &Vec) -> Result, + value_defined_cost_fn: Option<&impl Fn(&[u8]) -> Option>, ) -> CostResult<(), Error> where KB: AsRef<[u8]>, @@ -117,6 +119,7 @@ where aux, options, old_specialized_cost, + value_defined_cost_fn, &mut |_costs, _old_value, _value| Ok((false, None)), &mut |_a, key_bytes_to_remove, value_bytes_to_remove| { Ok(( @@ -139,7 +142,7 @@ where /// ``` /// # let mut store = grovedb_merk::test_utils::TempMerk::new(); /// # store.apply_with_costs_just_in_time_value_update::<_, Vec<_>>( - /// &[(vec![4,5,6], Op::Put(vec![0], BasicMerk))], + /// &[(vec![4,5,6], Op::Put(vec![0], BasicMerkNode))], /// &[], /// None, /// &|k, v| Ok(0), @@ -149,11 +152,11 @@ where /// /// use grovedb_costs::storage_cost::removal::StorageRemovedBytes::NoStorageRemoval; /// use grovedb_merk::Op; - /// use grovedb_merk::TreeFeatureType::BasicMerk; + /// use grovedb_merk::TreeFeatureType::BasicMerkNode; /// /// let batch = &[ /// // puts value [4,5,6] to key[1,2,3] - /// (vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerk)), + /// (vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerkNode)), /// // deletes key [4,5,6] /// (vec![4, 5, 6], Op::Delete), /// ]; @@ -173,6 +176,7 @@ where aux: &AuxMerkBatch, options: Option, old_specialized_cost: &impl Fn(&Vec, &Vec) -> Result, + value_defined_cost_fn: Option<&impl Fn(&[u8]) -> Option>, update_tree_value_based_on_costs: &mut impl FnMut( &StorageCost, &Vec, @@ -218,6 +222,7 @@ where aux, options, old_specialized_cost, + value_defined_cost_fn, update_tree_value_based_on_costs, section_removal_bytes, ) @@ -235,7 +240,7 @@ where /// ``` /// # let mut store = grovedb_merk::test_utils::TempMerk::new(); /// # store.apply_with_costs_just_in_time_value_update::<_, Vec<_>>( - /// &[(vec![4,5,6], Op::Put(vec![0], BasicMerk))], + /// &[(vec![4,5,6], Op::Put(vec![0], BasicMerkNode))], /// &[], /// None, /// &|k, v| Ok(0), @@ -245,11 +250,11 @@ where /// /// use grovedb_costs::storage_cost::removal::StorageRemovedBytes::NoStorageRemoval; /// use grovedb_merk::Op; - /// use grovedb_merk::TreeFeatureType::BasicMerk; + /// use grovedb_merk::TreeFeatureType::BasicMerkNode; /// /// let batch = &[ /// // puts value [4,5,6] to key [1,2,3] - /// (vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerk)), + /// (vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerkNode)), /// // deletes key [4,5,6] /// (vec![4, 5, 6], Op::Delete), /// ]; @@ -258,17 +263,19 @@ where /// &[], /// None, /// &|k, v| Ok(0), + /// &|v| None, /// &mut |s, o, v| Ok((false, None)), /// &mut |s, k, v| Ok((NoStorageRemoval, NoStorageRemoval)) /// ).unwrap().expect(""); /// } /// ``` - pub fn apply_unchecked( + pub fn apply_unchecked( &mut self, batch: &MerkBatch, aux: &AuxMerkBatch, options: Option, old_specialized_cost: &C, + value_defined_cost_fn: Option<&V>, update_tree_value_based_on_costs: &mut U, section_removal_bytes: &mut R, ) -> CostResult<(), Error> @@ -276,6 +283,7 @@ where KB: AsRef<[u8]>, KA: AsRef<[u8]>, C: Fn(&Vec, &Vec) -> Result, + V: Fn(&[u8]) -> Option, U: FnMut( &StorageCost, &Vec, @@ -294,6 +302,7 @@ where batch, self.source(), old_specialized_cost, + value_defined_cost_fn, update_tree_value_based_on_costs, section_removal_bytes, ) diff --git a/merk/src/merk/chunks.rs b/merk/src/merk/chunks.rs index 7e8c588e..4f6564ef 100644 --- a/merk/src/merk/chunks.rs +++ b/merk/src/merk/chunks.rs @@ -214,6 +214,7 @@ mod tests { use crate::{ proofs::chunk::{verify_leaf, verify_trunk}, test_utils::*, + tree::kv::ValueDefinedCostType, }; #[test] @@ -275,6 +276,7 @@ mod tests { .get_storage_context(SubtreePath::empty(), Some(&batch)) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .unwrap(); @@ -293,6 +295,7 @@ mod tests { .get_storage_context(SubtreePath::empty(), None) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .unwrap(); @@ -311,6 +314,7 @@ mod tests { .get_storage_context(SubtreePath::empty(), None) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .unwrap(); diff --git a/merk/src/merk/get.rs b/merk/src/merk/get.rs index 573b895a..f0c25b42 100644 --- a/merk/src/merk/get.rs +++ b/merk/src/merk/get.rs @@ -1,7 +1,12 @@ use grovedb_costs::{CostContext, CostResult, CostsExt, OperationCost}; use grovedb_storage::StorageContext; -use crate::{tree::TreeNode, CryptoHash, Error, Error::StorageError, Merk, TreeFeatureType}; +use crate::{ + tree::{kv::ValueDefinedCostType, TreeNode}, + CryptoHash, Error, + Error::StorageError, + Merk, TreeFeatureType, +}; impl<'db, S> Merk where @@ -16,8 +21,12 @@ where /// /// Note that this is essentially the same as a normal RocksDB `get`, so /// should be a fast operation and has almost no tree overhead. - pub fn exists(&self, key: &[u8]) -> CostResult { - self.has_node_direct(key) + pub fn exists( + &self, + key: &[u8], + value_defined_cost_fn: Option Option>, + ) -> CostResult { + self.has_node_direct(key, value_defined_cost_fn) } /// Returns if the value at the given key exists @@ -26,8 +35,12 @@ where /// should be a fast operation and has almost no tree overhead. /// Contrary to a simple exists, this traverses the tree and can be faster /// if the tree is cached, but slower if it is not - pub fn exists_by_traversing_tree(&self, key: &[u8]) -> CostResult { - self.has_node(key) + pub fn exists_by_traversing_tree( + &self, + key: &[u8], + value_defined_cost_fn: Option Option>, + ) -> CostResult { + self.has_node(key, value_defined_cost_fn) } /// Gets a value for the given key. If the key is not found, `None` is @@ -35,19 +48,32 @@ where /// /// Note that this is essentially the same as a normal RocksDB `get`, so /// should be a fast operation and has almost no tree overhead. - pub fn get(&self, key: &[u8], allow_cache: bool) -> CostResult>, Error> { + pub fn get( + &self, + key: &[u8], + allow_cache: bool, + value_defined_cost_fn: Option Option>, + ) -> CostResult>, Error> { if allow_cache { - self.get_node_fn(key, |node| { - node.value_as_slice() - .to_vec() - .wrap_with_cost(Default::default()) - }) + self.get_node_fn( + key, + |node| { + node.value_as_slice() + .to_vec() + .wrap_with_cost(Default::default()) + }, + value_defined_cost_fn, + ) } else { - self.get_node_direct_fn(key, |node| { - node.value_as_slice() - .to_vec() - .wrap_with_cost(Default::default()) - }) + self.get_node_direct_fn( + key, + |node| { + node.value_as_slice() + .to_vec() + .wrap_with_cost(Default::default()) + }, + value_defined_cost_fn, + ) } } @@ -56,25 +82,35 @@ where &self, key: &[u8], allow_cache: bool, + value_defined_cost_fn: Option Option>, ) -> CostResult, Error> { if allow_cache { - self.get_node_fn(key, |node| { - node.feature_type().wrap_with_cost(Default::default()) - }) + self.get_node_fn( + key, + |node| node.feature_type().wrap_with_cost(Default::default()), + value_defined_cost_fn, + ) } else { - self.get_node_direct_fn(key, |node| { - node.feature_type().wrap_with_cost(Default::default()) - }) + self.get_node_direct_fn( + key, + |node| node.feature_type().wrap_with_cost(Default::default()), + value_defined_cost_fn, + ) } } /// Gets a hash of a node by a given key, `None` is returned in case /// when node not found by the key. - pub fn get_hash(&self, key: &[u8], allow_cache: bool) -> CostResult, Error> { + pub fn get_hash( + &self, + key: &[u8], + allow_cache: bool, + value_defined_cost_fn: Option Option>, + ) -> CostResult, Error> { if allow_cache { - self.get_node_fn(key, |node| node.hash()) + self.get_node_fn(key, |node| node.hash(), value_defined_cost_fn) } else { - self.get_node_direct_fn(key, |node| node.hash()) + self.get_node_direct_fn(key, |node| node.hash(), value_defined_cost_fn) } } @@ -84,15 +120,20 @@ where &self, key: &[u8], allow_cache: bool, + value_defined_cost_fn: Option Option>, ) -> CostResult, Error> { if allow_cache { - self.get_node_fn(key, |node| { - (*node.value_hash()).wrap_with_cost(OperationCost::default()) - }) + self.get_node_fn( + key, + |node| (*node.value_hash()).wrap_with_cost(OperationCost::default()), + value_defined_cost_fn, + ) } else { - self.get_node_direct_fn(key, |node| { - (*node.value_hash()).wrap_with_cost(OperationCost::default()) - }) + self.get_node_direct_fn( + key, + |node| (*node.value_hash()).wrap_with_cost(OperationCost::default()), + value_defined_cost_fn, + ) } } @@ -102,15 +143,20 @@ where &self, key: &[u8], allow_cache: bool, + value_defined_cost_fn: Option Option>, ) -> CostResult, Error> { if allow_cache { - self.get_node_fn(key, |node| { - (*node.inner.kv.hash()).wrap_with_cost(OperationCost::default()) - }) + self.get_node_fn( + key, + |node| (*node.inner.kv.hash()).wrap_with_cost(OperationCost::default()), + value_defined_cost_fn, + ) } else { - self.get_node_direct_fn(key, |node| { - (*node.inner.kv.hash()).wrap_with_cost(OperationCost::default()) - }) + self.get_node_direct_fn( + key, + |node| (*node.inner.kv.hash()).wrap_with_cost(OperationCost::default()), + value_defined_cost_fn, + ) } } @@ -120,27 +166,44 @@ where &self, key: &[u8], allow_cache: bool, + value_defined_cost_fn: Option Option>, ) -> CostResult, CryptoHash)>, Error> { if allow_cache { - self.get_node_fn(key, |node| { - (node.value_as_slice().to_vec(), *node.value_hash()) - .wrap_with_cost(OperationCost::default()) - }) + self.get_node_fn( + key, + |node| { + (node.value_as_slice().to_vec(), *node.value_hash()) + .wrap_with_cost(OperationCost::default()) + }, + value_defined_cost_fn, + ) } else { - self.get_node_direct_fn(key, |node| { - (node.value_as_slice().to_vec(), *node.value_hash()) - .wrap_with_cost(OperationCost::default()) - }) + self.get_node_direct_fn( + key, + |node| { + (node.value_as_slice().to_vec(), *node.value_hash()) + .wrap_with_cost(OperationCost::default()) + }, + value_defined_cost_fn, + ) } } /// See if a node's field exists - fn has_node_direct(&self, key: &[u8]) -> CostResult { - TreeNode::get(&self.storage, key).map_ok(|x| x.is_some()) + fn has_node_direct( + &self, + key: &[u8], + value_defined_cost_fn: Option Option>, + ) -> CostResult { + TreeNode::get(&self.storage, key, value_defined_cost_fn).map_ok(|x| x.is_some()) } /// See if a node's field exists - fn has_node(&self, key: &[u8]) -> CostResult { + fn has_node( + &self, + key: &[u8], + value_defined_cost_fn: Option Option>, + ) -> CostResult { self.use_tree(move |maybe_tree| { let mut cursor = match maybe_tree { None => return Ok(false).wrap_with_cost(Default::default()), // empty tree @@ -162,7 +225,7 @@ where match maybe_child { None => { // fetch from RocksDB - break self.has_node_direct(key); + break self.has_node_direct(key, value_defined_cost_fn); } Some(child) => cursor = child, // traverse to child } @@ -171,18 +234,28 @@ where } /// Generic way to get a node's field - fn get_node_direct_fn(&self, key: &[u8], f: F) -> CostResult, Error> + fn get_node_direct_fn( + &self, + key: &[u8], + f: F, + value_defined_cost_fn: Option Option>, + ) -> CostResult, Error> where F: FnOnce(&TreeNode) -> CostContext, { - TreeNode::get(&self.storage, key).flat_map_ok(|maybe_node| { + TreeNode::get(&self.storage, key, value_defined_cost_fn).flat_map_ok(|maybe_node| { let mut cost = OperationCost::default(); Ok(maybe_node.map(|node| f(&node).unwrap_add_cost(&mut cost))).wrap_with_cost(cost) }) } /// Generic way to get a node's field - fn get_node_fn(&self, key: &[u8], f: F) -> CostResult, Error> + fn get_node_fn( + &self, + key: &[u8], + f: F, + value_defined_cost_fn: Option Option>, + ) -> CostResult, Error> where F: FnOnce(&TreeNode) -> CostContext, { @@ -207,7 +280,7 @@ where match maybe_child { None => { // fetch from RocksDB - break self.get_node_direct_fn(key, f); + break self.get_node_direct_fn(key, f, value_defined_cost_fn); } Some(child) => cursor = child, // traverse to child } @@ -218,7 +291,9 @@ where #[cfg(test)] mod test { - use crate::{test_utils::TempMerk, Op, TreeFeatureType::BasicMerk}; + use crate::{ + test_utils::TempMerk, tree::kv::ValueDefinedCostType, Op, TreeFeatureType::BasicMerkNode, + }; #[test] fn test_has_node_with_empty_tree() { @@ -226,11 +301,14 @@ mod test { let key = b"something"; - let result = merk.has_node(key).unwrap().unwrap(); + let result = merk + .has_node(key, None::<&fn(&[u8]) -> Option>) + .unwrap() + .unwrap(); assert!(!result); - let batch_entry = (key, Op::Put(vec![123; 60], BasicMerk)); + let batch_entry = (key, Op::Put(vec![123; 60], BasicMerkNode)); let batch = vec![batch_entry]; @@ -238,7 +316,10 @@ mod test { .unwrap() .expect("should ..."); - let result = merk.has_node(key).unwrap().unwrap(); + let result = merk + .has_node(key, None::<&fn(&[u8]) -> Option>) + .unwrap() + .unwrap(); assert!(result); } diff --git a/merk/src/merk/mod.rs b/merk/src/merk/mod.rs index 89982ef0..9c025dcb 100644 --- a/merk/src/merk/mod.rs +++ b/merk/src/merk/mod.rs @@ -494,7 +494,10 @@ where /// Loads the Merk from the base root key /// The base root key should only be used if the Merk tree is independent /// Meaning that it doesn't have a parent Merk - pub(crate) fn load_base_root(&mut self) -> CostResult<(), Error> { + pub(crate) fn load_base_root( + &mut self, + value_defined_cost_fn: Option Option>, + ) -> CostResult<(), Error> { self.storage .get_root(ROOT_KEY_KEY) .map(|root_result| root_result.map_err(Error::StorageError)) @@ -503,12 +506,14 @@ where if let Some(tree_root_key) = tree_root_key_opt { // Trying to build a tree out of it, costs will be accumulated because // `Tree::get` returns `CostContext` and this call happens inside `flat_map_ok`. - TreeNode::get(&self.storage, tree_root_key).map_ok(|tree| { - if let Some(t) = tree.as_ref() { - self.root_tree_key = Cell::new(Some(t.key().to_vec())); - } - self.tree = Cell::new(tree); - }) + TreeNode::get(&self.storage, tree_root_key, value_defined_cost_fn).map_ok( + |tree| { + if let Some(t) = tree.as_ref() { + self.root_tree_key = Cell::new(Some(t.key().to_vec())); + } + self.tree = Cell::new(tree); + }, + ) } else { Ok(()).wrap_with_cost(Default::default()) } @@ -518,12 +523,15 @@ where /// Loads the Merk from it's parent root key /// The base root key should only be used if the Merk tree is independent /// Meaning that it doesn't have a parent Merk - pub(crate) fn load_root(&mut self) -> CostResult<(), Error> { + pub(crate) fn load_root( + &mut self, + value_defined_cost_fn: Option Option>, + ) -> CostResult<(), Error> { // In case of successful seek for root key check if it exists if let Some(tree_root_key) = self.root_tree_key.get_mut() { // Trying to build a tree out of it, costs will be accumulated because // `Tree::get` returns `CostContext` and this call happens inside `flat_map_ok`. - TreeNode::get(&self.storage, tree_root_key).map_ok(|tree| { + TreeNode::get(&self.storage, tree_root_key, value_defined_cost_fn).map_ok(|tree| { self.tree = Cell::new(tree); }) } else { @@ -533,11 +541,15 @@ where } } -fn fetch_node<'db>(db: &impl StorageContext<'db>, key: &[u8]) -> Result, Error> { +fn fetch_node<'db>( + db: &impl StorageContext<'db>, + key: &[u8], + value_defined_cost_fn: Option Option>, +) -> Result, Error> { let bytes = db.get(key).unwrap().map_err(StorageError)?; // TODO: get_pinned ? if let Some(bytes) = bytes { Ok(Some( - TreeNode::decode(key.to_vec(), &bytes).map_err(EdError)?, + TreeNode::decode(key.to_vec(), &bytes, value_defined_cost_fn).map_err(EdError)?, )) } else { Ok(None) @@ -556,7 +568,10 @@ mod test { use tempfile::TempDir; use super::{Merk, RefWalker}; - use crate::{merk::source::MerkSource, test_utils::*, Op, TreeFeatureType::BasicMerk}; + use crate::{ + merk::source::MerkSource, test_utils::*, tree::kv::ValueDefinedCostType, Op, + TreeFeatureType::BasicMerkNode, + }; // TODO: Close and then reopen test @@ -656,7 +671,7 @@ mod test { let mut merk = TempMerk::new(); merk.apply::, _>( &[], - &[(vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerk), None)], + &[(vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerkNode), None)], None, ) .unwrap() @@ -672,27 +687,55 @@ mod test { let mut merk = TempMerk::new(); // no root - assert!(merk.get(&[1, 2, 3], true).unwrap().unwrap().is_none()); + assert!(merk + .get( + &[1, 2, 3], + true, + None::<&fn(&[u8]) -> Option> + ) + .unwrap() + .unwrap() + .is_none()); // cached - merk.apply::<_, Vec<_>>(&[(vec![5, 5, 5], Op::Put(vec![], BasicMerk))], &[], None) + merk.apply::<_, Vec<_>>( + &[(vec![5, 5, 5], Op::Put(vec![], BasicMerkNode))], + &[], + None, + ) + .unwrap() + .unwrap(); + assert!(merk + .get( + &[1, 2, 3], + true, + None::<&fn(&[u8]) -> Option> + ) .unwrap() - .unwrap(); - assert!(merk.get(&[1, 2, 3], true).unwrap().unwrap().is_none()); + .unwrap() + .is_none()); // uncached merk.apply::<_, Vec<_>>( &[ - (vec![0, 0, 0], Op::Put(vec![], BasicMerk)), - (vec![1, 1, 1], Op::Put(vec![], BasicMerk)), - (vec![2, 2, 2], Op::Put(vec![], BasicMerk)), + (vec![0, 0, 0], Op::Put(vec![], BasicMerkNode)), + (vec![1, 1, 1], Op::Put(vec![], BasicMerkNode)), + (vec![2, 2, 2], Op::Put(vec![], BasicMerkNode)), ], &[], None, ) .unwrap() .unwrap(); - assert!(merk.get(&[3, 3, 3], true).unwrap().unwrap().is_none()); + assert!(merk + .get( + &[3, 3, 3], + true, + None::<&fn(&[u8]) -> Option> + ) + .unwrap() + .unwrap() + .is_none()); } // TODO: what this test should do? @@ -706,6 +749,7 @@ mod test { .get_storage_context(SubtreePath::empty(), None) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .expect("cannot open merk"); @@ -729,6 +773,7 @@ mod test { .get_storage_context(SubtreePath::empty(), None) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .expect("cannot open merk"); @@ -746,10 +791,18 @@ mod test { nodes: &mut Vec>, ) { nodes.push(node.tree().encode()); - if let Some(c) = node.walk(true).unwrap().unwrap() { + if let Some(c) = node + .walk(true, None::<&fn(&[u8]) -> Option>) + .unwrap() + .unwrap() + { collect(c, nodes); } - if let Some(c) = node.walk(false).unwrap().unwrap() { + if let Some(c) = node + .walk(false, None::<&fn(&[u8]) -> Option>) + .unwrap() + .unwrap() + { collect(c, nodes); } } @@ -765,6 +818,7 @@ mod test { .get_storage_context(SubtreePath::empty(), Some(&batch)) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .expect("cannot open merk"); @@ -782,6 +836,7 @@ mod test { .get_storage_context(SubtreePath::empty(), None) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .expect("cannot open merk"); @@ -801,6 +856,7 @@ mod test { .get_storage_context(SubtreePath::empty(), None) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .expect("cannot open merk"); @@ -841,6 +897,7 @@ mod test { .get_storage_context(SubtreePath::empty(), Some(&batch)) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .expect("cannot open merk"); @@ -860,6 +917,7 @@ mod test { .get_storage_context(SubtreePath::empty(), None) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .expect("cannot open merk"); @@ -873,6 +931,7 @@ mod test { .get_storage_context(SubtreePath::empty(), None) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .expect("cannot open merk"); @@ -894,19 +953,20 @@ mod test { .get_storage_context(SubtreePath::empty(), Some(&batch)) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .expect("cannot open merk"); merk.apply::<_, Vec<_>>( - &[(b"9".to_vec(), Op::Put(b"a".to_vec(), BasicMerk))], + &[(b"9".to_vec(), Op::Put(b"a".to_vec(), BasicMerkNode))], &[], None, ) .unwrap() .expect("should insert successfully"); merk.apply::<_, Vec<_>>( - &[(b"10".to_vec(), Op::Put(b"a".to_vec(), BasicMerk))], + &[(b"10".to_vec(), Op::Put(b"a".to_vec(), BasicMerkNode))], &[], None, ) @@ -914,21 +974,29 @@ mod test { .expect("should insert successfully"); let result = merk - .get(b"10".as_slice(), true) + .get( + b"10".as_slice(), + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get successfully"); assert_eq!(result, Some(b"a".to_vec())); // Update the node merk.apply::<_, Vec<_>>( - &[(b"10".to_vec(), Op::Put(b"b".to_vec(), BasicMerk))], + &[(b"10".to_vec(), Op::Put(b"b".to_vec(), BasicMerkNode))], &[], None, ) .unwrap() .expect("should insert successfully"); let result = merk - .get(b"10".as_slice(), true) + .get( + b"10".as_slice(), + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get successfully"); assert_eq!(result, Some(b"b".to_vec())); @@ -943,20 +1011,25 @@ mod test { .get_storage_context(SubtreePath::empty(), None) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .expect("cannot open merk"); // Update the node after dropping merk merk.apply::<_, Vec<_>>( - &[(b"10".to_vec(), Op::Put(b"c".to_vec(), BasicMerk))], + &[(b"10".to_vec(), Op::Put(b"c".to_vec(), BasicMerkNode))], &[], None, ) .unwrap() .expect("should insert successfully"); let result = merk - .get(b"10".as_slice(), true) + .get( + b"10".as_slice(), + true, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("should get successfully"); assert_eq!(result, Some(b"c".to_vec())); diff --git a/merk/src/merk/open.rs b/merk/src/merk/open.rs index 115926ab..af15d596 100644 --- a/merk/src/merk/open.rs +++ b/merk/src/merk/open.rs @@ -4,6 +4,7 @@ use grovedb_costs::CostResult; use grovedb_storage::StorageContext; use crate::{ + tree::kv::ValueDefinedCostType, Error, Merk, MerkType, MerkType::{BaseMerk, LayeredMerk, StandaloneMerk}, }; @@ -24,7 +25,11 @@ where } /// Open standalone tree - pub fn open_standalone(storage: S, is_sum_tree: bool) -> CostResult { + pub fn open_standalone( + storage: S, + is_sum_tree: bool, + value_defined_cost_fn: Option Option>, + ) -> CostResult { let mut merk = Self { tree: Cell::new(None), root_tree_key: Cell::new(None), @@ -33,11 +38,15 @@ where is_sum_tree, }; - merk.load_base_root().map_ok(|_| merk) + merk.load_base_root(value_defined_cost_fn).map_ok(|_| merk) } /// Open base tree - pub fn open_base(storage: S, is_sum_tree: bool) -> CostResult { + pub fn open_base( + storage: S, + is_sum_tree: bool, + value_defined_cost_fn: Option Option>, + ) -> CostResult { let mut merk = Self { tree: Cell::new(None), root_tree_key: Cell::new(None), @@ -46,7 +55,7 @@ where is_sum_tree, }; - merk.load_base_root().map_ok(|_| merk) + merk.load_base_root(value_defined_cost_fn).map_ok(|_| merk) } /// Open layered tree with root key @@ -54,6 +63,7 @@ where storage: S, root_key: Option>, is_sum_tree: bool, + value_defined_cost_fn: Option Option>, ) -> CostResult { let mut merk = Self { tree: Cell::new(None), @@ -63,7 +73,7 @@ where is_sum_tree, }; - merk.load_root().map_ok(|_| merk) + merk.load_root(value_defined_cost_fn).map_ok(|_| merk) } } @@ -77,7 +87,7 @@ mod test { }; use tempfile::TempDir; - use crate::{Merk, Op, TreeFeatureType::BasicMerk}; + use crate::{tree::kv::ValueDefinedCostType, Merk, Op, TreeFeatureType::BasicMerkNode}; #[test] fn test_reopen_root_hash() { @@ -92,12 +102,13 @@ mod test { .get_storage_context(SubtreePath::from(test_prefix.as_ref()), Some(&batch)) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .unwrap(); merk.apply::<_, Vec<_>>( - &[(vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerk))], + &[(vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerkNode))], &[], None, ) @@ -116,6 +127,7 @@ mod test { .get_storage_context(SubtreePath::from(test_prefix.as_ref()), None) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .unwrap(); @@ -132,6 +144,7 @@ mod test { .get_storage_context(SubtreePath::empty(), Some(&batch)) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ); // Opening not existing merk should cost only root key seek (except context // creation) @@ -142,7 +155,7 @@ mod test { let mut merk = merk_fee_context.unwrap().unwrap(); merk.apply::<_, Vec<_>>( - &[(vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerk))], + &[(vec![1, 2, 3], Op::Put(vec![4, 5, 6], BasicMerkNode))], &[], None, ) @@ -159,6 +172,7 @@ mod test { .get_storage_context(SubtreePath::empty(), None) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ); // Opening existing merk should cost two seeks. (except context creation) diff --git a/merk/src/merk/restore.rs b/merk/src/merk/restore.rs index eb4cea61..e6ac22e2 100644 --- a/merk/src/merk/restore.rs +++ b/merk/src/merk/restore.rs @@ -39,6 +39,7 @@ use grovedb_storage::{Batch, StorageContext}; use super::Merk; #[cfg(feature = "full")] use crate::merk::source::MerkSource; +use crate::tree::kv::ValueDefinedCostType; #[cfg(feature = "full")] use crate::{ error::Error, @@ -50,7 +51,7 @@ use crate::{ tree::{combine_hash, value_hash, Link, RefWalker, TreeNode}, CryptoHash, Error::{CostsError, EdError, StorageError}, - TreeFeatureType::BasicMerk, + TreeFeatureType::BasicMerkNode, }; #[cfg(feature = "full")] @@ -116,7 +117,9 @@ impl<'db, S: StorageContext<'db>> Restorer { self.rewrite_trunk_child_heights()?; } - self.merk.load_base_root().unwrap()?; + self.merk + .load_base_root(None:: Option>) + .unwrap()?; Ok(self.merk) } @@ -136,7 +139,7 @@ impl<'db, S: StorageContext<'db>> Restorer { tree.visit_refs(&mut |proof_node| { if let Some((mut node, key)) = match &proof_node.node { Node::KV(key, value) => Some(( - TreeNode::new(key.clone(), value.clone(), None, BasicMerk).unwrap(), + TreeNode::new(key.clone(), value.clone(), None, BasicMerkNode).unwrap(), key, )), Node::KVValueHash(key, value, value_hash) => Some(( @@ -144,7 +147,7 @@ impl<'db, S: StorageContext<'db>> Restorer { key.clone(), value.clone(), *value_hash, - BasicMerk, + BasicMerkNode, ) .unwrap(), key, @@ -270,8 +273,12 @@ impl<'db, S: StorageContext<'db>> Restorer { fn rewrite_parent_link(&mut self, leaf: &ProofTree) -> Result<(), Error> { let parent_keys = self.parent_keys.as_mut().unwrap(); let parent_key = parent_keys.peek().unwrap().clone(); - let mut parent = crate::merk::fetch_node(&self.merk.storage, parent_key.as_slice())? - .expect("Could not find parent of leaf chunk"); + let mut parent = crate::merk::fetch_node( + &self.merk.storage, + parent_key.as_slice(), + None:: Option>, + )? + .expect("Could not find parent of leaf chunk"); let is_left_child = self.remaining_chunks_unchecked() % 2 == 0; if let Some(Link::Reference { ref mut key, .. }) = parent.link_mut(is_left_child) { @@ -305,16 +312,25 @@ impl<'db, S: StorageContext<'db>> Restorer { return Ok(node.tree().child_heights()); } - let mut cloned_node = - TreeNode::decode(node.tree().key().to_vec(), node.tree().encode().as_slice()) - .map_err(EdError)?; + let mut cloned_node = TreeNode::decode( + node.tree().key().to_vec(), + node.tree().encode().as_slice(), + None:: Option>, + ) + .map_err(EdError)?; - let left_child = node.walk(true).unwrap()?.unwrap(); + let left_child = node + .walk(true, None::<&fn(&[u8]) -> Option>) + .unwrap()? + .unwrap(); let left_child_heights = recurse(left_child, remaining_depth - 1, batch)?; let left_height = left_child_heights.0.max(left_child_heights.1) + 1; *cloned_node.link_mut(true).unwrap().child_heights_mut() = left_child_heights; - let right_child = node.walk(false).unwrap()?.unwrap(); + let right_child = node + .walk(false, None::<&fn(&[u8]) -> Option>) + .unwrap()? + .unwrap(); let right_child_heights = recurse(right_child, remaining_depth - 1, batch)?; let right_height = right_child_heights.0.max(right_child_heights.1) + 1; *cloned_node.link_mut(false).unwrap().child_heights_mut() = right_child_heights; @@ -327,7 +343,9 @@ impl<'db, S: StorageContext<'db>> Restorer { Ok((left_height, right_height)) } - self.merk.load_base_root().unwrap()?; + self.merk + .load_base_root(None:: Option>) + .unwrap()?; let mut batch = self.merk.storage.new_batch(); @@ -415,6 +433,7 @@ mod tests { .get_immediate_storage_context(SubtreePath::empty(), &tx) .unwrap(), false, + None::<&fn(&[u8]) -> Option>, ) .unwrap() .unwrap(); @@ -432,7 +451,13 @@ mod tests { let ctx = storage .get_immediate_storage_context(SubtreePath::empty(), &tx) .unwrap(); - let merk = Merk::open_base(ctx, false).unwrap().unwrap(); + let merk = Merk::open_base( + ctx, + false, + None::<&fn(&[u8]) -> Option>, + ) + .unwrap() + .unwrap(); let mut restorer = Merk::restore(merk, original.root_hash().unwrap()); assert_eq!(restorer.remaining_chunks(), None); @@ -466,8 +491,8 @@ mod tests { fn restore_2_left_heavy() { restore_test( &[ - &[(vec![0], Op::Put(vec![], BasicMerk))], - &[(vec![1], Op::Put(vec![], BasicMerk))], + &[(vec![0], Op::Put(vec![], BasicMerkNode))], + &[(vec![1], Op::Put(vec![], BasicMerkNode))], ], 2, ); @@ -477,8 +502,8 @@ mod tests { fn restore_2_right_heavy() { restore_test( &[ - &[(vec![1], Op::Put(vec![], BasicMerk))], - &[(vec![0], Op::Put(vec![], BasicMerk))], + &[(vec![1], Op::Put(vec![], BasicMerkNode))], + &[(vec![0], Op::Put(vec![], BasicMerkNode))], ], 2, ); diff --git a/merk/src/merk/source.rs b/merk/src/merk/source.rs index 709063a6..46782bdc 100644 --- a/merk/src/merk/source.rs +++ b/merk/src/merk/source.rs @@ -2,7 +2,7 @@ use grovedb_costs::CostResult; use grovedb_storage::StorageContext; use crate::{ - tree::{Fetch, TreeNode}, + tree::{kv::ValueDefinedCostType, Fetch, TreeNode}, Error, Link, Merk, }; @@ -37,8 +37,12 @@ impl<'s, 'db, S> Fetch for MerkSource<'s, S> where S: StorageContext<'db>, { - fn fetch(&self, link: &Link) -> CostResult { - TreeNode::get(self.storage, link.key()) + fn fetch( + &self, + link: &Link, + value_defined_cost_fn: Option<&impl Fn(&[u8]) -> Option>, + ) -> CostResult { + TreeNode::get(self.storage, link.key(), value_defined_cost_fn) .map_ok(|x| x.ok_or(Error::KeyNotFoundError("Key not found for fetch"))) .flatten() } diff --git a/merk/src/proofs/chunk.rs b/merk/src/proofs/chunk.rs index 077683cd..5e852ec1 100644 --- a/merk/src/proofs/chunk.rs +++ b/merk/src/proofs/chunk.rs @@ -43,12 +43,13 @@ use { #[cfg(feature = "full")] use super::{Node, Op}; +use crate::tree::kv::ValueDefinedCostType; #[cfg(feature = "full")] use crate::{ error::Error, tree::{Fetch, RefWalker}, Error::EdError, - TreeFeatureType::BasicMerk, + TreeFeatureType::BasicMerkNode, }; /// The minimum number of layers the trunk will be guaranteed to have before @@ -95,7 +96,10 @@ where depth: usize, ) -> CostResult { let mut cost = OperationCost::default(); - let maybe_left = match self.walk(true).unwrap_add_cost(&mut cost) { + let maybe_left = match self + .walk(true, None::<&fn(&[u8]) -> Option>) + .unwrap_add_cost(&mut cost) + { Ok(maybe_left) => maybe_left, Err(e) => { return Err(e).wrap_with_cost(cost); @@ -159,7 +163,11 @@ where // traverse left let has_left_child = self.tree().link(true).is_some(); if has_left_child { - let mut left = cost_return_on_error!(&mut cost, self.walk(true)).unwrap(); + let mut left = cost_return_on_error!( + &mut cost, + self.walk(true, None::<&fn(&[u8]) -> Option>) + ) + .unwrap(); cost_return_on_error!( &mut cost, left.traverse_for_trunk(proof, remaining_depth - 1, is_leftmost) @@ -174,7 +182,10 @@ where } // traverse right - if let Some(mut right) = cost_return_on_error!(&mut cost, self.walk(false)) { + if let Some(mut right) = cost_return_on_error!( + &mut cost, + self.walk(false, None::<&fn(&[u8]) -> Option>) + ) { cost_return_on_error!( &mut cost, right.traverse_for_trunk(proof, remaining_depth - 1, false) @@ -199,7 +210,7 @@ pub(crate) fn get_next_chunk( let mut chunk = Vec::with_capacity(512); let mut stack = Vec::with_capacity(32); - let mut node = TreeNode::new(vec![], vec![], None, BasicMerk).unwrap_add_cost(&mut cost); + let mut node = TreeNode::new(vec![], vec![], None, BasicMerkNode).unwrap_add_cost(&mut cost); while iter.valid().unwrap_add_cost(&mut cost) { let key = iter.key().unwrap_add_cost(&mut cost).unwrap(); @@ -213,7 +224,13 @@ pub(crate) fn get_next_chunk( let encoded_node = iter.value().unwrap_add_cost(&mut cost).unwrap(); cost_return_on_error_no_add!( &cost, - TreeNode::decode_into(&mut node, vec![], encoded_node).map_err(EdError) + TreeNode::decode_into( + &mut node, + vec![], + encoded_node, + None:: Option> + ) + .map_err(EdError) ); // TODO: Only use the KVValueHash if needed, saves 32 bytes @@ -459,7 +476,7 @@ mod tests { #[test] fn one_node_tree_trunk_roundtrip() { - let mut tree = BaseTree::new(vec![0], vec![], None, BasicMerk).unwrap(); + let mut tree = BaseTree::new(vec![0], vec![], None, BasicMerkNode).unwrap(); tree.commit(&mut NoopCommit {}, &|_, _| Ok(0)) .unwrap() .unwrap(); @@ -480,11 +497,11 @@ mod tests { // 0 // \ // 1 - let mut tree = BaseTree::new(vec![0], vec![], None, BasicMerk) + let mut tree = BaseTree::new(vec![0], vec![], None, BasicMerkNode) .unwrap() .attach( false, - Some(BaseTree::new(vec![1], vec![], None, BasicMerk).unwrap()), + Some(BaseTree::new(vec![1], vec![], None, BasicMerkNode).unwrap()), ); tree.commit(&mut NoopCommit {}, &|_, _| Ok(0)) .unwrap() @@ -505,11 +522,11 @@ mod tests { // 1 // / // 0 - let mut tree = BaseTree::new(vec![1], vec![], None, BasicMerk) + let mut tree = BaseTree::new(vec![1], vec![], None, BasicMerkNode) .unwrap() .attach( true, - Some(BaseTree::new(vec![0], vec![], None, BasicMerk).unwrap()), + Some(BaseTree::new(vec![0], vec![], None, BasicMerkNode).unwrap()), ); tree.commit(&mut NoopCommit {}, &|_, _| Ok(0)) .unwrap() @@ -530,15 +547,15 @@ mod tests { // 1 // / \ // 0 2 - let mut tree = BaseTree::new(vec![1], vec![], None, BasicMerk) + let mut tree = BaseTree::new(vec![1], vec![], None, BasicMerkNode) .unwrap() .attach( true, - Some(BaseTree::new(vec![0], vec![], None, BasicMerk).unwrap()), + Some(BaseTree::new(vec![0], vec![], None, BasicMerkNode).unwrap()), ) .attach( false, - Some(BaseTree::new(vec![2], vec![], None, BasicMerk).unwrap()), + Some(BaseTree::new(vec![2], vec![], None, BasicMerkNode).unwrap()), ); tree.commit(&mut NoopCommit {}, &|_, _| Ok(0)) .unwrap() diff --git a/merk/src/proofs/encoding.rs b/merk/src/proofs/encoding.rs index b0e82833..6b5f95b0 100644 --- a/merk/src/proofs/encoding.rs +++ b/merk/src/proofs/encoding.rs @@ -465,7 +465,7 @@ mod test { use super::super::{Node, Op}; use crate::{ tree::HASH_LENGTH, - TreeFeatureType::{BasicMerk, SummedMerk}, + TreeFeatureType::{BasicMerkNode, SummedMerkNode}, }; #[test] @@ -567,7 +567,7 @@ mod test { vec![1, 2, 3], vec![4, 5, 6], [0; 32], - BasicMerk, + BasicMerkNode, )); assert_eq!(op.encoding_length(), 43); @@ -585,7 +585,7 @@ mod test { vec![1, 2, 3], vec![4, 5, 6], [0; 32], - SummedMerk(6), + SummedMerkNode(6), )); assert_eq!(op.encoding_length(), 44); @@ -683,7 +683,7 @@ mod test { vec![1, 2, 3], vec![4, 5, 6], [0; 32], - BasicMerk, + BasicMerkNode, )); assert_eq!(op.encoding_length(), 43); @@ -701,7 +701,7 @@ mod test { vec![1, 2, 3], vec![4, 5, 6], [0; 32], - SummedMerk(5), + SummedMerkNode(5), )); assert_eq!(op.encoding_length(), 44); @@ -860,7 +860,7 @@ mod test { vec![1, 2, 3], vec![4, 5, 6], [0; 32], - BasicMerk + BasicMerkNode )) ); @@ -875,7 +875,7 @@ mod test { vec![1, 2, 3], vec![4, 5, 6], [0; 32], - SummedMerk(6) + SummedMerkNode(6) )) ); } @@ -960,7 +960,7 @@ mod test { vec![1, 2, 3], vec![4, 5, 6], [0; 32], - BasicMerk + BasicMerkNode )) ); @@ -975,7 +975,7 @@ mod test { vec![1, 2, 3], vec![4, 5, 6], [0; 32], - SummedMerk(6) + SummedMerkNode(6) )) ); } diff --git a/merk/src/proofs/query/mod.rs b/merk/src/proofs/query/mod.rs index b9ba161e..40aea905 100644 --- a/merk/src/proofs/query/mod.rs +++ b/merk/src/proofs/query/mod.rs @@ -67,6 +67,7 @@ use {super::Op, std::collections::LinkedList}; use super::Node; #[cfg(any(feature = "full", feature = "verify"))] use crate::error::Error; +use crate::tree::kv::ValueDefinedCostType; #[cfg(feature = "full")] use crate::tree::{Fetch, Link, RefWalker}; @@ -752,14 +753,15 @@ where left_to_right: bool, ) -> CostResult { if !query.is_empty() { - self.walk(left).flat_map_ok(|child_opt| { - if let Some(mut child) = child_opt { - child.create_proof(query, limit, offset, left_to_right) - } else { - Ok((LinkedList::new(), (true, true), limit, offset)) - .wrap_with_cost(Default::default()) - } - }) + self.walk(left, None::<&fn(&[u8]) -> Option>) + .flat_map_ok(|child_opt| { + if let Some(mut child) = child_opt { + child.create_proof(query, limit, offset, left_to_right) + } else { + Ok((LinkedList::new(), (true, true), limit, offset)) + .wrap_with_cost(Default::default()) + } + }) } else if let Some(link) = self.tree().link(left) { let mut proof = LinkedList::new(); proof.push_back(if left_to_right { @@ -793,7 +795,7 @@ mod test { }, test_utils::make_tree_seq, tree::{NoopCommit, PanicSource, RefWalker, TreeNode}, - TreeFeatureType::BasicMerk, + TreeFeatureType::BasicMerkNode, }; fn compare_result_tuples( @@ -808,15 +810,15 @@ mod test { } fn make_3_node_tree() -> TreeNode { - let mut tree = TreeNode::new(vec![5], vec![5], None, BasicMerk) + let mut tree = TreeNode::new(vec![5], vec![5], None, BasicMerkNode) .unwrap() .attach( true, - Some(TreeNode::new(vec![3], vec![3], None, BasicMerk).unwrap()), + Some(TreeNode::new(vec![3], vec![3], None, BasicMerkNode).unwrap()), ) .attach( false, - Some(TreeNode::new(vec![7], vec![7], None, BasicMerk).unwrap()), + Some(TreeNode::new(vec![7], vec![7], None, BasicMerkNode).unwrap()), ); tree.commit(&mut NoopCommit {}, &|_, _| Ok(0)) .unwrap() @@ -825,9 +827,9 @@ mod test { } fn make_6_node_tree() -> TreeNode { - let two_tree = TreeNode::new(vec![2], vec![2], None, BasicMerk).unwrap(); - let four_tree = TreeNode::new(vec![4], vec![4], None, BasicMerk).unwrap(); - let mut three_tree = TreeNode::new(vec![3], vec![3], None, BasicMerk) + let two_tree = TreeNode::new(vec![2], vec![2], None, BasicMerkNode).unwrap(); + let four_tree = TreeNode::new(vec![4], vec![4], None, BasicMerkNode).unwrap(); + let mut three_tree = TreeNode::new(vec![3], vec![3], None, BasicMerkNode) .unwrap() .attach(true, Some(two_tree)) .attach(false, Some(four_tree)); @@ -836,8 +838,8 @@ mod test { .unwrap() .expect("commit failed"); - let seven_tree = TreeNode::new(vec![7], vec![7], None, BasicMerk).unwrap(); - let mut eight_tree = TreeNode::new(vec![8], vec![8], None, BasicMerk) + let seven_tree = TreeNode::new(vec![7], vec![7], None, BasicMerkNode).unwrap(); + let mut eight_tree = TreeNode::new(vec![8], vec![8], None, BasicMerkNode) .unwrap() .attach(true, Some(seven_tree)); eight_tree @@ -845,7 +847,7 @@ mod test { .unwrap() .expect("commit failed"); - let mut root_tree = TreeNode::new(vec![5], vec![5], None, BasicMerk) + let mut root_tree = TreeNode::new(vec![5], vec![5], None, BasicMerkNode) .unwrap() .attach(true, Some(three_tree)) .attach(false, Some(eight_tree)); @@ -1533,26 +1535,26 @@ mod test { #[test] fn doc_proof() { - let mut tree = TreeNode::new(vec![5], vec![5], None, BasicMerk) + let mut tree = TreeNode::new(vec![5], vec![5], None, BasicMerkNode) .unwrap() .attach( true, Some( - TreeNode::new(vec![2], vec![2], None, BasicMerk) + TreeNode::new(vec![2], vec![2], None, BasicMerkNode) .unwrap() .attach( true, - Some(TreeNode::new(vec![1], vec![1], None, BasicMerk).unwrap()), + Some(TreeNode::new(vec![1], vec![1], None, BasicMerkNode).unwrap()), ) .attach( false, Some( - TreeNode::new(vec![4], vec![4], None, BasicMerk) + TreeNode::new(vec![4], vec![4], None, BasicMerkNode) .unwrap() .attach( true, Some( - TreeNode::new(vec![3], vec![3], None, BasicMerk) + TreeNode::new(vec![3], vec![3], None, BasicMerkNode) .unwrap(), ), ), @@ -1563,24 +1565,24 @@ mod test { .attach( false, Some( - TreeNode::new(vec![9], vec![9], None, BasicMerk) + TreeNode::new(vec![9], vec![9], None, BasicMerkNode) .unwrap() .attach( true, Some( - TreeNode::new(vec![7], vec![7], None, BasicMerk) + TreeNode::new(vec![7], vec![7], None, BasicMerkNode) .unwrap() .attach( true, Some( - TreeNode::new(vec![6], vec![6], None, BasicMerk) + TreeNode::new(vec![6], vec![6], None, BasicMerkNode) .unwrap(), ), ) .attach( false, Some( - TreeNode::new(vec![8], vec![8], None, BasicMerk) + TreeNode::new(vec![8], vec![8], None, BasicMerkNode) .unwrap(), ), ), @@ -1589,12 +1591,12 @@ mod test { .attach( false, Some( - TreeNode::new(vec![11], vec![11], None, BasicMerk) + TreeNode::new(vec![11], vec![11], None, BasicMerkNode) .unwrap() .attach( true, Some( - TreeNode::new(vec![10], vec![10], None, BasicMerk) + TreeNode::new(vec![10], vec![10], None, BasicMerkNode) .unwrap(), ), ), @@ -5740,7 +5742,7 @@ mod test { #[test] fn verify_ops() { - let mut tree = TreeNode::new(vec![5], vec![5], None, BasicMerk).unwrap(); + let mut tree = TreeNode::new(vec![5], vec![5], None, BasicMerkNode).unwrap(); tree.commit(&mut NoopCommit {}, &|_, _| Ok(0)) .unwrap() .expect("commit failed"); @@ -5766,7 +5768,7 @@ mod test { #[test] #[should_panic(expected = "verify failed")] fn verify_ops_mismatched_hash() { - let mut tree = TreeNode::new(vec![5], vec![5], None, BasicMerk).unwrap(); + let mut tree = TreeNode::new(vec![5], vec![5], None, BasicMerkNode).unwrap(); tree.commit(&mut NoopCommit {}, &|_, _| Ok(0)) .unwrap() .expect("commit failed"); diff --git a/merk/src/test_utils/mod.rs b/merk/src/test_utils/mod.rs index a9a087b5..6abe167e 100644 --- a/merk/src/test_utils/mod.rs +++ b/merk/src/test_utils/mod.rs @@ -39,9 +39,12 @@ use rand::prelude::*; pub use temp_merk::TempMerk; use crate::{ - tree::{kv::KV, BatchEntry, MerkBatch, NoopCommit, Op, PanicSource, TreeNode, Walker}, + tree::{ + kv::{ValueDefinedCostType, KV}, + BatchEntry, MerkBatch, NoopCommit, Op, PanicSource, TreeNode, Walker, + }, Merk, - TreeFeatureType::{BasicMerk, SummedMerk}, + TreeFeatureType::{BasicMerkNode, SummedMerkNode}, }; /// Assert tree invariants @@ -85,6 +88,7 @@ pub fn apply_memonly_unchecked(tree: TreeNode, batch: &MerkBatch>) -> Tr is_sum_node, )) }, + None::<&fn(&[u8]) -> Option>, &mut |_, _, _| Ok((false, None)), &mut |_flags, key_bytes_to_remove, value_bytes_to_remove| { Ok(( @@ -137,6 +141,7 @@ pub fn apply_to_memonly( is_sum_tree, )) }, + None::<&fn(&[u8]) -> Option>, &mut |_, _, _| Ok((false, None)), &mut |_flags, key_bytes_to_remove, value_bytes_to_remove| { Ok(( @@ -171,7 +176,7 @@ pub const fn seq_key(n: u64) -> [u8; 8] { /// Create batch entry with Put op using key n and a fixed value pub fn put_entry(n: u64) -> BatchEntry> { - (seq_key(n).to_vec(), Op::Put(vec![123; 60], BasicMerk)) + (seq_key(n).to_vec(), Op::Put(vec![123; 60], BasicMerkNode)) } /// Create batch entry with Delete op using key n @@ -235,9 +240,9 @@ pub fn make_tree_rand( let value = vec![123; 60]; let feature_type = if is_sum_tree { - SummedMerk(0) + SummedMerkNode(0) } else { - BasicMerk + BasicMerkNode }; let mut tree = TreeNode::new(vec![0; 20], value, None, feature_type).unwrap(); @@ -264,7 +269,7 @@ pub fn make_tree_seq(node_count: u64) -> TreeNode { }; let value = vec![123; 60]; - let mut tree = TreeNode::new(vec![0; 20], value, None, BasicMerk).unwrap(); + let mut tree = TreeNode::new(vec![0; 20], value, None, BasicMerkNode).unwrap(); let batch_count = node_count / batch_size; for i in 0..batch_count { @@ -288,6 +293,7 @@ where .get_storage_context(SubtreePath::empty(), Some(batch)) .unwrap(), false, + None:: Option>, ) .unwrap() .unwrap() @@ -305,6 +311,7 @@ where .get_storage_context(SubtreePath::empty(), None) .unwrap(), false, + None:: Option>, ) .unwrap() .unwrap() diff --git a/merk/src/test_utils/temp_merk.rs b/merk/src/test_utils/temp_merk.rs index 0fb4724e..25e5b75c 100644 --- a/merk/src/test_utils/temp_merk.rs +++ b/merk/src/test_utils/temp_merk.rs @@ -39,6 +39,7 @@ use grovedb_storage::{ Storage, }; +use crate::tree::kv::ValueDefinedCostType; #[cfg(feature = "full")] use crate::Merk; @@ -62,7 +63,13 @@ impl TempMerk { .get_storage_context(SubtreePath::empty(), Some(batch)) .unwrap(); - let merk = Merk::open_base(context, false).unwrap().unwrap(); + let merk = Merk::open_base( + context, + false, + None:: Option>, + ) + .unwrap() + .unwrap(); TempMerk { storage, merk, @@ -82,7 +89,13 @@ impl TempMerk { .storage .get_storage_context(SubtreePath::empty(), Some(self.batch)) .unwrap(); - self.merk = Merk::open_base(context, false).unwrap().unwrap(); + self.merk = Merk::open_base( + context, + false, + None:: Option>, + ) + .unwrap() + .unwrap(); } } diff --git a/merk/src/tree/encoding.rs b/merk/src/tree/encoding.rs index d7a56ce1..e635cd8c 100644 --- a/merk/src/tree/encoding.rs +++ b/merk/src/tree/encoding.rs @@ -39,6 +39,7 @@ use grovedb_storage::StorageContext; #[cfg(feature = "full")] use super::TreeNode; +use crate::tree::kv::ValueDefinedCostType; #[cfg(feature = "full")] use crate::{ error::{Error, Error::EdError}, @@ -50,12 +51,20 @@ use crate::{ impl TreeNode { /// Decode given bytes and set as Tree fields. Set key to value of given /// key. - pub fn decode_raw(bytes: &[u8], key: Vec) -> Result { - TreeNode::decode(key, bytes).map_err(EdError) + pub fn decode_raw( + bytes: &[u8], + key: Vec, + value_defined_cost_fn: Option Option>, + ) -> Result { + TreeNode::decode(key, bytes, value_defined_cost_fn).map_err(EdError) } /// Get value from storage given key. - pub(crate) fn get<'db, S, K>(storage: &S, key: K) -> CostResult, Error> + pub(crate) fn get<'db, S, K>( + storage: &S, + key: K, + value_defined_cost_fn: Option Option>, + ) -> CostResult, Error> where S: StorageContext<'db>, K: AsRef<[u8]>, @@ -66,7 +75,7 @@ impl TreeNode { let tree_opt = cost_return_on_error_no_add!( &cost, tree_bytes - .map(|x| TreeNode::decode_raw(&x, key.as_ref().to_vec())) + .map(|x| TreeNode::decode_raw(&x, key.as_ref().to_vec(), value_defined_cost_fn)) .transpose() ); @@ -111,18 +120,33 @@ impl TreeNode { #[inline] /// Decode bytes from reader, set as Tree fields and set key to given key - pub fn decode_into(&mut self, key: Vec, input: &[u8]) -> ed::Result<()> { + pub fn decode_into( + &mut self, + key: Vec, + input: &[u8], + value_defined_cost_fn: Option Option>, + ) -> ed::Result<()> { let mut tree_inner: TreeNodeInner = Decode::decode(input)?; tree_inner.kv.key = key; + if let Some(value_defined_cost_fn) = value_defined_cost_fn { + tree_inner.kv.value_defined_cost = value_defined_cost_fn(input); + } self.inner = Box::new(tree_inner); Ok(()) } #[inline] /// Decode input and set as Tree fields. Set the key as the given key. - pub fn decode(key: Vec, input: &[u8]) -> ed::Result { + pub fn decode( + key: Vec, + input: &[u8], + value_defined_cost_fn: Option Option>, + ) -> ed::Result { let mut tree_inner: TreeNodeInner = Decode::decode(input)?; tree_inner.kv.key = key; + if let Some(value_defined_cost_fn) = value_defined_cost_fn { + tree_inner.kv.value_defined_cost = value_defined_cost_fn(input); + } Ok(TreeNode::new_with_tree_inner(tree_inner)) } } @@ -131,12 +155,12 @@ impl TreeNode { #[cfg(test)] mod tests { use super::{super::Link, *}; - use crate::TreeFeatureType::{BasicMerk, SummedMerk}; + use crate::TreeFeatureType::{BasicMerkNode, SummedMerkNode}; #[test] fn encode_leaf_tree() { let tree = - TreeNode::from_fields(vec![0], vec![1], [55; 32], None, None, BasicMerk).unwrap(); + TreeNode::from_fields(vec![0], vec![1], [55; 32], None, None, BasicMerkNode).unwrap(); assert_eq!(tree.encoding_length(), 68); assert_eq!( tree.value_encoding_length_with_parent_to_child_reference(), @@ -163,10 +187,10 @@ mod tests { Some(Link::Modified { pending_writes: 1, child_heights: (123, 124), - tree: TreeNode::new(vec![2], vec![3], None, BasicMerk).unwrap(), + tree: TreeNode::new(vec![2], vec![3], None, BasicMerkNode).unwrap(), }), None, - BasicMerk, + BasicMerkNode, ) .unwrap(); tree.encode(); @@ -182,10 +206,10 @@ mod tests { hash: [66; 32], sum: None, child_heights: (123, 124), - tree: TreeNode::new(vec![2], vec![3], None, BasicMerk).unwrap(), + tree: TreeNode::new(vec![2], vec![3], None, BasicMerkNode).unwrap(), }), None, - BasicMerk, + BasicMerkNode, ) .unwrap(); assert_eq!( @@ -211,10 +235,10 @@ mod tests { hash: [66; 32], sum: Some(10), child_heights: (123, 124), - tree: TreeNode::new(vec![2], vec![3], None, BasicMerk).unwrap(), + tree: TreeNode::new(vec![2], vec![3], None, BasicMerkNode).unwrap(), }), None, - SummedMerk(5), + SummedMerkNode(5), ) .unwrap(); assert_eq!( @@ -243,7 +267,7 @@ mod tests { key: vec![2], }), None, - BasicMerk, + BasicMerkNode, ) .unwrap(); assert_eq!( @@ -276,10 +300,15 @@ mod tests { 131, 208, 25, 73, 98, 245, 209, 227, 170, 26, 72, 212, 134, 166, 126, 39, 98, 166, 199, 149, 144, 21, 1, ]; - let tree = TreeNode::decode(vec![0], bytes.as_slice()).expect("should decode correctly"); + let tree = TreeNode::decode( + vec![0], + bytes.as_slice(), + None::<&fn(&[u8]) -> Option>, + ) + .expect("should decode correctly"); assert_eq!(tree.key(), &[0]); assert_eq!(tree.value_as_slice(), &[1]); - assert_eq!(tree.inner.kv.feature_type, BasicMerk); + assert_eq!(tree.inner.kv.feature_type, BasicMerkNode); } #[test] @@ -291,7 +320,12 @@ mod tests { 55, 55, 55, 55, 55, 55, 32, 34, 236, 157, 87, 27, 167, 116, 207, 158, 131, 208, 25, 73, 98, 245, 209, 227, 170, 26, 72, 212, 134, 166, 126, 39, 98, 166, 199, 149, 144, 21, 1, ]; - let tree = TreeNode::decode(vec![0], bytes.as_slice()).expect("should decode correctly"); + let tree = TreeNode::decode( + vec![0], + bytes.as_slice(), + None::<&fn(&[u8]) -> Option>, + ) + .expect("should decode correctly"); assert_eq!(tree.key(), &[0]); assert_eq!(tree.value_as_slice(), &[1]); if let Some(Link::Reference { @@ -312,7 +346,11 @@ mod tests { #[test] fn decode_invalid_bytes_as_tree() { let bytes = vec![2, 3, 4, 5]; - let tree = TreeNode::decode(vec![0], bytes.as_slice()); + let tree = TreeNode::decode( + vec![0], + bytes.as_slice(), + None::<&fn(&[u8]) -> Option>, + ); assert!(matches!(tree, Err(_))); } } diff --git a/merk/src/tree/kv.rs b/merk/src/tree/kv.rs index 8855855a..ff020abc 100644 --- a/merk/src/tree/kv.rs +++ b/merk/src/tree/kv.rs @@ -45,7 +45,7 @@ use crate::tree::kv::ValueDefinedCostType::{LayeredValueDefinedCost, Specialized use crate::{ tree::{ hash::{combine_hash, kv_digest_to_kv_hash, value_hash, HASH_LENGTH_X2}, - tree_feature_type::{TreeFeatureType, TreeFeatureType::BasicMerk}, + tree_feature_type::{TreeFeatureType, TreeFeatureType::BasicMerkNode}, }, Link, HASH_LENGTH_U32, HASH_LENGTH_U32_X2, }; @@ -237,19 +237,6 @@ impl KV { self.wrap_with_cost(cost) } - /// Replaces the `KV`'s value with the given value, updates the hash, - /// value hash and returns the modified `KV`. - /// This is used when we want a fixed cost, for example in sum trees - #[inline] - pub fn put_value_with_fixed_cost_then_update( - mut self, - value: Vec, - value_cost: u32, - ) -> CostContext { - self.value_defined_cost = Some(SpecializedValueDefinedCost(value_cost)); - self.put_value_then_update(value) - } - /// Replaces the `KV`'s value with the given value, does not update the /// hashes, value hash and returns the modified `KV`. /// This is used when we want a fixed cost, for example in sum trees @@ -263,37 +250,6 @@ impl KV { self.put_value_no_update_of_hashes(value) } - /// Replaces the `KV`'s value with the given value and value hash, - /// updates the hash and returns the modified `KV`. - #[inline] - pub fn put_value_and_reference_value_hash_then_update( - mut self, - value: Vec, - reference_value_hash: CryptoHash, - ) -> CostContext { - let mut cost = OperationCost::default(); - let actual_value_hash = value_hash(value.as_slice()).unwrap_add_cost(&mut cost); - let combined_value_hash = - combine_hash(&actual_value_hash, &reference_value_hash).unwrap_add_cost(&mut cost); - self.value = value; - self.value_hash = combined_value_hash; - self.hash = kv_digest_to_kv_hash(self.key(), self.value_hash()).unwrap_add_cost(&mut cost); - self.wrap_with_cost(cost) - } - - /// Replaces the `KV`'s value with the given value and value hash, - /// updates the hash and returns the modified `KV`. - #[inline] - pub fn put_value_with_reference_value_hash_and_value_cost_then_update( - mut self, - value: Vec, - reference_value_hash: CryptoHash, - value_cost: u32, - ) -> CostContext { - self.value_defined_cost = Some(LayeredValueDefinedCost(value_cost)); - self.put_value_and_reference_value_hash_then_update(value, reference_value_hash) - } - /// Returns the key as a slice. #[inline] pub fn key(&self) -> &[u8] { @@ -573,7 +529,7 @@ impl Decode for KV { let mut kv = Self { key: Vec::with_capacity(0), value: Vec::with_capacity(128), - feature_type: BasicMerk, + feature_type: BasicMerkNode, value_defined_cost: None, hash: NULL_HASH, value_hash: NULL_HASH, @@ -604,11 +560,11 @@ impl Terminated for KV {} #[cfg(test)] mod test { use super::*; - use crate::tree::tree_feature_type::TreeFeatureType::SummedMerk; + use crate::tree::tree_feature_type::TreeFeatureType::SummedMerkNode; #[test] fn new_kv() { - let kv = KV::new(vec![1, 2, 3], vec![4, 5, 6], None, BasicMerk).unwrap(); + let kv = KV::new(vec![1, 2, 3], vec![4, 5, 6], None, BasicMerkNode).unwrap(); assert_eq!(kv.key(), &[1, 2, 3]); assert_eq!(kv.value_as_slice(), &[4, 5, 6]); @@ -617,7 +573,7 @@ mod test { #[test] fn with_value() { - let kv = KV::new(vec![1, 2, 3], vec![4, 5, 6], None, BasicMerk) + let kv = KV::new(vec![1, 2, 3], vec![4, 5, 6], None, BasicMerkNode) .unwrap() .put_value_then_update(vec![7, 8, 9]) .unwrap(); @@ -629,7 +585,7 @@ mod test { #[test] fn encode_and_decode_kv() { - let kv = KV::new(vec![1, 2, 3], vec![4, 5, 6], None, BasicMerk).unwrap(); + let kv = KV::new(vec![1, 2, 3], vec![4, 5, 6], None, BasicMerkNode).unwrap(); let mut encoded_kv = vec![]; kv.encode_into(&mut encoded_kv).expect("encoded"); let mut decoded_kv = KV::decode(encoded_kv.as_slice()).unwrap(); @@ -637,7 +593,7 @@ mod test { assert_eq!(kv, decoded_kv); - let kv = KV::new(vec![1, 2, 3], vec![4, 5, 6], None, SummedMerk(20)).unwrap(); + let kv = KV::new(vec![1, 2, 3], vec![4, 5, 6], None, SummedMerkNode(20)).unwrap(); let mut encoded_kv = vec![]; kv.encode_into(&mut encoded_kv).expect("encoded"); let mut decoded_kv = KV::decode(encoded_kv.as_slice()).unwrap(); diff --git a/merk/src/tree/link.rs b/merk/src/tree/link.rs index 8cdb948b..ab26159b 100644 --- a/merk/src/tree/link.rs +++ b/merk/src/tree/link.rs @@ -486,11 +486,11 @@ mod test { super::{hash::NULL_HASH, TreeNode}, *, }; - use crate::TreeFeatureType::BasicMerk; + use crate::TreeFeatureType::BasicMerkNode; #[test] fn from_modified_tree() { - let tree = TreeNode::new(vec![0], vec![1], None, BasicMerk).unwrap(); + let tree = TreeNode::new(vec![0], vec![1], None, BasicMerkNode).unwrap(); let link = Link::from_modified_tree(tree); assert!(link.is_modified()); assert_eq!(link.height(), 1); @@ -507,7 +507,7 @@ mod test { let link = Link::maybe_from_modified_tree(None); assert!(link.is_none()); - let tree = TreeNode::new(vec![0], vec![1], None, BasicMerk).unwrap(); + let tree = TreeNode::new(vec![0], vec![1], None, BasicMerkNode).unwrap(); let link = Link::maybe_from_modified_tree(Some(tree)); assert!(link.expect("expected link").is_modified()); } @@ -519,7 +519,7 @@ mod test { let child_heights = (0, 0); let pending_writes = 1; let key = vec![0]; - let tree = || TreeNode::new(vec![0], vec![1], None, BasicMerk).unwrap(); + let tree = || TreeNode::new(vec![0], vec![1], None, BasicMerkNode).unwrap(); let reference = Link::Reference { hash, @@ -585,7 +585,7 @@ mod test { Link::Modified { pending_writes: 1, child_heights: (1, 1), - tree: TreeNode::new(vec![0], vec![1], None, BasicMerk).unwrap(), + tree: TreeNode::new(vec![0], vec![1], None, BasicMerkNode).unwrap(), } .hash(); } @@ -596,7 +596,7 @@ mod test { Link::Modified { pending_writes: 1, child_heights: (1, 1), - tree: TreeNode::new(vec![0], vec![1], None, BasicMerk).unwrap(), + tree: TreeNode::new(vec![0], vec![1], None, BasicMerkNode).unwrap(), } .into_reference(); } @@ -608,7 +608,7 @@ mod test { hash: [1; 32], sum: None, child_heights: (1, 1), - tree: TreeNode::new(vec![0], vec![1], None, BasicMerk).unwrap(), + tree: TreeNode::new(vec![0], vec![1], None, BasicMerkNode).unwrap(), } .into_reference(); } diff --git a/merk/src/tree/mod.rs b/merk/src/tree/mod.rs index 5f15ed94..c57d3644 100644 --- a/merk/src/tree/mod.rs +++ b/merk/src/tree/mod.rs @@ -487,8 +487,8 @@ impl TreeNode { #[inline] pub fn sum(&self) -> Result, Error> { match self.inner.kv.feature_type { - TreeFeatureType::BasicMerk => Ok(None), - TreeFeatureType::SummedMerk(value) => value + TreeFeatureType::BasicMerkNode => Ok(None), + TreeFeatureType::SummedMerkNode(value) => value .checked_add(self.child_sum(true)) .and_then(|a| a.checked_add(self.child_sum(false))) .ok_or(Overflow("sum is overflowing")) @@ -944,7 +944,15 @@ impl TreeNode { /// Fetches the child on the given side using the given data source, and /// places it in the child slot (upgrading the link from `Link::Reference` /// to `Link::Loaded`). - pub fn load(&mut self, left: bool, source: &S) -> CostResult<(), Error> { + pub fn load( + &mut self, + left: bool, + source: &S, + value_defined_cost_fn: Option<&V>, + ) -> CostResult<(), Error> + where + V: Fn(&[u8]) -> Option, + { // TODO: return Err instead of panic? let link = self.link(left).expect("Expected link"); let (child_heights, hash, sum) = match link { @@ -958,7 +966,7 @@ impl TreeNode { }; let mut cost = OperationCost::default(); - let tree = cost_return_on_error!(&mut cost, source.fetch(link)); + let tree = cost_return_on_error!(&mut cost, source.fetch(link, value_defined_cost_fn)); debug_assert_eq!(tree.key(), link.key()); *self.slot_mut(left) = Some(Link::Loaded { tree, @@ -986,11 +994,13 @@ mod test { use grovedb_costs::storage_cost::removal::StorageRemovedBytes::NoStorageRemoval; use super::{commit::NoopCommit, hash::NULL_HASH, TreeNode}; - use crate::tree::{tree_feature_type::TreeFeatureType::SummedMerk, TreeFeatureType::BasicMerk}; + use crate::tree::{ + tree_feature_type::TreeFeatureType::SummedMerkNode, TreeFeatureType::BasicMerkNode, + }; #[test] fn build_tree() { - let tree = TreeNode::new(vec![1], vec![101], None, BasicMerk).unwrap(); + let tree = TreeNode::new(vec![1], vec![101], None, BasicMerkNode).unwrap(); assert_eq!(tree.key(), &[1]); assert_eq!(tree.value_as_slice(), &[101]); assert!(tree.child(true).is_none()); @@ -1002,13 +1012,13 @@ mod test { let tree = tree.attach( true, - Some(TreeNode::new(vec![2], vec![102], None, BasicMerk).unwrap()), + Some(TreeNode::new(vec![2], vec![102], None, BasicMerkNode).unwrap()), ); assert_eq!(tree.key(), &[1]); assert_eq!(tree.child(true).unwrap().key(), &[2]); assert!(tree.child(false).is_none()); - let tree = TreeNode::new(vec![3], vec![103], None, BasicMerk) + let tree = TreeNode::new(vec![3], vec![103], None, BasicMerkNode) .unwrap() .attach(false, Some(tree)); assert_eq!(tree.key(), &[3]); @@ -1019,29 +1029,29 @@ mod test { #[should_panic] #[test] fn attach_existing() { - TreeNode::new(vec![0], vec![1], None, BasicMerk) + TreeNode::new(vec![0], vec![1], None, BasicMerkNode) .unwrap() .attach( true, - Some(TreeNode::new(vec![2], vec![3], None, BasicMerk).unwrap()), + Some(TreeNode::new(vec![2], vec![3], None, BasicMerkNode).unwrap()), ) .attach( true, - Some(TreeNode::new(vec![4], vec![5], None, BasicMerk).unwrap()), + Some(TreeNode::new(vec![4], vec![5], None, BasicMerkNode).unwrap()), ); } #[test] fn modify() { - let tree = TreeNode::new(vec![0], vec![1], None, BasicMerk) + let tree = TreeNode::new(vec![0], vec![1], None, BasicMerkNode) .unwrap() .attach( true, - Some(TreeNode::new(vec![2], vec![3], None, BasicMerk).unwrap()), + Some(TreeNode::new(vec![2], vec![3], None, BasicMerkNode).unwrap()), ) .attach( false, - Some(TreeNode::new(vec![4], vec![5], None, BasicMerk).unwrap()), + Some(TreeNode::new(vec![4], vec![5], None, BasicMerkNode).unwrap()), ); let tree = tree.walk(true, |left_opt| { @@ -1053,7 +1063,7 @@ mod test { let tree = tree.walk(true, |left_opt| { assert!(left_opt.is_none()); - Some(TreeNode::new(vec![2], vec![3], None, BasicMerk).unwrap()) + Some(TreeNode::new(vec![2], vec![3], None, BasicMerkNode).unwrap()) }); assert_eq!(tree.link(true).unwrap().key(), &[2]); @@ -1067,11 +1077,11 @@ mod test { #[test] fn child_and_link() { - let mut tree = TreeNode::new(vec![0], vec![1], None, BasicMerk) + let mut tree = TreeNode::new(vec![0], vec![1], None, BasicMerkNode) .unwrap() .attach( true, - Some(TreeNode::new(vec![2], vec![3], None, BasicMerk).unwrap()), + Some(TreeNode::new(vec![2], vec![3], None, BasicMerkNode).unwrap()), ); assert!(tree.link(true).expect("expected link").is_modified()); assert!(tree.child(true).is_some()); @@ -1095,11 +1105,11 @@ mod test { #[test] fn child_hash() { - let mut tree = TreeNode::new(vec![0], vec![1], None, BasicMerk) + let mut tree = TreeNode::new(vec![0], vec![1], None, BasicMerkNode) .unwrap() .attach( true, - Some(TreeNode::new(vec![2], vec![3], None, BasicMerk).unwrap()), + Some(TreeNode::new(vec![2], vec![3], None, BasicMerkNode).unwrap()), ); tree.commit(&mut NoopCommit {}, &|_, _| Ok(0)) .unwrap() @@ -1116,7 +1126,7 @@ mod test { #[test] fn hash() { - let tree = TreeNode::new(vec![0], vec![1], None, BasicMerk).unwrap(); + let tree = TreeNode::new(vec![0], vec![1], None, BasicMerkNode).unwrap(); assert_eq!( tree.hash().unwrap(), [ @@ -1128,13 +1138,13 @@ mod test { #[test] fn child_pending_writes() { - let tree = TreeNode::new(vec![0], vec![1], None, BasicMerk).unwrap(); + let tree = TreeNode::new(vec![0], vec![1], None, BasicMerkNode).unwrap(); assert_eq!(tree.child_pending_writes(true), 0); assert_eq!(tree.child_pending_writes(false), 0); let tree = tree.attach( true, - Some(TreeNode::new(vec![2], vec![3], None, BasicMerk).unwrap()), + Some(TreeNode::new(vec![2], vec![3], None, BasicMerkNode).unwrap()), ); assert_eq!(tree.child_pending_writes(true), 1); assert_eq!(tree.child_pending_writes(false), 0); @@ -1142,7 +1152,7 @@ mod test { #[test] fn height_and_balance() { - let tree = TreeNode::new(vec![0], vec![1], None, BasicMerk).unwrap(); + let tree = TreeNode::new(vec![0], vec![1], None, BasicMerkNode).unwrap(); assert_eq!(tree.height(), 1); assert_eq!(tree.child_height(true), 0); assert_eq!(tree.child_height(false), 0); @@ -1150,7 +1160,7 @@ mod test { let tree = tree.attach( true, - Some(TreeNode::new(vec![2], vec![3], None, BasicMerk).unwrap()), + Some(TreeNode::new(vec![2], vec![3], None, BasicMerkNode).unwrap()), ); assert_eq!(tree.height(), 2); assert_eq!(tree.child_height(true), 1); @@ -1167,11 +1177,11 @@ mod test { #[test] fn commit() { - let mut tree = TreeNode::new(vec![0], vec![1], None, BasicMerk) + let mut tree = TreeNode::new(vec![0], vec![1], None, BasicMerkNode) .unwrap() .attach( false, - Some(TreeNode::new(vec![2], vec![3], None, BasicMerk).unwrap()), + Some(TreeNode::new(vec![2], vec![3], None, BasicMerkNode).unwrap()), ); tree.commit(&mut NoopCommit {}, &|_, _| Ok(0)) .unwrap() @@ -1182,11 +1192,11 @@ mod test { #[test] fn sum_tree() { - let mut tree = TreeNode::new(vec![0], vec![1], None, SummedMerk(3)) + let mut tree = TreeNode::new(vec![0], vec![1], None, SummedMerkNode(3)) .unwrap() .attach( false, - Some(TreeNode::new(vec![2], vec![3], None, SummedMerk(5)).unwrap()), + Some(TreeNode::new(vec![2], vec![3], None, SummedMerkNode(5)).unwrap()), ); tree.commit(&mut NoopCommit {}, &|_, _| Ok(0)) .unwrap() diff --git a/merk/src/tree/ops.rs b/merk/src/tree/ops.rs index 82055b11..4b3a3d34 100644 --- a/merk/src/tree/ops.rs +++ b/merk/src/tree/ops.rs @@ -150,7 +150,11 @@ pub struct PanicSource {} #[cfg(feature = "full")] impl Fetch for PanicSource { - fn fetch(&self, _link: &Link) -> CostResult { + fn fetch( + &self, + _link: &Link, + value_defined_cost_fn: Option<&impl Fn(&[u8]) -> Option>, + ) -> CostResult { unreachable!("'fetch' should not have been called") } } @@ -165,16 +169,18 @@ where /// not require a non-empty tree. /// /// Keys in batch must be sorted and unique. - pub fn apply_to, C, U, R>( + pub fn apply_to, C, V, U, R>( maybe_tree: Option, batch: &MerkBatch, source: S, old_tree_cost: &C, + value_defined_cost_fn: Option<&V>, update_tree_value_based_on_costs: &mut U, section_removal_bytes: &mut R, ) -> CostContext, KeyUpdates), Error>> where C: Fn(&Vec, &Vec) -> Result, + V: Fn(&[u8]) -> Option, U: FnMut( &StorageCost, &Vec, @@ -201,6 +207,7 @@ where batch, source, old_tree_cost, + value_defined_cost_fn, update_tree_value_based_on_costs, section_removal_bytes, ) @@ -226,6 +233,7 @@ where tree.apply_sorted( batch, old_tree_cost, + value_defined_cost_fn, update_tree_value_based_on_costs, section_removal_bytes ) @@ -241,15 +249,17 @@ where /// Builds a `Tree` from a batch of operations. /// /// Keys in batch must be sorted and unique. - fn build, C, U, R>( + fn build, C, V, U, R>( batch: &MerkBatch, source: S, old_tree_cost: &C, + value_defined_cost_fn: Option<&V>, update_tree_value_based_on_costs: &mut U, section_removal_bytes: &mut R, ) -> CostResult, Error> where C: Fn(&Vec, &Vec) -> Result, + V: Fn(&[u8]) -> Option, U: FnMut( &StorageCost, &Vec, @@ -276,6 +286,7 @@ where left_batch, source.clone(), old_tree_cost, + value_defined_cost_fn, update_tree_value_based_on_costs, section_removal_bytes ) @@ -288,6 +299,7 @@ where tree.apply_sorted( right_batch, old_tree_cost, + value_defined_cost_fn, update_tree_value_based_on_costs, section_removal_bytes ) @@ -300,6 +312,7 @@ where right_batch, source.clone(), old_tree_cost, + value_defined_cost_fn, update_tree_value_based_on_costs, section_removal_bytes ) @@ -370,6 +383,7 @@ where None ), old_tree_cost, + value_defined_cost_fn, update_tree_value_based_on_costs, section_removal_bytes, ) @@ -387,6 +401,7 @@ where self.apply_sorted( batch, &|_, _| Ok(0), + None::<&fn(&[u8]) -> Option>, &mut |_, _, _| Ok((false, None)), &mut |_flags, key_bytes_to_remove, value_bytes_to_remove| { Ok(( @@ -401,15 +416,17 @@ where /// `Walker::apply`_to, but requires a populated tree. /// /// Keys in batch must be sorted and unique. - fn apply_sorted, C, U, R>( + fn apply_sorted, C, V, U, R>( self, batch: &MerkBatch, old_specialized_cost: &C, + value_defined_cost_fn: Option<&V>, update_tree_value_based_on_costs: &mut U, section_removal_bytes: &mut R, ) -> CostResult<(Option, KeyUpdates), Error> where C: Fn(&Vec, &Vec) -> Result, + V: Fn(&[u8]) -> Option, U: FnMut( &StorageCost, &Vec, @@ -527,7 +544,8 @@ where needs_value_verification: false, }; - let maybe_tree = cost_return_on_error!(&mut cost, self.remove()); + let maybe_tree = + cost_return_on_error!(&mut cost, self.remove(value_defined_cost_fn)); #[rustfmt::skip] let (maybe_tree, mut key_updates) @@ -538,6 +556,7 @@ where &batch[..index], source.clone(), old_specialized_cost, + value_defined_cost_fn, update_tree_value_based_on_costs, section_removal_bytes ) @@ -551,6 +570,7 @@ where &batch[index + 1..], source.clone(), old_specialized_cost, + value_defined_cost_fn, update_tree_value_based_on_costs, section_removal_bytes ) @@ -593,6 +613,7 @@ where exclusive, KeyUpdates::new(new_keys, updated_keys, LinkedList::default(), None), old_specialized_cost, + value_defined_cost_fn, update_tree_value_based_on_costs, section_removal_bytes, ) @@ -604,13 +625,14 @@ where /// /// This recursion executes serially in the same thread, but in the future /// will be dispatched to workers in other threads. - fn recurse, C, U, R>( + fn recurse, C, V, U, R>( self, batch: &MerkBatch, mid: usize, exclusive: bool, mut key_updates: KeyUpdates, old_tree_cost: &C, + value_defined_cost_fn: Option<&V>, update_tree_value_based_on_costs: &mut U, section_removal_bytes: &mut R, ) -> CostResult<(Option, KeyUpdates), Error> @@ -621,6 +643,7 @@ where &Vec, &mut Vec, ) -> Result<(bool, Option), Error>, + V: Fn(&[u8]) -> Option, R: FnMut(&Vec, u32, u32) -> Result<(StorageRemovedBytes, StorageRemovedBytes), Error>, { let mut cost = OperationCost::default(); @@ -638,26 +661,31 @@ where let source = self.clone_source(); cost_return_on_error!( &mut cost, - self.walk(true, |maybe_left| { - Self::apply_to( - maybe_left, - left_batch, - source, - old_tree_cost, - update_tree_value_based_on_costs, - section_removal_bytes, - ) - .map_ok(|(maybe_left, mut key_updates_left)| { - key_updates.new_keys.append(&mut key_updates_left.new_keys); - key_updates - .updated_keys - .append(&mut key_updates_left.updated_keys); - key_updates - .deleted_keys - .append(&mut key_updates_left.deleted_keys); - maybe_left - }) - }) + self.walk( + true, + |maybe_left| { + Self::apply_to( + maybe_left, + left_batch, + source, + old_tree_cost, + value_defined_cost_fn, + update_tree_value_based_on_costs, + section_removal_bytes, + ) + .map_ok(|(maybe_left, mut key_updates_left)| { + key_updates.new_keys.append(&mut key_updates_left.new_keys); + key_updates + .updated_keys + .append(&mut key_updates_left.updated_keys); + key_updates + .deleted_keys + .append(&mut key_updates_left.deleted_keys); + maybe_left + }) + }, + value_defined_cost_fn + ) ) } else { self @@ -667,32 +695,37 @@ where let source = tree.clone_source(); cost_return_on_error!( &mut cost, - tree.walk(false, |maybe_right| { - Self::apply_to( - maybe_right, - right_batch, - source, - old_tree_cost, - update_tree_value_based_on_costs, - section_removal_bytes, - ) - .map_ok(|(maybe_right, mut key_updates_right)| { - key_updates.new_keys.append(&mut key_updates_right.new_keys); - key_updates - .updated_keys - .append(&mut key_updates_right.updated_keys); - key_updates - .deleted_keys - .append(&mut key_updates_right.deleted_keys); - maybe_right - }) - }) + tree.walk( + false, + |maybe_right| { + Self::apply_to( + maybe_right, + right_batch, + source, + old_tree_cost, + value_defined_cost_fn, + update_tree_value_based_on_costs, + section_removal_bytes, + ) + .map_ok(|(maybe_right, mut key_updates_right)| { + key_updates.new_keys.append(&mut key_updates_right.new_keys); + key_updates + .updated_keys + .append(&mut key_updates_right.updated_keys); + key_updates + .deleted_keys + .append(&mut key_updates_right.deleted_keys); + maybe_right + }) + }, + value_defined_cost_fn + ) ) } else { tree }; - let tree = cost_return_on_error!(&mut cost, tree.maybe_balance()); + let tree = cost_return_on_error!(&mut cost, tree.maybe_balance(value_defined_cost_fn)); let new_root_key = tree.tree().key(); @@ -715,7 +748,10 @@ where /// Checks if the tree is unbalanced and if so, applies AVL tree rotation(s) /// to rebalance the tree and its subtrees. Returns the root node of the /// balanced tree after applying the rotations. - fn maybe_balance(self) -> CostResult { + fn maybe_balance(self, value_defined_cost_fn: Option<&V>) -> CostResult + where + V: Fn(&[u8]) -> Option, + { let mut cost = OperationCost::default(); let balance_factor = self.balance_factor(); @@ -729,37 +765,55 @@ where let tree = if left == (self.tree().link(left).unwrap().balance_factor() > 0) { cost_return_on_error!( &mut cost, - self.walk_expect(left, |child| child.rotate(!left).map_ok(Option::Some)) + self.walk_expect( + left, + |child| child + .rotate(!left, value_defined_cost_fn) + .map_ok(Option::Some), + value_defined_cost_fn + ) ) } else { self }; - let rotate = tree.rotate(left).unwrap_add_cost(&mut cost); + let rotate = tree + .rotate(left, value_defined_cost_fn) + .unwrap_add_cost(&mut cost); rotate.wrap_with_cost(cost) } /// Applies an AVL tree rotation, a constant-time operation which only needs /// to swap pointers in order to rebalance a tree. - fn rotate(self, left: bool) -> CostResult { + fn rotate(self, left: bool, value_defined_cost_fn: Option<&V>) -> CostResult + where + V: Fn(&[u8]) -> Option, + { let mut cost = OperationCost::default(); - let (tree, child) = cost_return_on_error!(&mut cost, self.detach_expect(left)); - let (child, maybe_grandchild) = cost_return_on_error!(&mut cost, child.detach(!left)); + let (tree, child) = + cost_return_on_error!(&mut cost, self.detach_expect(left, value_defined_cost_fn)); + let (child, maybe_grandchild) = + cost_return_on_error!(&mut cost, child.detach(!left, value_defined_cost_fn)); // attach grandchild to self tree.attach(left, maybe_grandchild) - .maybe_balance() + .maybe_balance(value_defined_cost_fn) .flat_map_ok(|tree| { // attach self to child, return child - child.attach(!left, Some(tree)).maybe_balance() + child + .attach(!left, Some(tree)) + .maybe_balance(value_defined_cost_fn) }) .add_cost(cost) } /// Removes the root node from the tree. Rearranges and rebalances /// descendants (if any) in order to maintain a valid tree. - pub fn remove(self) -> CostResult, Error> { + pub fn remove(self, value_defined_cost_fn: Option<&V>) -> CostResult, Error> + where + V: Fn(&[u8]) -> Option, + { let mut cost = OperationCost::default(); let tree = self.tree(); @@ -769,14 +823,20 @@ where let maybe_tree = if has_left && has_right { // two children, promote edge of taller child - let (tree, tall_child) = cost_return_on_error!(&mut cost, self.detach_expect(left)); - let (_, short_child) = cost_return_on_error!(&mut cost, tree.detach_expect(!left)); - let promoted = - cost_return_on_error!(&mut cost, tall_child.promote_edge(!left, short_child)); + let (tree, tall_child) = + cost_return_on_error!(&mut cost, self.detach_expect(left, value_defined_cost_fn)); + let (_, short_child) = + cost_return_on_error!(&mut cost, tree.detach_expect(!left, value_defined_cost_fn)); + let promoted = cost_return_on_error!( + &mut cost, + tall_child.promote_edge(!left, short_child, value_defined_cost_fn) + ); Some(promoted) } else if has_left || has_right { // single child, promote it - Some(cost_return_on_error!(&mut cost, self.detach_expect(left)).1) + Some( + cost_return_on_error!(&mut cost, self.detach_expect(left, value_defined_cost_fn)).1, + ) } else { // no child None @@ -789,31 +849,49 @@ where /// reattaches it at the top in order to fill in a gap when removing a root /// node from a tree with both left and right children. Attaches `attach` on /// the opposite side. Returns the promoted node. - fn promote_edge(self, left: bool, attach: Self) -> CostResult { - self.remove_edge(left).flat_map_ok(|(edge, maybe_child)| { - edge.attach(!left, maybe_child) - .attach(left, Some(attach)) - .maybe_balance() - }) + fn promote_edge( + self, + left: bool, + attach: Self, + value_defined_cost_fn: Option<&V>, + ) -> CostResult + where + V: Fn(&[u8]) -> Option, + { + self.remove_edge(left, value_defined_cost_fn) + .flat_map_ok(|(edge, maybe_child)| { + edge.attach(!left, maybe_child) + .attach(left, Some(attach)) + .maybe_balance(value_defined_cost_fn) + }) } /// Traverses to the tree's edge on the given side and detaches it /// (reattaching its child, if any, to its former parent). Return value is /// `(edge, maybe_updated_tree)`. - fn remove_edge(self, left: bool) -> CostResult<(Self, Option), Error> { + fn remove_edge( + self, + left: bool, + value_defined_cost_fn: Option<&V>, + ) -> CostResult<(Self, Option), Error> + where + V: Fn(&[u8]) -> Option, + { let mut cost = OperationCost::default(); if self.tree().link(left).is_some() { // this node is not the edge, recurse - let (tree, child) = cost_return_on_error!(&mut cost, self.detach_expect(left)); - let (edge, maybe_child) = cost_return_on_error!(&mut cost, child.remove_edge(left)); + let (tree, child) = + cost_return_on_error!(&mut cost, self.detach_expect(left, value_defined_cost_fn)); + let (edge, maybe_child) = + cost_return_on_error!(&mut cost, child.remove_edge(left, value_defined_cost_fn)); tree.attach(left, maybe_child) - .maybe_balance() + .maybe_balance(value_defined_cost_fn) .map_ok(|tree| (edge, Some(tree))) .add_cost(cost) } else { // this node is the edge, detach its child if present - self.detach(!left) + self.detach(!left, value_defined_cost_fn) } } } @@ -824,13 +902,13 @@ mod test { use super::*; use crate::{ test_utils::{apply_memonly, assert_tree_invariants, del_entry, make_tree_seq, seq_key}, - tree::{tree_feature_type::TreeFeatureType::BasicMerk, *}, + tree::{tree_feature_type::TreeFeatureType::BasicMerkNode, *}, }; #[test] fn simple_insert() { - let batch = [(b"foo2".to_vec(), Op::Put(b"bar2".to_vec(), BasicMerk))]; - let tree = TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerk).unwrap(); + let batch = [(b"foo2".to_vec(), Op::Put(b"bar2".to_vec(), BasicMerkNode))]; + let tree = TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerkNode).unwrap(); let (maybe_walker, key_updates) = Walker::new(tree, PanicSource {}) .apply_sorted_without_costs(&batch) .unwrap() @@ -845,8 +923,8 @@ mod test { #[test] fn simple_update() { - let batch = [(b"foo".to_vec(), Op::Put(b"bar2".to_vec(), BasicMerk))]; - let tree = TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerk).unwrap(); + let batch = [(b"foo".to_vec(), Op::Put(b"bar2".to_vec(), BasicMerkNode))]; + let tree = TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerkNode).unwrap(); let (maybe_walker, key_updates) = Walker::new(tree, PanicSource {}) .apply_sorted_without_costs(&batch) .unwrap() @@ -872,9 +950,10 @@ mod test { hash: [123; 32], sum: None, child_heights: (0, 0), - tree: TreeNode::new(b"foo2".to_vec(), b"bar2".to_vec(), None, BasicMerk).unwrap(), + tree: TreeNode::new(b"foo2".to_vec(), b"bar2".to_vec(), None, BasicMerkNode) + .unwrap(), }), - BasicMerk, + BasicMerkNode, ) .unwrap(); let (maybe_walker, key_updates) = Walker::new(tree, PanicSource {}) @@ -897,7 +976,7 @@ mod test { #[test] fn delete_non_existent() { let batch = [(b"foo2".to_vec(), Op::Delete)]; - let tree = TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerk).unwrap(); + let tree = TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerkNode).unwrap(); Walker::new(tree, PanicSource {}) .apply_sorted_without_costs(&batch) .unwrap() @@ -907,7 +986,7 @@ mod test { #[test] fn delete_only_node() { let batch = [(b"foo".to_vec(), Op::Delete)]; - let tree = TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerk).unwrap(); + let tree = TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerkNode).unwrap(); let (maybe_walker, key_updates) = Walker::new(tree, PanicSource {}) .apply_sorted_without_costs(&batch) .unwrap() @@ -977,11 +1056,12 @@ mod test { #[test] fn apply_empty_none() { - let (maybe_tree, key_updates) = Walker::::apply_to::, _, _, _>( + let (maybe_tree, key_updates) = Walker::::apply_to::, _, _, _, _>( None, &[], PanicSource {}, &|_, _| Ok(0), + None::<&fn(&[u8]) -> Option>, &mut |_, _, _| Ok((false, None)), &mut |_flags, key_bytes_to_remove, value_bytes_to_remove| { Ok(( @@ -999,12 +1079,13 @@ mod test { #[test] fn insert_empty_single() { - let batch = vec![(vec![0], Op::Put(vec![1], BasicMerk))]; + let batch = vec![(vec![0], Op::Put(vec![1], BasicMerkNode))]; let (maybe_tree, key_updates) = Walker::::apply_to( None, &batch, PanicSource {}, &|_, _| Ok(0), + None::<&fn(&[u8]) -> Option>, &mut |_, _, _| Ok((false, None)), &mut |_flags, key_bytes_to_remove, value_bytes_to_remove| { Ok(( @@ -1025,12 +1106,13 @@ mod test { #[test] fn insert_updated_single() { - let batch = vec![(vec![0], Op::Put(vec![1], BasicMerk))]; + let batch = vec![(vec![0], Op::Put(vec![1], BasicMerkNode))]; let (maybe_tree, key_updates) = Walker::::apply_to( None, &batch, PanicSource {}, &|_, _| Ok(0), + None::<&fn(&[u8]) -> Option>, &mut |_, _, _| Ok((false, None)), &mut |_flags, key_bytes_to_remove, value_bytes_to_remove| { Ok(( @@ -1046,14 +1128,15 @@ mod test { let maybe_walker = maybe_tree.map(|tree| Walker::::new(tree, PanicSource {})); let batch = vec![ - (vec![0], Op::Put(vec![2], BasicMerk)), - (vec![1], Op::Put(vec![2], BasicMerk)), + (vec![0], Op::Put(vec![2], BasicMerkNode)), + (vec![1], Op::Put(vec![2], BasicMerkNode)), ]; let (maybe_tree, key_updates) = Walker::::apply_to( maybe_walker, &batch, PanicSource {}, &|_, _| Ok(0), + None::<&fn(&[u8]) -> Option>, &mut |_, _, _| Ok((false, None)), &mut |_flags, key_bytes_to_remove, value_bytes_to_remove| { Ok(( @@ -1074,15 +1157,16 @@ mod test { #[test] fn insert_updated_multiple() { let batch = vec![ - (vec![0], Op::Put(vec![1], BasicMerk)), - (vec![1], Op::Put(vec![2], BasicMerk)), - (vec![2], Op::Put(vec![3], BasicMerk)), + (vec![0], Op::Put(vec![1], BasicMerkNode)), + (vec![1], Op::Put(vec![2], BasicMerkNode)), + (vec![2], Op::Put(vec![3], BasicMerkNode)), ]; let (maybe_tree, key_updates) = Walker::::apply_to( None, &batch, PanicSource {}, &|_, _| Ok(0), + None::<&fn(&[u8]) -> Option>, &mut |_, _, _| Ok((false, None)), &mut |_flags, key_bytes_to_remove, value_bytes_to_remove| { Ok(( @@ -1098,8 +1182,8 @@ mod test { let maybe_walker = maybe_tree.map(|tree| Walker::::new(tree, PanicSource {})); let batch = vec![ - (vec![0], Op::Put(vec![5], BasicMerk)), - (vec![1], Op::Put(vec![8], BasicMerk)), + (vec![0], Op::Put(vec![5], BasicMerkNode)), + (vec![1], Op::Put(vec![8], BasicMerkNode)), (vec![2], Op::Delete), ]; let (maybe_tree, key_updates) = Walker::::apply_to( @@ -1107,6 +1191,7 @@ mod test { &batch, PanicSource {}, &|_, _| Ok(0), + None::<&fn(&[u8]) -> Option>, &mut |_, _, _| Ok((false, None)), &mut |_flags, key_bytes_to_remove, value_bytes_to_remove| { Ok(( @@ -1127,8 +1212,8 @@ mod test { #[test] fn insert_root_single() { - let tree = TreeNode::new(vec![5], vec![123], None, BasicMerk).unwrap(); - let batch = vec![(vec![6], Op::Put(vec![123], BasicMerk))]; + let tree = TreeNode::new(vec![5], vec![123], None, BasicMerkNode).unwrap(); + let batch = vec![(vec![6], Op::Put(vec![123], BasicMerkNode))]; let tree = apply_memonly(tree, &batch); assert_eq!(tree.key(), &[5]); assert!(tree.child(true).is_none()); @@ -1137,10 +1222,10 @@ mod test { #[test] fn insert_root_double() { - let tree = TreeNode::new(vec![5], vec![123], None, BasicMerk).unwrap(); + let tree = TreeNode::new(vec![5], vec![123], None, BasicMerkNode).unwrap(); let batch = vec![ - (vec![4], Op::Put(vec![123], BasicMerk)), - (vec![6], Op::Put(vec![123], BasicMerk)), + (vec![4], Op::Put(vec![123], BasicMerkNode)), + (vec![6], Op::Put(vec![123], BasicMerkNode)), ]; let tree = apply_memonly(tree, &batch); assert_eq!(tree.key(), &[5]); @@ -1150,12 +1235,12 @@ mod test { #[test] fn insert_rebalance() { - let tree = TreeNode::new(vec![5], vec![123], None, BasicMerk).unwrap(); + let tree = TreeNode::new(vec![5], vec![123], None, BasicMerkNode).unwrap(); - let batch = vec![(vec![6], Op::Put(vec![123], BasicMerk))]; + let batch = vec![(vec![6], Op::Put(vec![123], BasicMerkNode))]; let tree = apply_memonly(tree, &batch); - let batch = vec![(vec![7], Op::Put(vec![123], BasicMerk))]; + let batch = vec![(vec![7], Op::Put(vec![123], BasicMerkNode))]; let tree = apply_memonly(tree, &batch); assert_eq!(tree.key(), &[6]); @@ -1165,10 +1250,10 @@ mod test { #[test] fn insert_100_sequential() { - let mut tree = TreeNode::new(vec![0], vec![123], None, BasicMerk).unwrap(); + let mut tree = TreeNode::new(vec![0], vec![123], None, BasicMerkNode).unwrap(); for i in 0..100 { - let batch = vec![(vec![i + 1], Op::Put(vec![123], BasicMerk))]; + let batch = vec![(vec![i + 1], Op::Put(vec![123], BasicMerkNode))]; tree = apply_memonly(tree, &batch); } diff --git a/merk/src/tree/tree_feature_type.rs b/merk/src/tree/tree_feature_type.rs index e99ca310..c1fceed3 100644 --- a/merk/src/tree/tree_feature_type.rs +++ b/merk/src/tree/tree_feature_type.rs @@ -39,16 +39,16 @@ use ed::{Decode, Encode}; use integer_encoding::{VarInt, VarIntReader, VarIntWriter}; #[cfg(any(feature = "full", feature = "verify"))] -use crate::tree::tree_feature_type::TreeFeatureType::{BasicMerk, SummedMerk}; +use crate::tree::tree_feature_type::TreeFeatureType::{BasicMerkNode, SummedMerkNode}; #[cfg(any(feature = "full", feature = "verify"))] #[derive(Copy, Clone, PartialEq, Eq, Debug)] /// Basic or summed pub enum TreeFeatureType { - /// Basic Merk - BasicMerk, - /// Summed Merk - SummedMerk(i64), + /// Basic Merk Tree Node + BasicMerkNode, + /// Summed Merk Tree Node + SummedMerkNode(i64), } #[cfg(feature = "full")] @@ -57,23 +57,23 @@ impl TreeFeatureType { /// Get length of encoded SummedMerk pub fn sum_length(&self) -> Option { match self { - BasicMerk => None, - SummedMerk(m) => Some(m.encode_var_vec().len() as u32), + BasicMerkNode => None, + SummedMerkNode(m) => Some(m.encode_var_vec().len() as u32), } } #[inline] /// Is sum feature? pub fn is_sum_feature(&self) -> bool { - matches!(self, SummedMerk(_)) + matches!(self, SummedMerkNode(_)) } #[inline] /// Get encoding cost of self pub(crate) fn encoding_cost(&self) -> usize { match self { - BasicMerk => 1, - SummedMerk(_sum) => 9, + BasicMerkNode => 1, + SummedMerkNode(_sum) => 9, } } } @@ -85,11 +85,11 @@ impl Encode for TreeFeatureType { #[inline] fn encode_into(&self, dest: &mut W) -> ed::Result<()> { match self { - BasicMerk => { + BasicMerkNode => { dest.write_all(&[0])?; Ok(()) } - SummedMerk(sum) => { + SummedMerkNode(sum) => { dest.write_all(&[1])?; dest.write_varint(sum.to_owned())?; Ok(()) @@ -100,8 +100,8 @@ impl Encode for TreeFeatureType { #[inline] fn encoding_length(&self) -> ed::Result { match self { - BasicMerk => Ok(1), - SummedMerk(sum) => { + BasicMerkNode => Ok(1), + SummedMerkNode(sum) => { let encoded_sum = sum.encode_var_vec(); // 1 for the enum type // encoded_sum.len() for the length of the encoded vector @@ -118,10 +118,10 @@ impl Decode for TreeFeatureType { let mut feature_type: [u8; 1] = [0]; input.read_exact(&mut feature_type)?; match feature_type { - [0] => Ok(BasicMerk), + [0] => Ok(BasicMerkNode), [1] => { let encoded_sum: i64 = input.read_varint()?; - Ok(SummedMerk(encoded_sum)) + Ok(SummedMerkNode(encoded_sum)) } _ => Err(ed::Error::UnexpectedByte(55)), } diff --git a/merk/src/tree/walk/fetch.rs b/merk/src/tree/walk/fetch.rs index a13be66c..08b66d99 100644 --- a/merk/src/tree/walk/fetch.rs +++ b/merk/src/tree/walk/fetch.rs @@ -35,6 +35,8 @@ use grovedb_costs::CostResult; use super::super::{Link, TreeNode}; #[cfg(feature = "full")] use crate::error::Error; +#[cfg(feature = "full")] +use crate::tree::kv::ValueDefinedCostType; #[cfg(feature = "full")] /// A source of data to be used by the tree when encountering a pruned node. @@ -43,5 +45,9 @@ use crate::error::Error; pub trait Fetch { /// Called when the tree needs to fetch a node with the given `Link`. The /// `link` value will always be a `Link::Reference` variant. - fn fetch(&self, link: &Link) -> CostResult; + fn fetch( + &self, + link: &Link, + value_defined_cost_fn: Option<&impl Fn(&[u8]) -> Option>, + ) -> CostResult; } diff --git a/merk/src/tree/walk/mod.rs b/merk/src/tree/walk/mod.rs index d4a6a537..35e476f3 100644 --- a/merk/src/tree/walk/mod.rs +++ b/merk/src/tree/walk/mod.rs @@ -77,7 +77,14 @@ where /// Similar to `Tree#detach`, but yields a `Walker` which fetches from the /// same source as `self`. Returned tuple is `(updated_self, /// maybe_child_walker)`. - pub fn detach(mut self, left: bool) -> CostResult<(Self, Option), Error> { + pub fn detach( + mut self, + left: bool, + value_defined_cost_fn: Option<&V>, + ) -> CostResult<(Self, Option), Error> + where + V: Fn(&[u8]) -> Option, + { let mut cost = OperationCost::default(); let link = match self.tree.link(left) { @@ -96,7 +103,10 @@ where Some(Link::Reference { .. }) => (), _ => unreachable!("Expected Some(Link::Reference)"), } - cost_return_on_error!(&mut cost, self.source.fetch(&link.unwrap())) + cost_return_on_error!( + &mut cost, + self.source.fetch(&link.unwrap(), value_defined_cost_fn) + ) }; let child = self.wrap(child); @@ -106,29 +116,44 @@ where /// Similar to `Tree#detach_expect`, but yields a `Walker` which fetches /// from the same source as `self`. Returned tuple is `(updated_self, /// child_walker)`. - pub fn detach_expect(self, left: bool) -> CostResult<(Self, Self), Error> { - self.detach(left).map_ok(|(walker, maybe_child)| { - if let Some(child) = maybe_child { - (walker, child) - } else { - panic!( - "Expected {} child, got None", - if left { "left" } else { "right" } - ); - } - }) + pub fn detach_expect( + self, + left: bool, + value_defined_cost_fn: Option<&V>, + ) -> CostResult<(Self, Self), Error> + where + V: Fn(&[u8]) -> Option, + { + self.detach(left, value_defined_cost_fn) + .map_ok(|(walker, maybe_child)| { + if let Some(child) = maybe_child { + (walker, child) + } else { + panic!( + "Expected {} child, got None", + if left { "left" } else { "right" } + ); + } + }) } /// Similar to `Tree#walk`, but yields a `Walker` which fetches from the /// same source as `self`. - pub fn walk(self, left: bool, f: F) -> CostResult + pub fn walk( + self, + left: bool, + f: F, + value_defined_cost_fn: Option<&V>, + ) -> CostResult where F: FnOnce(Option) -> CostResult, Error>, T: Into, + V: Fn(&[u8]) -> Option, { let mut cost = OperationCost::default(); - let (mut walker, maybe_child) = cost_return_on_error!(&mut cost, self.detach(left)); + let (mut walker, maybe_child) = + cost_return_on_error!(&mut cost, self.detach(left, value_defined_cost_fn)); let new_child = match f(maybe_child).unwrap_add_cost(&mut cost) { Ok(x) => x.map(|t| t.into()), Err(e) => return Err(e).wrap_with_cost(cost), @@ -139,14 +164,21 @@ where /// Similar to `Tree#walk_expect` but yields a `Walker` which fetches from /// the same source as `self`. - pub fn walk_expect(self, left: bool, f: F) -> CostResult + pub fn walk_expect( + self, + left: bool, + f: F, + value_defined_cost_fn: Option<&V>, + ) -> CostResult where F: FnOnce(Self) -> CostResult, Error>, T: Into, + V: Fn(&[u8]) -> Option, { let mut cost = OperationCost::default(); - let (mut walker, child) = cost_return_on_error!(&mut cost, self.detach_expect(left)); + let (mut walker, child) = + cost_return_on_error!(&mut cost, self.detach_expect(left, value_defined_cost_fn)); let new_child = match f(child).unwrap_add_cost(&mut cost) { Ok(x) => x.map(|t| t.into()), Err(e) => return Err(e).wrap_with_cost(cost), @@ -368,34 +400,42 @@ mod test { use grovedb_costs::{storage_cost::removal::StorageRemovedBytes::NoStorageRemoval, CostsExt}; use super::{super::NoopCommit, *}; - use crate::tree::{TreeFeatureType::BasicMerk, TreeNode}; + use crate::tree::{TreeFeatureType::BasicMerkNode, TreeNode}; #[derive(Clone)] struct MockSource {} impl Fetch for MockSource { - fn fetch(&self, link: &Link) -> CostResult { - TreeNode::new(link.key().to_vec(), b"foo".to_vec(), None, BasicMerk).map(Ok) + fn fetch( + &self, + link: &Link, + _value_defined_cost_fn: Option<&impl Fn(&[u8]) -> Option>, + ) -> CostResult { + TreeNode::new(link.key().to_vec(), b"foo".to_vec(), None, BasicMerkNode).map(Ok) } } #[test] fn walk_modified() { - let tree = TreeNode::new(b"test".to_vec(), b"abc".to_vec(), None, BasicMerk) + let tree = TreeNode::new(b"test".to_vec(), b"abc".to_vec(), None, BasicMerkNode) .unwrap() .attach( true, - Some(TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerk).unwrap()), + Some(TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerkNode).unwrap()), ); let source = MockSource {}; let walker = Walker::new(tree, source); let walker = walker - .walk(true, |child| -> CostResult, Error> { - assert_eq!(child.expect("should have child").tree().key(), b"foo"); - Ok(None).wrap_with_cost(Default::default()) - }) + .walk( + true, + |child| -> CostResult, Error> { + assert_eq!(child.expect("should have child").tree().key(), b"foo"); + Ok(None).wrap_with_cost(Default::default()) + }, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("walk failed"); assert!(walker.into_inner().child(true).is_none()); @@ -403,11 +443,11 @@ mod test { #[test] fn walk_stored() { - let mut tree = TreeNode::new(b"test".to_vec(), b"abc".to_vec(), None, BasicMerk) + let mut tree = TreeNode::new(b"test".to_vec(), b"abc".to_vec(), None, BasicMerkNode) .unwrap() .attach( true, - Some(TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerk).unwrap()), + Some(TreeNode::new(b"foo".to_vec(), b"bar".to_vec(), None, BasicMerkNode).unwrap()), ); tree.commit(&mut NoopCommit {}, &|_, _| Ok(0)) .unwrap() @@ -417,10 +457,14 @@ mod test { let walker = Walker::new(tree, source); let walker = walker - .walk(true, |child| -> CostResult, Error> { - assert_eq!(child.expect("should have child").tree().key(), b"foo"); - Ok(None).wrap_with_cost(Default::default()) - }) + .walk( + true, + |child| -> CostResult, Error> { + assert_eq!(child.expect("should have child").tree().key(), b"foo"); + Ok(None).wrap_with_cost(Default::default()) + }, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("walk failed"); assert!(walker.into_inner().child(true).is_none()); @@ -439,7 +483,7 @@ mod test { sum: None, }), None, - BasicMerk, + BasicMerkNode, ) .unwrap(); @@ -447,10 +491,14 @@ mod test { let walker = Walker::new(tree, source); let walker = walker - .walk_expect(true, |child| -> CostResult, Error> { - assert_eq!(child.tree().key(), b"foo"); - Ok(None).wrap_with_cost(Default::default()) - }) + .walk_expect( + true, + |child| -> CostResult, Error> { + assert_eq!(child.tree().key(), b"foo"); + Ok(None).wrap_with_cost(Default::default()) + }, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("walk failed"); assert!(walker.into_inner().child(true).is_none()); @@ -458,16 +506,20 @@ mod test { #[test] fn walk_none() { - let tree = TreeNode::new(b"test".to_vec(), b"abc".to_vec(), None, BasicMerk).unwrap(); + let tree = TreeNode::new(b"test".to_vec(), b"abc".to_vec(), None, BasicMerkNode).unwrap(); let source = MockSource {}; let walker = Walker::new(tree, source); walker - .walk(true, |child| -> CostResult, Error> { - assert!(child.is_none()); - Ok(None).wrap_with_cost(Default::default()) - }) + .walk( + true, + |child| -> CostResult, Error> { + assert!(child.is_none()); + Ok(None).wrap_with_cost(Default::default()) + }, + None::<&fn(&[u8]) -> Option>, + ) .unwrap() .expect("walk failed"); } diff --git a/merk/src/tree/walk/ref_walker.rs b/merk/src/tree/walk/ref_walker.rs index 26c23b0a..d9fb1bcd 100644 --- a/merk/src/tree/walk/ref_walker.rs +++ b/merk/src/tree/walk/ref_walker.rs @@ -36,6 +36,7 @@ use super::{ super::{Link, TreeNode}, Fetch, }; +use crate::tree::kv::ValueDefinedCostType; #[cfg(feature = "full")] use crate::Error; @@ -73,7 +74,14 @@ where /// Traverses to the child on the given side (if any), fetching from the /// source if pruned. When fetching, the link is upgraded from /// `Link::Reference` to `Link::Loaded`. - pub fn walk(&mut self, left: bool) -> CostResult>, Error> { + pub fn walk( + &mut self, + left: bool, + value_defined_cost_fn: Option<&V>, + ) -> CostResult>, Error> + where + V: Fn(&[u8]) -> Option, + { let link = match self.tree.link(left) { None => return Ok(None).wrap_with_cost(Default::default()), Some(link) => link, @@ -84,7 +92,7 @@ where Link::Reference { .. } => { let load_res = self .tree - .load(left, &self.source) + .load(left, &self.source, value_defined_cost_fn) .unwrap_add_cost(&mut cost); if let Err(e) = load_res { return Err(e).wrap_with_cost(cost);