update tests and api for try map ref

Author: stegaBOB

sdk/pinocchio/src/account_info.rs

@@ -1,7 +1,6 @@
 //! Data structures to represent account information.
 
 use core::{
-    convert::Infallible,
     marker::PhantomData,
     mem::ManuallyDrop,
     ptr::{write, NonNull},
@@ -662,25 +661,36 @@ impl<'a, T: ?Sized> Ref<'a, T> {
     pub fn map<U: ?Sized, F>(orig: Ref<'a, T>, f: F) -> Ref<'a, U>
     where
         F: FnOnce(&T) -> &U,
-    {
-        Self::try_map::<_, _, Infallible>(orig, |x| Ok(f(x))).unwrap()
-    }
-
-    /// Maps a reference to a new type in a fallible map, returning an error if the map fails.
-    #[inline]
-    pub fn try_map<U: ?Sized, F, E>(orig: Ref<'a, T>, f: F) -> Result<Ref<'a, U>, E>
-    where
-        F: FnOnce(&T) -> Result<&U, E>,
     {
         // Avoid decrementing the borrow flag on Drop.
         let orig = ManuallyDrop::new(orig);
-
-        Ok(Ref {
-            value: NonNull::from(f(&*orig)?),
+        Ref {
+            value: NonNull::from(f(&*orig)),
             state: orig.state,
             borrow_shift: orig.borrow_shift,
             marker: PhantomData,
-        })
+        }
+    }
+
+    /// Tries to makes a new `Ref` for a component of the borrowed data.
+    /// On failure, the original guard is returned alongside with the error
+    /// returned by the closure.
+    #[inline]
+    pub fn try_map<U: ?Sized, E>(
+        orig: Ref<'a, T>,
+        f: impl FnOnce(&T) -> Result<&U, E>,
+    ) -> Result<Ref<'a, U>, (Self, E)> {
+        // Avoid decrementing the borrow flag on Drop.
+        let orig = ManuallyDrop::new(orig);
+        match f(&*orig) {
+            Ok(value) => Ok(Ref {
+                value: NonNull::from(value),
+                state: orig.state,
+                borrow_shift: orig.borrow_shift,
+                marker: PhantomData,
+            }),
+            Err(e) => Err((ManuallyDrop::into_inner(orig), e)),
+        }
     }
 
     /// Filters and maps a reference to a new type.
@@ -743,26 +753,36 @@ impl<'a, T: ?Sized> RefMut<'a, T> {
     pub fn map<U: ?Sized, F>(orig: RefMut<'a, T>, f: F) -> RefMut<'a, U>
     where
         F: FnOnce(&mut T) -> &mut U,
-    {
-        Self::try_map::<_, _, Infallible>(orig, |x| Ok(f(x))).unwrap()
-    }
-
-    /// Maps a mutable reference to a new type in a fallible map, returning an error if the map fails.
-    #[inline]
-    pub fn try_map<U: ?Sized, F, E>(orig: RefMut<'a, T>, f: F) -> Result<RefMut<'a, U>, E>
-    where
-        F: FnOnce(&mut T) -> Result<&mut U, E>,
     {
         // Avoid decrementing the borrow flag on Drop.
         let mut orig = ManuallyDrop::new(orig);
-        let value = f(&mut *orig)?;
-
-        Ok(RefMut {
-            value: NonNull::from(value),
+        RefMut {
+            value: NonNull::from(f(&mut *orig)),
             state: orig.state,
             borrow_bitmask: orig.borrow_bitmask,
             marker: PhantomData,
-        })
+        }
+    }
+
+    /// Tries to makes a new `RefMut` for a component of the borrowed data.
+    /// On failure, the original guard is returned alongside with the error
+    /// returned by the closure.
+    #[inline]
+    pub fn try_map<U: ?Sized, E>(
+        orig: RefMut<'a, T>,
+        f: impl FnOnce(&mut T) -> Result<&mut U, E>,
+    ) -> Result<RefMut<'a, U>, (Self, E)> {
+        // Avoid decrementing the borrow flag on Drop.
+        let mut orig = ManuallyDrop::new(orig);
+        match f(&mut *orig) {
+            Ok(value) => Ok(RefMut {
+                value: NonNull::from(value),
+                state: orig.state,
+                borrow_bitmask: orig.borrow_bitmask,
+                marker: PhantomData,
+            }),
+            Err(e) => Err((ManuallyDrop::into_inner(orig), e)),
+        }
     }
 
     /// Filters and maps a mutable reference to a new type.
@@ -773,17 +793,13 @@ impl<'a, T: ?Sized> RefMut<'a, T> {
     {
         // Avoid decrementing the mutable borrow flag on Drop.
         let mut orig = ManuallyDrop::new(orig);
-
         match f(&mut *orig) {
-            Some(value) => {
-                let value = NonNull::from(value);
-                Ok(RefMut {
-                    value,
-                    state: orig.state,
-                    borrow_bitmask: orig.borrow_bitmask,
-                    marker: PhantomData,
-                })
-            }
+            Some(value) => Ok(RefMut {
+                value: NonNull::from(value),
+                state: orig.state,
+                borrow_bitmask: orig.borrow_bitmask,
+                marker: PhantomData,
+            }),
             None => Err(ManuallyDrop::into_inner(orig)),
         }
     }
@@ -841,6 +857,19 @@ mod tests {
         assert_eq!(state, NOT_BORROWED - (1 << DATA_BORROW_SHIFT));
         assert_eq!(*new_ref, 3);
 
+        let Ok(new_ref) = Ref::try_map::<_, u8>(new_ref, |_| Ok(&4)) else {
+            unreachable!()
+        };
+
+        assert_eq!(state, NOT_BORROWED - (1 << DATA_BORROW_SHIFT));
+        assert_eq!(*new_ref, 4);
+
+        let (new_ref, err) = Ref::try_map::<u8, u8>(new_ref, |_| Err(5)).unwrap_err();
+        assert_eq!(state, NOT_BORROWED - (1 << DATA_BORROW_SHIFT));
+        assert_eq!(err, 5);
+        // Unchanged
+        assert_eq!(*new_ref, 4);
+
         let new_ref = Ref::filter_map(new_ref, |_| Option::<&u8>::None);
 
         assert_eq!(state, NOT_BORROWED - (1 << DATA_BORROW_SHIFT));