Fix/chunk lifetimes

grovedb/src/lib.rs

@@ -199,6 +199,7 @@ use grovedb_merk::{
     tree::{combine_hash, value_hash},
     BatchEntry, CryptoHash, KVIterator, Merk,
 };
+use grovedb_merk::ChunkProducer;
 use grovedb_path::SubtreePath;
 #[cfg(feature = "full")]
 use grovedb_storage::rocksdb_storage::PrefixedRocksDbImmediateStorageContext;
@@ -223,6 +224,7 @@ pub use crate::error::Error;
 #[cfg(feature = "full")]
 use crate::util::{root_merk_optional_tx, storage_context_optional_tx};
 use crate::Error::MerkError;
+use crate::replication::util_encode_vec_ops;
 
 #[cfg(feature = "full")]
 type Hash = [u8; 32];
@@ -317,13 +319,48 @@ impl GroveDb {
         }
     }
 
+    fn open_transactional_merk_by_prefix<'db>(
+        &'db self,
+        prefix: SubtreePrefix,
+        root_key: Option<Vec<u8>>,
+        is_sum_tree: bool,
+        tx: &'db Transaction,
+        batch: Option<&'db StorageBatch>,
+    ) -> CostResult<Merk<PrefixedRocksDbTransactionContext>, Error>
+    {
+        let mut cost = OperationCost::default();
+        let storage = self
+            .db
+            .get_transactional_storage_context_by_subtree_prefix(prefix, batch, tx)
+            .unwrap_add_cost(&mut cost);
+        if root_key.is_some() {
+            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 by prefix with given root key".to_owned())
+            }).add_cost(cost)
+        }
+        else {
+            Merk::open_base(
+                storage,
+                false,
+                Some(&Element::value_defined_cost_for_serialized_value),
+            ).map_err(|_| {
+                Error::CorruptedData("cannot open a root subtree by prefix".to_owned())
+            }).add_cost(cost)
+        }
+    }
+
     /// Opens a Merk at given path for with direct write access. Intended for
     /// replication purposes.
-    fn open_merk_for_replication<'db, 'b, B>(
+    fn open_merk_for_replication<'tx, 'db: 'tx, 'b, B>(
         &'db self,
         path: SubtreePath<'b, B>,
-        tx: &'db Transaction,
-    ) -> Result<Merk<PrefixedRocksDbImmediateStorageContext<'db>>, Error>
+        tx: &'tx Transaction<'db>,
+    ) -> Result<(Merk<PrefixedRocksDbImmediateStorageContext<'tx>>, Option<Vec<u8>>, bool), Error>
     where
         B: AsRef<[u8]> + 'b,
     {
@@ -350,29 +387,37 @@ impl GroveDb {
                 .unwrap()?;
             let is_sum_tree = element.is_sum_tree();
             if let Element::Tree(root_key, _) | Element::SumTree(root_key, ..) = element {
+                Ok((
                 Merk::open_layered_with_root_key(
                     storage,
-                    root_key,
+                    root_key.clone(),
                     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()
+                .unwrap()?,
+                    root_key,
+                    is_sum_tree
+                ))
             } else {
                 Err(Error::CorruptedPath(
                     "cannot open a subtree as parent exists but is not a tree",
                 ))
             }
         } else {
+            Ok((
             Merk::open_base(
                 storage,
                 false,
                 None::<&fn(&[u8]) -> Option<ValueDefinedCostType>>,
             )
             .map_err(|_| Error::CorruptedData("cannot open a the root subtree".to_owned()))
-            .unwrap()
+            .unwrap()?,
+                None,
+                false
+            ))
         }
     }
 
@@ -437,6 +482,40 @@ impl GroveDb {
         }
     }
 
+    fn open_non_transactional_merk_by_prefix<'db>(
+        &'db self,
+        prefix: SubtreePrefix,
+        root_key: Option<Vec<u8>>,
+        is_sum_tree: bool,
+        batch: Option<&'db StorageBatch>,
+    ) -> CostResult<Merk<PrefixedRocksDbStorageContext>, Error>
+    {
+        let mut cost = OperationCost::default();
+        let storage = self
+            .db
+            .get_storage_context_by_subtree_prefix(prefix, batch)
+            .unwrap_add_cost(&mut cost);
+        if root_key.is_some() {
+            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 by prefix with given root key".to_owned())
+            }).add_cost(cost)
+        }
+        else {
+            Merk::open_base(
+                storage,
+                false,
+                Some(&Element::value_defined_cost_for_serialized_value),
+            ).map_err(|_| {
+                Error::CorruptedData("cannot open a root subtree by prefix".to_owned())
+            }).add_cost(cost)
+        }
+    }
+
     /// Creates a checkpoint
     pub fn create_checkpoint<P: AsRef<Path>>(&self, path: P) -> Result<(), Error> {
         self.db.create_checkpoint(path).map_err(|e| e.into())