fixes and tests

grovedb-version/src/version/v2.rs

@@ -13,7 +13,7 @@ use crate::version::{
 };
 
 pub const GROVE_V2: GroveVersion = GroveVersion {
-    protocol_version: 0,
+    protocol_version: 1,
     grovedb_versions: GroveDBVersions {
         apply_batch: GroveDBApplyBatchVersions {
             apply_batch_structure: 0,
@@ -186,7 +186,7 @@ pub const GROVE_V2: GroveVersion = GroveVersion {
     },
     merk_versions: MerkVersions {
         average_case_costs: MerkAverageCaseCostsVersions {
-            add_average_case_merk_propagate: 1,
+            add_average_case_merk_propagate: 1, // changed
         },
     },
 };

grovedb/Cargo.toml

@@ -45,6 +45,7 @@ criterion = "0.5.1"
 hex = "0.4.3"
 pretty_assertions = "1.4.0"
 rand = "0.8.5"
+assert_matches = "1.5.0"
 
 [[bench]]
 name = "insertion_benchmark"

grovedb/src/batch/mod.rs

@@ -1864,6 +1864,28 @@ impl GroveDb {
                                                                 aggregate_data,
                                                             }
                                                             .into();
+                                                    } else if let Element::BigSumTree(.., flags) =
+                                                        element
+                                                    {
+                                                        *mutable_occupied_entry =
+                                                            GroveOp::InsertTreeWithRootHash {
+                                                                hash: root_hash,
+                                                                root_key: calculated_root_key,
+                                                                flags: flags.clone(),
+                                                                aggregate_data,
+                                                            }
+                                                            .into();
+                                                    } else if let Element::CountTree(.., flags) =
+                                                        element
+                                                    {
+                                                        *mutable_occupied_entry =
+                                                            GroveOp::InsertTreeWithRootHash {
+                                                                hash: root_hash,
+                                                                root_key: calculated_root_key,
+                                                                flags: flags.clone(),
+                                                                aggregate_data,
+                                                            }
+                                                            .into();
                                                     } else {
                                                         return Err(Error::InvalidBatchOperation(
                                                             "insertion of element under a non tree",

grovedb/src/element/constructor.rs

@@ -29,6 +29,18 @@ impl Element {
         Element::new_sum_tree(Default::default())
     }
 
+    #[cfg(feature = "full")]
+    /// Set element to default empty big sum tree without flags
+    pub fn empty_big_sum_tree() -> Self {
+        Element::new_big_sum_tree(Default::default())
+    }
+
+    #[cfg(feature = "full")]
+    /// Set element to default empty count tree without flags
+    pub fn empty_count_tree() -> Self {
+        Element::new_count_tree(Default::default())
+    }
+
     #[cfg(feature = "full")]
     /// Set element to default empty sum tree with flags
     pub fn empty_sum_tree_with_flags(flags: Option<ElementFlags>) -> Self {

grovedb/src/element/helpers.rs

@@ -21,7 +21,7 @@ use grovedb_version::{check_grovedb_v0, error::GroveVersionError, version::Grove
 #[cfg(feature = "full")]
 use integer_encoding::VarInt;
 
-use crate::element::BIG_SUM_TREE_COST_SIZE;
+use crate::element::{BIG_SUM_TREE_COST_SIZE, COUNT_TREE_COST_SIZE};
 #[cfg(feature = "full")]
 use crate::reference_path::path_from_reference_path_type;
 #[cfg(any(feature = "full", feature = "verify"))]
@@ -45,6 +45,16 @@ impl Element {
         }
     }
 
+    #[cfg(any(feature = "full", feature = "verify"))]
+    /// Decoded the integer value in the CountTree element type, returns 1 for
+    /// everything else
+    pub fn count_value_or_default(&self) -> u64 {
+        match self {
+            Element::CountTree(_, count_value, _) => *count_value,
+            _ => 1,
+        }
+    }
+
     #[cfg(any(feature = "full", feature = "verify"))]
     /// Decoded the integer value in the SumItem element type, returns 0 for
     /// everything else
@@ -232,7 +242,7 @@ impl Element {
             TreeType::NormalTree => Ok(BasicMerkNode),
             TreeType::SumTree => Ok(SummedMerkNode(self.sum_value_or_default())),
             TreeType::BigSumTree => Ok(BigSummedMerkNode(self.big_sum_value_or_default())),
-            TreeType::CountTree => Ok(CountedMerkNode(1)),
+            TreeType::CountTree => Ok(CountedMerkNode(self.count_value_or_default())),
         }
     }
 
@@ -388,6 +398,17 @@ impl Element {
                     key_len, value_len, node_type,
                 )
             }
+            Element::CountTree(_, _count_value, flags) => {
+                let flags_len = flags.map_or(0, |flags| {
+                    let flags_len = flags.len() as u32;
+                    flags_len + flags_len.required_space() as u32
+                });
+                let value_len = COUNT_TREE_COST_SIZE + flags_len;
+                let key_len = key.len() as u32;
+                KV::layered_value_byte_cost_size_for_key_and_value_lengths(
+                    key_len, value_len, node_type,
+                )
+            }
             Element::SumItem(.., flags) => {
                 let flags_len = flags.map_or(0, |flags| {
                     let flags_len = flags.len() as u32;
@@ -414,6 +435,7 @@ impl Element {
             Element::SumTree(..) => Ok(SUM_TREE_COST_SIZE),
             Element::BigSumTree(..) => Ok(BIG_SUM_TREE_COST_SIZE),
             Element::SumItem(..) => Ok(SUM_ITEM_COST_SIZE),
+            Element::CountTree(..) => Ok(COUNT_TREE_COST_SIZE),
             _ => Err(Error::CorruptedCodeExecution(
                 "trying to get tree cost from non tree element",
             )),
@@ -436,6 +458,7 @@ impl Element {
             Element::Tree(..) => Some(LayeredValueDefinedCost(cost)),
             Element::SumTree(..) => Some(LayeredValueDefinedCost(cost)),
             Element::BigSumTree(..) => Some(LayeredValueDefinedCost(cost)),
+            Element::CountTree(..) => Some(LayeredValueDefinedCost(cost)),
             Element::SumItem(..) => Some(SpecializedValueDefinedCost(cost)),
             _ => None,
         }

grovedb/src/lib.rs

@@ -501,7 +501,8 @@ impl GroveDb {
             let tree_type = element.tree_type();
             if let Element::Tree(root_key, _)
             | Element::SumTree(root_key, ..)
-            | Element::BigSumTree(root_key, ..) = element
+            | Element::BigSumTree(root_key, ..)
+            | Element::CountTree(root_key, ..) = element
             {
                 let tree_type = tree_type.expect("expected tree type");
                 Merk::open_layered_with_root_key(
@@ -715,7 +716,7 @@ impl GroveDb {
                     grove_version
                 )
             );
-            let (root_hash, root_key, sum) = cost_return_on_error!(
+            let (root_hash, root_key, aggregate_data) = cost_return_on_error!(
                 &mut cost,
                 child_tree
                     .root_hash_key_and_aggregate_data()
@@ -728,7 +729,7 @@ impl GroveDb {
                     parent_key,
                     root_key,
                     root_hash,
-                    sum,
+                    aggregate_data,
                     grove_version,
                 )
             );
@@ -832,6 +833,19 @@ impl GroveDb {
                     None,
                     grove_version,
                 )
+            } else if let Element::CountTree(.., flag) = element {
+                let tree = Element::new_count_tree_with_flags_and_count_value(
+                    maybe_root_key,
+                    aggregate_data.as_u64(),
+                    flag,
+                );
+                tree.insert_subtree(
+                    parent_tree,
+                    key.as_ref(),
+                    root_tree_hash,
+                    None,
+                    grove_version,
+                )
             } else {
                 Err(Error::InvalidPath(
                     "can only propagate on tree items".to_owned(),
@@ -912,6 +926,25 @@ impl GroveDb {
                         merk_feature_type,
                         grove_version,
                     )
+                } else if let Element::CountTree(.., flag) = element {
+                    let tree = Element::new_count_tree_with_flags_and_count_value(
+                        maybe_root_key,
+                        aggregate_data.as_u64(),
+                        flag,
+                    );
+                    let merk_feature_type = cost_return_on_error!(
+                        &mut cost,
+                        tree.get_feature_type(parent_tree.tree_type)
+                            .wrap_with_cost(OperationCost::default())
+                    );
+                    tree.insert_subtree_into_batch_operations(
+                        key,
+                        root_tree_hash,
+                        true,
+                        batch_operations,
+                        merk_feature_type,
+                        grove_version,
+                    )
                 } else {
                     Err(Error::InvalidPath(
                         "can only propagate on tree items".to_owned(),

grovedb/src/operations/get/mod.rs

@@ -455,7 +455,10 @@ impl GroveDb {
             }
             .unwrap_add_cost(&mut cost);
             match element {
-                Ok(Element::Tree(..)) | Ok(Element::SumTree(..)) => Ok(()).wrap_with_cost(cost),
+                Ok(Element::Tree(..))
+                | Ok(Element::SumTree(..))
+                | Ok(Element::BigSumTree(..))
+                | Ok(Element::CountTree(..)) => Ok(()).wrap_with_cost(cost),
                 Ok(_) | Err(Error::PathKeyNotFound(_)) => Err(error_fn()).wrap_with_cost(cost),
                 Err(e) => Err(e).wrap_with_cost(cost),
             }

grovedb/src/operations/insert/mod.rs

@@ -317,7 +317,10 @@ impl GroveDb {
                     )
                 );
             }
-            Element::Tree(ref value, _) | Element::SumTree(ref value, ..) => {
+            Element::Tree(ref value, _)
+            | Element::SumTree(ref value, ..)
+            | Element::BigSumTree(ref value, ..)
+            | Element::CountTree(ref value, ..) => {
                 if value.is_some() {
                     return Err(Error::InvalidCodeExecution(
                         "a tree should be empty at the moment of insertion when not using batches",
@@ -450,7 +453,10 @@ impl GroveDb {
                     )
                 );
             }
-            Element::Tree(ref value, _) | Element::SumTree(ref value, ..) => {
+            Element::Tree(ref value, _)
+            | Element::SumTree(ref value, ..)
+            | Element::BigSumTree(ref value, ..)
+            | Element::CountTree(ref value, ..) => {
                 if value.is_some() {
                     return Err(Error::InvalidCodeExecution(
                         "a tree should be empty at the moment of insertion when not using batches",
@@ -1593,6 +1599,87 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_one_insert_item_cost_under_count_tree() {
+        let grove_version = GroveVersion::latest();
+        let db = make_empty_grovedb();
+        let tx = db.start_transaction();
+
+        db.insert(
+            EMPTY_PATH,
+            b"tree",
+            Element::empty_count_tree(),
+            None,
+            Some(&tx),
+            grove_version,
+        )
+        .unwrap()
+        .unwrap();
+
+        let cost = db
+            .insert(
+                [b"tree".as_slice()].as_ref(),
+                b"key1",
+                Element::new_item(b"test".to_vec()),
+                None,
+                Some(&tx),
+                grove_version,
+            )
+            .cost_as_result()
+            .unwrap();
+
+        // Explanation for 152 storage_written_bytes
+
+        // Key -> 37 bytes
+        // 32 bytes for the key prefix
+        // 4 bytes for the key
+        // 1 byte for key_size (required space for 36)
+
+        // Value -> 81
+        //   1 for the enum type item
+        //   1 for size of test bytes
+        //   4 for test bytes
+        //   1 for the flag option (but no flags)
+        // 32 for node hash
+        // 32 for value hash (trees have this for free)
+        // 9 for Count node
+        // 1 byte for the value_size (required space for 1)
+
+        // Parent Hook -> 48
+        // Key Bytes 4
+        // Hash Size 32
+        // Key Length 1
+        // Count Merk 9
+        // Child Heights 2
+
+        // Total 37 + 81 + 48 = 166
+
+        // Explanation for replaced bytes
+
+        // Replaced parent Value -> 86
+        //   1 for the flag option (but no flags)
+        //   1 for the enum type
+        //   1 for an empty option
+        //   1 for the count merk
+        //   9 for the count
+        // 32 for node hash
+        // 40 for the parent hook
+        // 2 byte for the value_size
+        assert_eq!(
+            cost,
+            OperationCost {
+                seek_count: 5, // todo: verify this
+                storage_cost: StorageCost {
+                    added_bytes: 166,
+                    replaced_bytes: 87,
+                    removed_bytes: NoStorageRemoval
+                },
+                storage_loaded_bytes: 162, // todo: verify this
+                hash_node_calls: 8,        // todo: verify this
+            }
+        );
+    }
+
     #[test]
     fn test_one_insert_item_with_apple_flags_cost() {
         let grove_version = GroveVersion::latest();

grovedb/src/operations/proof/generate.rs

@@ -314,7 +314,10 @@ impl GroveDb {
                                 }
                                 has_a_result_at_level |= true;
                             }
-                            Ok(Element::Tree(Some(_), _)) | Ok(Element::SumTree(Some(_), ..))
+                            Ok(Element::Tree(Some(_), _))
+                            | Ok(Element::SumTree(Some(_), ..))
+                            | Ok(Element::BigSumTree(Some(_), ..))
+                            | Ok(Element::CountTree(Some(_), ..))
                                 if !done_with_results
                                     && query.has_subquery_or_matching_in_path_on_key(key) =>
                             {