fix test

grovedb/src/merk_cache.rs

@@ -3,19 +3,15 @@
 use std::{
     cell::{Cell, UnsafeCell},
     collections::{btree_map::Entry, BTreeMap},
-    ops::{Deref, DerefMut},
 };
 
-use grovedb_costs::{cost_return_on_error, cost_return_on_error_no_add, CostResult, CostsExt};
+use grovedb_costs::{cost_return_on_error, CostResult, CostsExt};
 use grovedb_merk::Merk;
-use grovedb_path::{SubtreePath, SubtreePathBuilder};
-use grovedb_storage::{
-    rocksdb_storage::{PrefixedRocksDbTransactionContext, RocksDbStorage},
-    StorageBatch,
-};
+use grovedb_path::SubtreePathBuilder;
+use grovedb_storage::{rocksdb_storage::PrefixedRocksDbTransactionContext, StorageBatch};
 use grovedb_version::version::GroveVersion;
 
-use crate::{Element, Error, GroveDb, Transaction};
+use crate::{Error, GroveDb, Transaction};
 
 type TxMerk<'db> = Merk<PrefixedRocksDbTransactionContext<'db>>;
 
@@ -52,6 +48,10 @@ impl<'db, 'b, B: AsRef<[u8]>> MerkCache<'db, 'b, B> {
         let mut cost = Default::default();
 
         // SAFETY: there are no other references to `merks` memory at the same time.
+        // Note while it's possible to have direct references to actual Merk trees,
+        // outside of the scope of this function, this map (`merks`) has
+        // indirect connection to them through `Box`, thus there are no overlapping
+        // references, and that is requirement of `UnsafeCell` we have there.
         let boxed_flag_merk = match unsafe {
             self.merks
                 .get()
@@ -91,27 +91,27 @@ impl<'db, 'b, B: AsRef<[u8]>> MerkCache<'db, 'b, B> {
         // references plus borrowing rules aren't violated (one `&mut` or many
         // `&` with no `&mut` at a time).
         //
-        // To make sure changes to the internal structure won't affect existing borrows
-        // we have an indirection in a form of `Box`, that allows us to move and update
+        // To make sure changes to the map won't affect existing borrows we have an
+        // indirection in a form of `Box`, that allows us to move and update
         // `MerkCache` with new subtrees and possible reallocations without breaking
         // `MerkHandle`'s references. We use `UnsafeCell` to connect lifetimes and check
         // in compile time that `MerkHandle`s won't outlive the cache, even though we
         // don't hold any references to it, but `&mut` reference would make this borrow
         // exclusive for the whole time of `MerkHandle`, so it shall go intially through
         // a shared reference.
         //
-        // Borrowing rules are covered slightly more complicated way:
-        // 1. Of a pair behind heap allocation only Merk is uniquely borrowed by
-        //    `MerkHandle`
-        // 2. Borrow flag is referenced by `MerkHandle` to be updated on `Drop` and is
-        //    referenced while taking a new `MerkHandle` to check if it was already
-        //    borrowed, that gives us two shared references to the same memory and
-        //    that's allowed, note we're not referring to the Merk part of the pair
-        // 3. Borrow flag's reference points to a heap allocated memory and will remain
-        //    valid just as the Merk reference to the memory right after the flag
+        // Borrowing rules are covered using a borrow flag of each Merk:
+        // 1. Borrow flag's reference points to a heap allocated memory and will remain
+        //    valid. Since the reference is shared and no need to obtain a `&mut`
+        //    reference this part of the memory is covered.
+        // 2. For the same reason the Merk's pointer can be converted to a reference,
+        //    because the memory behind the `Box` is valid and `MerkHandle` can't
+        //    outlive it since we use lifetime parameters.
+        // 3. We can get unique reference out of that pointer safely because of
+        //    borrowing flag.
         Ok(unsafe {
             MerkHandle {
-                merk: merk_ptr.as_mut().expect("`Box` contents are never null"),
+                merk: merk_ptr,
                 taken_handle: taken_handle_ref
                     .as_ref()
                     .expect("`Box` contents are never null"),
@@ -175,9 +175,18 @@ impl<'db, 'c> MerkHandle<'db, 'c> {
         if self.taken_handle.get() {
             panic!("Attempt to have double &mut borrow on Merk");
         }
+
         self.taken_handle.set(true);
-        let result = f(unsafe { self.merk.as_mut().expect("pointer to Box cannot be null") });
+
+        // SAFETY: here we want to have `&mut` reference to Merk out of a pointer, there
+        // is a checklist for that:
+        // 1. Memory is valid, because `MerkHandle` can't outlive `MerkCache` and heap
+        //    allocated Merks stay at their place for the whole `MerkCache` lifetime.
+        // 2. No other references exist because of `taken_handle` check above.
+        let result = f(unsafe { self.merk.as_mut().expect("`Box` contents are never null") });
+
         self.taken_handle.set(false);
+
         result
     }
 }

grovedb/src/tests/mod.rs

@@ -1190,8 +1190,6 @@ mod tests {
             )
             .unwrap();
 
-        dbg!(&result);
-
         assert!(matches!(
             result,
             Err(Error::CorruptedReferencePathKeyNotFound(_))

grovedb/src/util.rs

@@ -163,6 +163,10 @@ pub(crate) fn follow_reference<'db, 'b, 'c, B: AsRef<[u8]>>(
             &mut cost,
             referred_merk
                 .for_merk(|m| { Element::get(m, &referred_key, true, merk_cache.version) })
+                .map_err(|e| match e {
+                    Error::PathKeyNotFound(s) => Error::CorruptedReferencePathKeyNotFound(s),
+                    e => e,
+                })
         );
 
         match element {