fix: eliminate Owner poisoning in own_result by requiring value recovery on error

merk/src/owner.rs

@@ -64,26 +64,44 @@ impl<T> Owner<T> {
     }
 
     /// Takes temporary ownership of the contained value by passing it to `f`.
-    /// The function must return a result of the same type (the same value, or a
-    /// new value to take its place).
+    /// The closure must return `Ok(T)` on success or `Err((T, E))` on failure.
+    ///
+    /// **Error-safe:** In both `Ok` and `Err` paths the `Owner` retains a valid
+    /// value afterward -- on error, the value bundled in the `Err` variant is
+    /// restored.
     ///
-    /// # Warning
+    /// **Not unwind-safe:** If the closure panics, `Owner` will be left in a
+    /// poisoned state (inner = None) because the value was moved out before
+    /// the call and restoration happens only on normal return.
     ///
-    /// If `f` returns `Err`, the contained value has been consumed and the
-    /// `Owner` is left in a poisoned state (`inner = None`). Any subsequent
-    /// access (deref, `own`, `own_return`, etc.) will panic. Callers **must**
-    /// not use the `Owner` after `own_result` returns an error.
+    /// # Example
+    /// ```
+    /// # use grovedb_merk::owner::Owner;
+    /// let mut owner = Owner::new(42);
+    /// let result = owner.own_result(|v| {
+    ///     if v > 100 {
+    ///         Ok(v + 1)
+    ///     } else {
+    ///         Err((v, "too small"))   // value is returned alongside the error
+    ///     }
+    /// });
+    /// assert!(result.is_err());
+    /// assert_eq!(*owner, 42); // Owner still holds the original value
+    /// ```
     pub fn own_result<F, E>(&mut self, f: F) -> Result<(), E>
     where
-        F: FnOnce(T) -> Result<T, E>,
+        F: FnOnce(T) -> Result<T, (T, E)>,
     {
         let old_value = unwrap(self.inner.take());
         match f(old_value) {
             Ok(new_value) => {
                 self.inner = Some(new_value);
                 Ok(())
             }
-            Err(e) => Err(e),
+            Err((restored_value, e)) => {
+                self.inner = Some(restored_value);
+                Err(e)
+            }
         }
     }
 
@@ -112,8 +130,8 @@ fn unwrap<T>(option: Option<T>) -> T {
         Some(value) => value,
         None => panic!(
             "Owner is in a poisoned state (inner value is None). \
-             This can happen if `own_result` was called and the closure returned Err, \
-             consuming the value. The Owner must not be used after such an error."
+             This can happen if a closure passed to `own_result` panicked \
+             before returning Ok/Err."
         ),
     }
 }
@@ -125,17 +143,29 @@ mod tests {
     #[test]
     fn test_own_result_success_preserves_value() {
         let mut owner = Owner::new(42);
-        let result = owner.own_result(|v| Ok::<_, ()>(v + 1));
+        let result = owner.own_result(|v| Ok::<_, (i32, ())>(v + 1));
         assert!(result.is_ok());
         assert_eq!(*owner, 43);
     }
 
     #[test]
-    #[should_panic(expected = "Owner is in a poisoned state")]
-    fn test_own_result_error_poisons_owner() {
+    fn test_own_result_error_restores_value() {
+        let mut owner = Owner::new(42);
+        let result = owner.own_result(|v| Err::<i32, (i32, &str)>((v, "fail")));
+        assert!(result.is_err());
+        // Owner still holds the original value -- no poisoning
+        assert_eq!(*owner, 42);
+    }
+
+    #[test]
+    fn test_own_result_error_restores_modified_value() {
         let mut owner = Owner::new(42);
-        let _ = owner.own_result(|_v| Err::<i32, &str>("fail"));
-        // Accessing the poisoned Owner should panic with a clear message
-        let _ = *owner;
+        let result = owner.own_result(|v| {
+            let modified = v + 10;
+            Err::<i32, (i32, &str)>((modified, "fail"))
+        });
+        assert!(result.is_err());
+        // Owner holds the value returned in the Err variant
+        assert_eq!(*owner, 52);
     }
 }

merk/src/tree/mod.rs

@@ -755,7 +755,8 @@ impl TreeNode {
     /// Returns an error if there is already a child on the given side,
     /// indicating a corrupted tree state.
     #[inline]
-    pub fn attach(mut self, left: bool, maybe_child: Option<Self>) -> Result<Self, Error> {
+    #[allow(clippy::result_large_err)]
+    pub fn attach(mut self, left: bool, maybe_child: Option<Self>) -> Result<Self, (Self, Error)> {
         debug_assert_ne!(
             Some(self.key()),
             maybe_child.as_ref().map(|c| c.key()),
@@ -765,11 +766,14 @@ impl TreeNode {
         let slot = self.slot_mut(left);
 
         if slot.is_some() {
-            return Err(Error::CorruptedState(if left {
-                "Tried to attach to left tree slot, but it is already Some"
-            } else {
-                "Tried to attach to right tree slot, but it is already Some"
-            }));
+            return Err((
+                self,
+                Error::CorruptedState(if left {
+                    "Tried to attach to left tree slot, but it is already Some"
+                } else {
+                    "Tried to attach to right tree slot, but it is already Some"
+                }),
+            ));
         }
         *slot = Link::maybe_from_modified_tree(maybe_child);
 
@@ -833,7 +837,7 @@ impl TreeNode {
         F: FnOnce(Option<Self>) -> Option<Self>,
     {
         let (tree, maybe_child) = self.detach(left);
-        tree.attach(left, f(maybe_child))
+        tree.attach(left, f(maybe_child)).map_err(|(_tree, e)| e)
     }
 
     /// Like `walk`, but returns an error if there is no child on the given
@@ -844,7 +848,7 @@ impl TreeNode {
         F: FnOnce(Self) -> Option<Self>,
     {
         let (tree, child) = self.detach_expect(left)?;
-        tree.attach(left, f(child))
+        tree.attach(left, f(child)).map_err(|(_tree, e)| e)
     }
 
     /// Returns a mutable reference to the child slot for the given side.
@@ -860,6 +864,7 @@ impl TreeNode {
     /// Replaces the root node's value with the given value and returns the
     /// modified `Tree`.
     #[inline]
+    #[allow(clippy::result_large_err)]
     pub fn put_value(
         mut self,
         value: Vec<u8>,
@@ -885,7 +890,7 @@ impl TreeNode {
             (StorageRemovedBytes, StorageRemovedBytes),
             Error,
         >,
-    ) -> CostResult<Self, Error> {
+    ) -> CostContext<Result<Self, (Self, Error)>> {
         let mut cost = OperationCost::default();
 
         self.inner.kv = self.inner.kv.put_value_no_update_of_hashes(value);
@@ -895,15 +900,14 @@ impl TreeNode {
             // we are replacing a value
             // in this case there is a possibility that the client would want to update the
             // element flags based on the change of values
-            cost_return_on_error_no_add!(
-                cost,
-                self.just_in_time_tree_node_value_update(
-                    old_specialized_cost,
-                    get_temp_new_value_with_old_flags,
-                    update_tree_value_based_on_costs,
-                    section_removal_bytes
-                )
-            );
+            if let Err(e) = self.just_in_time_tree_node_value_update(
+                old_specialized_cost,
+                get_temp_new_value_with_old_flags,
+                update_tree_value_based_on_costs,
+                section_removal_bytes,
+            ) {
+                return Err((self, e)).wrap_with_cost(cost);
+            }
         }
 
         self.inner.kv = self.inner.kv.update_hashes().unwrap_add_cost(&mut cost);
@@ -913,6 +917,7 @@ impl TreeNode {
     /// Replaces the root node's value with the given value and returns the
     /// modified `Tree`.
     #[inline]
+    #[allow(clippy::result_large_err)]
     pub fn put_value_with_fixed_cost(
         mut self,
         value: Vec<u8>,
@@ -939,7 +944,7 @@ impl TreeNode {
             (StorageRemovedBytes, StorageRemovedBytes),
             Error,
         >,
-    ) -> CostResult<Self, Error> {
+    ) -> CostContext<Result<Self, (Self, Error)>> {
         let mut cost = OperationCost::default();
         self.inner.kv = self.inner.kv.put_value_with_fixed_cost_no_update_of_hashes(
             value,
@@ -951,15 +956,14 @@ impl TreeNode {
             // we are replacing a value
             // in this case there is a possibility that the client would want to update the
             // element flags based on the change of values
-            cost_return_on_error_no_add!(
-                cost,
-                self.just_in_time_tree_node_value_update(
-                    old_specialized_cost,
-                    get_temp_new_value_with_old_flags,
-                    update_tree_value_based_on_costs,
-                    section_removal_bytes
-                )
-            );
+            if let Err(e) = self.just_in_time_tree_node_value_update(
+                old_specialized_cost,
+                get_temp_new_value_with_old_flags,
+                update_tree_value_based_on_costs,
+                section_removal_bytes,
+            ) {
+                return Err((self, e)).wrap_with_cost(cost);
+            }
         }
 
         self.inner.kv = self.inner.kv.update_hashes().unwrap_add_cost(&mut cost);
@@ -969,6 +973,7 @@ impl TreeNode {
     /// Replaces the root node's value with the given value and value hash
     /// and returns the modified `Tree`.
     #[inline]
+    #[allow(clippy::result_large_err)]
     pub fn put_value_and_reference_value_hash(
         mut self,
         value: Vec<u8>,
@@ -995,7 +1000,7 @@ impl TreeNode {
             (StorageRemovedBytes, StorageRemovedBytes),
             Error,
         >,
-    ) -> CostResult<Self, Error> {
+    ) -> CostContext<Result<Self, (Self, Error)>> {
         let mut cost = OperationCost::default();
 
         self.inner.kv = self.inner.kv.put_value_no_update_of_hashes(value);
@@ -1005,15 +1010,14 @@ impl TreeNode {
             // we are replacing a value
             // in this case there is a possibility that the client would want to update the
             // element flags based on the change of values
-            cost_return_on_error_no_add!(
-                cost,
-                self.just_in_time_tree_node_value_update(
-                    old_specialized_cost,
-                    get_temp_new_value_with_old_flags,
-                    update_tree_value_based_on_costs,
-                    section_removal_bytes
-                )
-            );
+            if let Err(e) = self.just_in_time_tree_node_value_update(
+                old_specialized_cost,
+                get_temp_new_value_with_old_flags,
+                update_tree_value_based_on_costs,
+                section_removal_bytes,
+            ) {
+                return Err((self, e)).wrap_with_cost(cost);
+            }
         }
 
         self.inner.kv = self
@@ -1027,6 +1031,7 @@ impl TreeNode {
     /// Replaces the root node's value with the given value and value hash
     /// and returns the modified `Tree`.
     #[inline]
+    #[allow(clippy::result_large_err)]
     pub fn put_value_with_reference_value_hash_and_value_cost(
         mut self,
         value: Vec<u8>,
@@ -1054,7 +1059,7 @@ impl TreeNode {
             (StorageRemovedBytes, StorageRemovedBytes),
             Error,
         >,
-    ) -> CostResult<Self, Error> {
+    ) -> CostContext<Result<Self, (Self, Error)>> {
         let mut cost = OperationCost::default();
 
         self.inner.kv = self.inner.kv.put_value_with_fixed_cost_no_update_of_hashes(
@@ -1067,15 +1072,14 @@ impl TreeNode {
             // we are replacing a value
             // in this case there is a possibility that the client would want to update the
             // element flags based on the change of values
-            cost_return_on_error_no_add!(
-                cost,
-                self.just_in_time_tree_node_value_update(
-                    old_specialized_cost,
-                    get_temp_new_value_with_old_flags,
-                    update_tree_value_based_on_costs,
-                    section_removal_bytes
-                )
-            );
+            if let Err(e) = self.just_in_time_tree_node_value_update(
+                old_specialized_cost,
+                get_temp_new_value_with_old_flags,
+                update_tree_value_based_on_costs,
+                section_removal_bytes,
+            ) {
+                return Err((self, e)).wrap_with_cost(cost);
+            }
         }
 
         self.inner.kv = self

merk/src/tree/walk/mod.rs

@@ -130,6 +130,7 @@ where
 
     /// Similar to `Tree#walk`, but yields a `Walker` which fetches from the
     /// same source as `self`.
+    #[allow(clippy::result_large_err)]
     pub fn walk<F, T, V>(
         self,
         left: bool,
@@ -158,6 +159,7 @@ where
 
     /// Similar to `Tree#walk_expect` but yields a `Walker` which fetches from
     /// the same source as `self`.
+    #[allow(clippy::result_large_err)]
     pub fn walk_expect<F, T, V>(
         self,
         left: bool,
@@ -210,6 +212,7 @@ where
     ///
     /// Returns an error if the slot is already occupied, indicating a
     /// corrupted tree state.
+    #[allow(clippy::result_large_err)]
     pub fn attach<T>(mut self, left: bool, maybe_child: Option<T>) -> Result<Self, Error>
     where
         T: Into<TreeNode>,
@@ -220,6 +223,7 @@ where
     }
 
     /// Similar to `Tree#put_value`.
+    #[allow(clippy::result_large_err)]
     pub fn put_value(
         mut self,
         value: Vec<u8>,
@@ -264,6 +268,7 @@ where
     }
 
     /// Similar to `Tree#put_value_with_fixed_cost`.
+    #[allow(clippy::result_large_err)]
     pub fn put_value_with_fixed_cost(
         mut self,
         value: Vec<u8>,
@@ -310,6 +315,7 @@ where
     }
 
     /// Similar to `Tree#put_value_and_reference_value_hash`.
+    #[allow(clippy::result_large_err)]
     pub fn put_value_and_reference_value_hash(
         mut self,
         value: Vec<u8>,
@@ -356,6 +362,7 @@ where
     }
 
     /// Similar to `Tree#put_value_with_reference_value_hash_and_value_cost`.
+    #[allow(clippy::result_large_err)]
     pub fn put_value_with_reference_value_hash_and_value_cost(
         mut self,
         value: Vec<u8>,