Soundness review: Accessing short-lived references from long-lived objects

I occasionally want to pass a borrow down the call stack without giving it explicitly as a parameter to every intermediate function along the way. This is particularly important when there are unmodifiable traits in the picture, like Future::poll() — I'd like to write an executor that provides access to some kind of context object to the futures it's running, but also mutates it when they aren't being processed.

I wrote this to make this possible, but I'm not sure if it's completely sound, especially with regard to panics. Is this a reasonable approach to take, or is there a better option out there?

use core::ptr::NonNull;
use std::cell::Cell;
use std::fmt::{Debug, Formatter};

/// Allows long-lived types to access short-lived ones
pub struct ScopedMut<T:?Sized>(Cell<Option<NonNull<T>>>);

impl<T:?Sized> Default for ScopedMut<T> {
    fn default()->Self { ScopedMut(Cell::new(None)) }
}

impl<T:?Sized> ScopedMut<T> {
    pub fn try_with_mut<F,O>(&self, f:F)->Option<O> where F:FnOnce(&mut T)->O {
        let mut ptr = self.0.replace(None)?;
        // Safety: The referenced object is exclusively borrowed for the scope of the containing
        // lend() call, and the reconstructed reference cannot escape `F`.
        let result = unsafe { f(ptr.as_mut()) };

        // NB: Should never fail, as the only changes to `ptr` occur within `lend()` and
        // `try_with_mut()`; they both restore the original value before they return.
        assert!(self.0.replace(Some(ptr)).is_none());
        Some(result)
    }

    pub fn with_mut<F,O>(&self, f:F)->O where F:FnOnce(&mut T)->O {
        self.try_with_mut(f).expect("No lent value")
    }
    
    pub fn lend<F,O>(&self, item: &mut T, f:F)->O where F:FnOnce()->O {
        let new = item.into();
        let old = self.0.replace(Some(new));
        
        let result = f();
        
        assert_eq!(self.0.replace(old), Some(new));
        result
    }
}

impl<T:?Sized + Debug> Debug for ScopedMut<T> {
    fn fmt(&self, f: &mut Formatter)->Result<(), std::fmt::Error> {
        self.try_with_mut(|x| f.debug_tuple("ScopedMut").field(&&*x).finish())
            .unwrap_or_else(|| f.write_str("ScopedMut(_)"))
    }
}

pub struct ScopedRef<T:?Sized>(Cell<Option<NonNull<T>>>);

impl<T:?Sized> Default for ScopedRef<T> {
    fn default()->Self { ScopedRef(Cell::new(None)) }
}

impl<T:?Sized> ScopedRef<T> {
    pub fn try_with_ref<F,O>(&self, f:F)->Option<O> where F:FnOnce(&T)->O {
        let ptr = self.0.get()?;
        // Safety: The referenced object will be alive for the duration of the containing lend()
        // call, and the reference we construct here can't escape the body of `F`.
        let result = unsafe { f(ptr.as_ref()) };
        Some(result)
    }

    pub fn with_ref<F,O>(&self, f:F)->O where F:FnOnce(&T)->O {
        self.try_with_ref(f).expect("No lent value")
    }
    
    pub fn lend<F,O>(&self, item: &T, f:F)->O where F:FnOnce()->O {
        let new = item.into();
        let old = self.0.replace(Some(new));
        
        let result = f();
        
        assert_eq!(self.0.replace(old), Some(new));
        result
    }
}

impl<T:?Sized + Debug> Debug for ScopedRef<T> {
    fn fmt(&self, f: &mut Formatter)->Result<(), std::fmt::Error> {
        self.try_with_ref(|x| f.debug_tuple("ScopedRef").field(&x).finish())
            .unwrap_or_else(|| f.write_str("ScopedRef(_)"))
    }
}

// #[test]
fn ref_loop() {
    let scoped_ref = ScopedRef::default();

    let closure = || scoped_ref.with_ref(|&x| 2*x);

    for x in 0..10 {
        dbg!(scoped_ref.lend(&x, closure));
    }
}

// #[test]
fn mut_loop() {
    let scoped_mut = ScopedMut::default();

    let closure = || scoped_mut.with_mut(|x| *x *= 2);

    for mut x in 0..10 {
        scoped_mut.lend(&mut x, closure);
        dbg!(x);
    }
}

fn main() {
    ref_loop();
    mut_loop();
}

(Playground)

lend is not sound in the presence of panics. If a panic happens there is no value for lend to put back into the mutable reference. The take_mut crate which does a similar operation to the lend method aborts on panics. It uses catch_unwind instead of a panicking drop guard though, which I'm not sure is correct in the face of foreign exceptions (these are currently unstable)

2 Likes

Shouldn't restoring old be ok¹? If it's None, that's valid by definition; otherwise, we've got nested lends: The outer lend is holding exclusive access to that referent, which was never provided to f.

¹ I understand that the code doesn't currently do this; I'd need to do something with catch_unwind to make it happen.

Usually, using some local guard object's destructor is preferred for panic-safe cleanup, because they're less overhead AFAIK.

3 Likes

So e.g. something like

pub fn lend<F, O>(&self, item: &mut T, f: F) -> O
where
    F: FnOnce() -> O,
{
    let new = item.into();
    let old = self.0.replace(Some(new));

    struct Guard<'a, T: ?Sized> {
        old: Option<NonNull<T>>,
        new: NonNull<T>,
        this: &'a ScopedMut<T>,
    }
    let _guard = Guard {
        old,
        new,
        this: self,
    };

    let result = f();

    impl<T: ?Sized> Drop for Guard<'_, T> {
        fn drop(&mut self) {
            assert_eq!(self.this.0.replace(self.old), Some(self.new));
        }
    }
    
    result
}

Using this pattern for try_with_mut as well would make sense too IMO, to make it more well-behaved, even when it might not be necessary for safety.


Just for fun, the same thing syntactically nicer by using scopeguard::defer: Rust Playground

1 Like

Thanks. I realized that the guard operation is exactly the same for all the operations; here's the updated code for ScopedMut:

struct Guard<'a,T:?Sized> {
    cell: &'a Cell<Option<NonNull<T>>>,
    old: Option<NonNull<T>>,
    new: Option<NonNull<T>>
}

impl<'a, T:?Sized> Guard<'a, T> {
    fn replace(cell: &'a Cell<Option<NonNull<T>>>, new: Option<NonNull<T>>)->Self {
        Guard {
            cell, new, old: cell.replace(new)
        }
    }
}

impl<'a, T:?Sized> Drop for Guard<'a, T> {
    fn drop(&mut self) {
        // NB: This shouldn't ever fail, as all changes are made via `Guard::replace`
        assert_eq!(self.new, self.cell.replace(self.old));
    }
}

impl<T:?Sized> ScopedMut<T> {
    pub fn try_with_mut<F,O>(&self, f:F)->Option<O> where F:FnOnce(&mut T)->O {
        let guard = Guard::replace(&self.0, None);
        
        // Safety: The referenced object is exclusively borrowed for the scope of the containing
        // lend() call, and the reconstructed reference cannot escape `F`.
        unsafe { Some(f(guard.old?.as_mut())) }
    }

    pub fn with_mut<F,O>(&self, f:F)->O where F:FnOnce(&mut T)->O {
        self.try_with_mut(f).expect("No lent value")
    }
    
    pub fn lend<F,O>(&self, item: &mut T, f:F)->O where F:FnOnce()->O {
        let _guard = Guard::replace(&self.0, Some(item.into()));
        f()
    }
}
1 Like