wip

grovedb/src/merk_cache.rs

@@ -17,6 +17,11 @@ use crate::{Element, Error, GroveDb, Transaction};
 
 type TxMerk<'db> = Merk<PrefixedRocksDbTransactionContext<'db>>;
 
+struct CachedMerk<'db> {
+    updated: bool,
+    merk: TxMerk<'db>,
+}
+
 /// Merk caching structure.
 ///
 /// Since we usually postpone all writes to the very end with a single RocksDB
@@ -29,7 +34,7 @@ pub(crate) struct MerkCache<'db, 'b, B> {
     tx: &'db Transaction<'db>,
     batch: &'db StorageBatch,
     version: &'db GroveVersion,
-    merks: BTreeMap<SubtreePath<'b, B>, TxMerk<'db>>,
+    merks: BTreeMap<SubtreePath<'b, B>, CachedMerk<'db>>,
 }
 
 impl<'db, 'b, B: AsRef<[u8]>> MerkCache<'db, 'b, B> {
@@ -54,7 +59,7 @@ impl<'db, 'b, B: AsRef<[u8]>> MerkCache<'db, 'b, B> {
     fn get_merk_mut_internal<'s>(
         &'s mut self,
         path: SubtreePath<'b, B>,
-    ) -> CostResult<&'s mut TxMerk<'db>, Error> {
+    ) -> CostResult<&'s mut CachedMerk<'db>, Error> {
         let mut cost = Default::default();
 
         match self.merks.entry(path) {
@@ -69,7 +74,11 @@ impl<'db, 'b, B: AsRef<[u8]>> MerkCache<'db, 'b, B> {
                         self.version
                     )
                 );
-                Ok(e.insert(merk)).wrap_with_cost(cost)
+                Ok(e.insert(CachedMerk {
+                    merk,
+                    updated: false,
+                }))
+                .wrap_with_cost(cost)
             }
         }
     }
@@ -84,8 +93,8 @@ impl<'db, 'b, B: AsRef<[u8]>> MerkCache<'db, 'b, B> {
     pub(crate) fn get_multi_mut<'s, const N: usize>(
         &'s mut self,
         paths: [SubtreePath<'b, B>; N],
-    ) -> CostResult<[MerkHandle<'db, 's>; N], Error> {
-        let mut result_uninit = [const { MaybeUninit::<MerkHandle<'db, 's>>::uninit() }; N];
+    ) -> CostResult<[MerkHandle<'db, 's, 'b, B>; N], Error> {
+        let mut result_uninit = [const { MaybeUninit::<MerkHandle<'db, 's, 'b, B>>::uninit() }; N];
         let mut cost = Default::default();
 
         let unique_args: HashSet<_> = paths.iter().collect();
@@ -100,13 +109,14 @@ impl<'db, 'b, B: AsRef<[u8]>> MerkCache<'db, 'b, B> {
             // mutable references. The mandatory keys' uniqueness check above makes
             // sure no overlapping memory will be referenced.
             let merk_ref = unsafe {
-                MerkHandle {
-                    merk: (cost_return_on_error!(&mut cost, self.get_merk_mut_internal(path))
-                        as *mut TxMerk<'db>)
+                MerkHandle::new(
+                    path.clone(),
+                    (cost_return_on_error!(&mut cost, self.get_merk_mut_internal(path))
+                        as *mut CachedMerk<'db>)
                         .as_mut::<'s>()
                         .expect("not a null pointer"),
-                    version: &self.version,
-                }
+                    &self.version,
+                )
             };
             result_uninit[i].write(merk_ref);
         }
@@ -116,30 +126,32 @@ impl<'db, 'b, B: AsRef<[u8]>> MerkCache<'db, 'b, B> {
         // N in our case. `mem::transmute` would represent it better, however,
         // due to poor support of const generics in stable Rust we bypass
         // compile-time size checks with pointer casts.
-        let result = unsafe { (&result_uninit as *const _ as *const [MerkHandle; N]).read() };
+        let result = unsafe { (&result_uninit as *const _ as *const [MerkHandle<B>; N]).read() };
         mem::forget(result_uninit);
 
         Ok(result).wrap_with_cost(cost)
     }
 }
 
 /// Handle to a cached Merk.
-pub(crate) struct MerkHandle<'db, 'c> {
+pub(crate) struct MerkHandle<'db, 'c, 'b, B> {
+    path: SubtreePath<'b, B>,
     merk: &'c mut TxMerk<'db>,
     version: &'db GroveVersion,
+    updated: &'c mut bool,
 }
 
 /// It is allowed to dereference `MerkHandle` to regular Merks but in a
 /// non-mutable way since we want to track what have been done to those Merks.
-impl<'db, 'c> Deref for MerkHandle<'db, 'c> {
+impl<'db, 'c, 'b, B> Deref for MerkHandle<'db, 'c, 'b, B> {
     type Target = TxMerk<'db>;
 
     fn deref(&self) -> &Self::Target {
         &self.merk
     }
 }
 
-impl<'db, 'c> MerkHandle<'db, 'c> {
+impl<'db, 'c, 'b, B> MerkHandle<'db, 'c, 'b, B> {
     pub(crate) fn insert(
         &mut self,
         key: impl AsRef<[u8]>,
@@ -148,7 +160,20 @@ impl<'db, 'c> MerkHandle<'db, 'c> {
     ) -> CostResult<(), Error> {
         element
             .insert(self.merk, key, options, self.version)
-            .for_ok(|_| todo!())
+            .for_ok(|_| *self.updated = true)
+    }
+
+    fn new(
+        path: SubtreePath<'b, B>,
+        cached_merk: &'c mut CachedMerk<'db>,
+        version: &'db GroveVersion,
+    ) -> Self {
+        Self {
+            path,
+            merk: &mut cached_merk.merk,
+            version,
+            updated: &mut cached_merk.updated,
+        }
     }
 }