Compiler Explorer

Source code

use core::{
    convert::Infallible,
    marker::PhantomData,
    mem::ManuallyDrop,
    ptr::NonNull,
};

/// Reference to account data or lamports with checked borrow rules.
pub struct Ref<'a, T: ?Sized> {
    value: NonNull<T>,
    state: NonNull<u8>,
    /// Indicates the type of borrow (lamports or data) by representing the
    /// shift amount.
    borrow_shift: u8,
    /// The `value` raw pointer is only valid while the `&'a T` lives so we claim
    /// to hold a reference to it.
    marker: PhantomData<&'a T>,
}

#[unsafe(no_mangle)]
pub fn map<'a>(input: Ref<'a, String>) -> Ref<'a, str> {
    Ref::map(input, |i| i.as_str())
}

#[unsafe(no_mangle)]
pub fn map_inline<'a>(input: Ref<'a, String>) -> Ref<'a, str> {
    Ref::map_inline(input, |i| i.as_str())
}

impl<'a, T: ?Sized> Ref<'a, T> {
    /// Maps a reference to a new type.
    #[inline]
    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.
    #[inline]
    pub fn map_inline<U: ?Sized, F>(orig: Ref<'a, T>, f: F) -> Ref<'a, U>
    where
        F: FnOnce(&T) -> &U,
    {
        let orig = ManuallyDrop::new(orig);

Ref {
            value: NonNull::from(f(&*orig)),
            state: orig.state,
            borrow_shift: orig.borrow_shift,
            marker: PhantomData,
        }
    }

/// 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)?),
            state: orig.state,
            borrow_shift: orig.borrow_shift,
            marker: PhantomData,
        })
    }
}

impl<T: ?Sized> core::ops::Deref for Ref<'_, T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        unsafe { self.value.as_ref() }
    }
}

impl<T: ?Sized> Drop for Ref<'_, T> {
    // decrement the immutable borrow count
    fn drop(&mut self) {
        unsafe { *self.state.as_mut() += 1 << self.borrow_shift };
    }
}