Idea: safe and efficient mutable overlays thanks to the borrow checker (and a macro)

I’m sharing the following in the hope that someone will find this interesting, or might want to point out possible problems with/alternatives to this approach.

Let’s say that in our library we have a type Something that can be modified in a way that can be expressed as a sequence of “patches”. That something could be an image, a 3d model, a graph, a database, a text buffer, etc. Let’s say that we routinely have to examine different sequences of (slight) modifications based on the same original, but we don’t want to clone the original many times, only to modify each clone and drop it.

We could introduce a type Overlay<'_> that stores a shared reference to an original Something plus the modifications to apply to it. Ideally Overlay values would allow the same operations that the original Something type allows. It might be even possible to chain several overlays. (This might be inefficient for long chains, but then long chains of overlays might not be needed.)

The nice thing about Rust is that it allows to implement the above pattern safely without any runtime cost like runtime reference counting or locking! It’s a simple idea, but it seems to be relevant to my work, and I think that it gives a nice example of the power of the borrow checker. I wonder if there are any libraries in the wild that utilize this pattern.

One problem with the idea is that since an overlay keeps a reference to an original, returning an overlay from a function or storing it somewhere is not straightforward. Here comes an idea how to mitigate this situation: we could introduce an operation, let’s call it “absorb”, that allows to absorb one overlay into the original. (Since this mutates the original, this is only allowed if no other overlays of the same original exist.) In this way, instead of returning an overlay, we can absorb it into the original.

If this “absorb” operation was a function, it would have to take a mutable reference to the original plus it would have to take ownership of the overlay to be absorbed (which includes a reference to the orignal). The borrow checker will not allow this.

But with a macro it’s possible! (The listing below demonstrates the technique.)

// A modifiable overlay of a Vec that behaves itself as a Vec.
// This toy version only allows extending the original.
#[derive(Debug)]
struct VecOverlay<'a, T> {
    orig: &'a Vec<T>,
    ext: Vec<T>,
}

macro_rules! absorb {
    (&mut $v:expr, $o:expr $(,)?) => {{
        // Make sure that the overlay refers to the original.
        assert_eq!(&($v) as *const _, ($o).orig as *const _);
        // We took possession of o, so we can as well mutate it.
        let mut o = ($o);
        // Append the supplement to the original.
        ($v).append(&mut o.ext);
    }};
}

fn main() {
    let mut v = vec![0, 1, 2];
    let mut o = VecOverlay { orig: &v, ext: vec![3] };

    // The overlay can be modified.
    o.ext.push(4);

    // There can be another overlay ...
    let mut o2 = VecOverlay { orig: &v, ext: vec![10, 11] };
    // ... that can be modified as well.
    o2.ext.push(12);

    // Absorb one overlay.  The macro effectively takes a mutable reference to a Vec and AT THE
    // SAME TIME takes ownership of a VecOverlay (that must refer to the same Vec).
    absorb!(&mut v, o);

    dbg!(&v);
}

Here's a related pattern.

    let mut v = vec![0, 1, 2];
    VecOverlay::absorb_if(&mut v, vec![3], |mut o| {
        o.ext.push(4);
        let mut o2 = o.another(vec![10, 11]);
        o2.ext.push(12);
        Some(o.ext)
    });
    
    dbg!(&v);
impl<'a, T> VecOverlay<'a, T> {
    pub fn new(orig: &'a Vec<T>, ext: Vec<T>) -> Self {
        Self { orig, ext }
    }

    pub fn another(&self, ext: Vec<T>) -> Self {
        Self::new(self.orig, ext)
    }
    
    pub fn absorb_if<F: FnOnce(VecOverlay<'_, T>) -> Option<Vec<T>>>(
        orig: &mut Vec<T>,
        ext: Vec<T>,
        extender: F,
    ) {
        if let Some(ext) = extender(VecOverlay::new(&*orig, ext)) {
            orig.extend(ext);
        }
    }
}

Could be tweaked in a number of ways (just return Vec<T>, return VecOverlay<'_, T>, ...).

2 Likes

This is an interesting technique, however note that overlay fields would have to be public if it was defined in a different module.

Note that the assert_eq in your code is a runtime check. And while it’s not required for memory safety (not until you use unsafe in the overlay that is) it is absolutely necessary for correctness.

Besides, reference counting shouldn’t be expensive. Even Arc is fine for most purposes, and Rc is basically free. Absorption could be a problem but from your post it seems you want to store, return, and do other stuff on overlays, and with an Rc instead of a reference inside it’s trivial.

The Graphite vector art editor uses a similar approach where edits are stored as a sequence of nodes in a graph. As far as I know, they use IDs instead of references avoiding the borrow problem. The editor is open source so maybe you find a solution there.

Being able to replace the macro by a function/method that takes a closure would be nice, but it seems to me that this breaks one aspect of the pattern that I consider critical:

In the above snippet, I would like the closure to return a value of type Option<VecOverlay<'a, T>> (e.g. Some(o)) and not a value of type Option<Vec<T>> (e.g. Some(o.ext)). The reason is that when returning Some(o.ext) above there is nothing that guarantees that o.ext is part of an overlay of the original vector v. It could be actually any Vec<i32>, and this will work in this toy example, but one of the purposes of this pattern is to ensure that any overlay may be only applied to the original to which it refers to.

However, it seems that there is no way to make the borrow checker accept this.

I imagine that in production code with complicated types, the Something type would have an unsafe method absorb_bare_overlay(&mut self, owr: Overlay_without_reference). Then, the macro would only compare the reference in the overlay with the original, and then call that function which is only safe to be called when one is sure that the overlay fits the base to be applied to.

Although the borrow checker knows when they are the same or not, I believe that there is no formally compile-time way to do the check.

But then it’s just a pointer comparison, and I would be surprised if the compiler was not able to optimize it away, since both the creation of the overlay and its “absorption” necessarily happen within a single function.

My real potential use case for this pattern is a library for tensor network computations that I work on (it’s meant for “quantum-inspired” algorithms, not for machine learning). My Something would be a tensor network that stores a graph with one tensor per node. The tensors are essentially multi-dimensional arrays of numbers that can be anything between small and huge. My Overlay would be a “tensor network based on another one”.

Now I could simply store the arrays behind Arcs (and I might end up going this way), but I see the following problems:

  • While Arcs are not a performance problem in general, is this still true when managing lots of small objects (down to arrays of 4×4×4 f32, i.e. 256 bytes) while running 64 threads on a NUMA machine?

  • Putting the central and most numerous kind of object behind Arcs feels not very idiomatic in Rust.

In addition, it’s not like our networks will share tensors in unstructured ways. It’s rather that a single network will typically serve as basis for multiple short-lived networks.

This is how I came up with the “overlay” idea. I might actually not bother with deltas to the graph itself (which is almost never big), and only restict myself to allowing overlays use (copy-on-write) the tensors that are present in the original network.

I think the borrow checker only knows/cares whether they have related lifetimes. Different references may have the same lifetime.

On the other hand, sometimes a lifetime itself can be used as a tag; I’ve seen that pattern in dyngo. I wonder if that idea could be used here...

If it results in a better API then sure, go ahead. I initially got the impression you’re more concerned about runtime costs.

Also, how do you share references across threads? Do you use scoped threads?

It's possible. See the section on Branded Types in this paper.

2 Likes

You could use tests for that. Normally, you don't want to have test environments, such as assert, directly in the applied module, but rather store them directly in a #[test] module, which is called with cargo test. If you don't trust your functions and macros in principle, it is actually better to have them return an Option<T> or a Result<T,E> than to fiddle around with assert calls everywhere. Asserts cost performance in the long run, and many people like their programmes to run fast. Besides that - Cool concept. Oh, I've just had an idea. You don't need to write any complex test modules at first. Cargo test also tests your /// comments. So, to begin with, you can simply put your assertions in your documentation. This might even be easier at the start, because then you don't have to worry so much about Cargo's implementation of the test environment. Here's an example, how a test module for your crate could look like: gitlab

This is amazing! Thanks for the link to the research paper. And I thought that I understand possible uses of PhantomData...

I evolved your variant further towards “realism” (attached below). In particular:

  • I introduced a type for the base. The function that takes the closure is now a method of this MyVec type and is called update_with. Next to &mut self it only takes the closure.

  • The inner workings of VecOverlay are now private to the overlay module.

  • Next to the update_with method, an update_with macro is provided as well.

I’m quite satisfied with the result, but I’m still not 100% sure about the following:

  • Is it indeed always true that the check inside the update_with method is not necessary (and hence can be a debug_assert)?

  • Are there any use cases that truly require the macro variant? I believe not, although the macro seems useful for cases where overlays over multiple vecs are alive simultaneously.

#[macro_use]
pub mod overlay {
    use std::marker::PhantomData;

    #[derive(Debug, Clone)]
    pub struct VecOverlay<'brand, 'borrow, T> {
        orig: &'borrow MyVec<T>,
        ext: Vec<T>,
        _invariance: PhantomData<fn(&'brand ()) -> &'brand()>
    }

    impl<'brand, 'borrow, T> VecOverlay<'brand, 'borrow, T> {
        pub fn is_based_on(&self, myvec: &MyVec<T>) -> bool {
            self.orig as *const _ == myvec as *const _
        }

        #[doc(hidden)]
        // For use in ‘update_with’ macro.
        pub fn _extract_delta(self) -> Vec<T> {
            self.ext
        }

        pub fn push(&mut self, value: T) {
            self.ext.push(value);
        }
    }

    // This would be some non-trivial type that the overlay is based on.  In this toy example we
    // simply use a minimalistic newtype of Vec.
    #[derive(Debug, Clone)]
    pub struct MyVec<T>(pub Vec<T>);

    impl<T> MyVec<T> {
        #[doc(hidden)]
        /// Apply the delta of an overlay.  Attention: for correctness, the caller must be sure that
        /// the given overlay refers to ‘self’!
        ///
        /// For use in ‘update_with’ macro.  May have to be declared as unsafe if ignoring the above
        /// requirement could leave to UB.
        pub fn _apply_delta(&mut self, mut delta: Vec<T>) {
            self.0.append(&mut delta);
        }

        pub fn to_overlay(&self) -> VecOverlay<'_, '_, T>  {
            VecOverlay { orig: &self, ext: Vec::new(), _invariance: PhantomData }
        }

        pub fn update_with<F>(
            &mut self,
            maybe_make_overlay: F,
        ) where
            F: for<'x, 'y> FnOnce(VecOverlay<'x, 'y, T>) -> Option<VecOverlay<'x, 'y, T>>,
        {
            if let Some(vo) = maybe_make_overlay(self.to_overlay()) {
                // The following is ensured by the borrow checker:
                debug_assert!(vo.is_based_on(&self));
                self._apply_delta(vo._extract_delta());
            }
        }
    }

    // The macro variant is useful when overlays to multiple vectors are alive simultanously.
    #[allow(unused_macros)]
    macro_rules! update_with {
        (&mut $v:expr, $o:expr $(,)?) => {{
            // The following assert is crucial: the overlay could refer to a different MyVec!
            assert!($o.is_based_on(&$v));
            $v._apply_delta($o._extract_delta());
        }};
    }
}

use overlay::*;

fn main() {
    let mut v = MyVec(vec![0, 1, 2]);

    // Use of ‘update_with’ method: an arbitrary number of overlays over the same vector can be
    // handled.  In the end, one of them might be absorbed into the base.
    v.update_with(|mut o| {
        o.push(3);
        let mut o2 = o.clone();
        o2.push(12);
        o.push(4);
        Some(o)
    });

    dbg!(&v);
}

Later addition:

Hmm, I removed the 'brand and PhantomData bits and it still compiles and runs as before. Moreover, I cannot find a way to hack it to prove that what I removed was necessary.

But I remember that in the earlier version returning Some(o) from the closure was not accepted by the borrow checker. So this must be somehow related to the reorganization that happened inbetween.

I'll look at the rest in a bit, but...

Here you go.

     let mut v = MyVec(vec![0, 1, 2]);
+    // Get a `VecOverlay<'static, _>`....
+    let mut dummy = Box::leak(Box::new(MyVec(vec![])));
+    let overlay = dummy.to_overlay();

     // Use of ‘update_with’ method: an arbitrary number of overlays over the same vector can be
     // handled.  In the end, one of them might be absorbed into the base.
     v.update_with(|mut o| {
         o.push(3);
         let mut o2 = o.clone();
         o2.push(12);
         o.push(4);
-        Some(o)
+        // ...which can coerce to any other lifetime.
+        Some(overlay)
    });

    dbg!(&v);

(You could use a static instead of leaking but I'm lazy.)

1 Like

I believe so. But here's me thinking about it "out loud":

This allows the consumer to control the creation of 'brands.

        // Caller choses lifetime on `&self` which becomes the `'brand`
        // of the `VecOverlay`
        pub fn to_overlay(&self) -> VecOverlay<'_, '_, T>  {
            VecOverlay { orig: &self, ext: Vec::new(), _invariance: PhantomData }
        }

However, it's impossible to use that method to create a VecOverlay where the 'brand is shorter than the 'borrow, so this still fails:

    v.update_with(|mut o| {
        o.push(3);
        let mut o2 = o.clone();
        o2.push(12);
        o.push(4);
        static dummy: MyVec<i32> = MyVec(Vec::new());
        // The closure has to be able to handle every combination
        // of `'brand` and `'borrow`, but there are some combinations
        // which `to_overlay` cannot produce ==> error.
        Some(dummy.to_overlay())
    });

...but that feels pretty thin to me. The idea behind branding is that the caller can't know what 'brand they have. Ideally that method would be private -- make it so the only (safe) way to get a branded VecOverlay is within a closure that can handle any 'brand (like the where bounds of update_with). But that rules out the macro approach (unless you split up branded and unbranded overlays).

Thin or not, I think it's still enough to not need the check in update_with, as things are now. Be careful whenever you add more ways to create brands, and be careful to always require the closures to handle all possible combinations of the lifetimes.

The branding/GhostCell approach lacks ergonomics for that use case, definitely. Want 4 independent overlays? Nest four closures...[1] You could make a macro for that use case instead I suppose.

Outside of ergonomics... I can think of two (related) things possible with the macro approach that aren't currently possible with the closure approach, but you could fill those holes too.

  1. The MyVec is still directly shared-usable with the macro approach. You could mitigate this difference by giving shared access to the MyVec via VecOverlay.

  2. The macro approach -- really, the to_overlay method -- allows one to mess with overlays which are never applied by never calling the macro, and you can do this with only shared access to the MyVec.[2] Whereas the closure requires &mut MyVec<_>. You could mitigate this with a non-applying observe_with method or whatever.


If _apply_delta needs to be unsafe, your macro is unsound without hardening to make sure the types are what you expect they are.

struct Annoying<T>(Vec<T>);
impl<T> Annoying<T> {
    fn is_based_on<U>(&self, _: U) -> bool {
        true
    }
    fn _extract_delta(self) -> Vec<T> {
        self.0
    }
}

// ...

    let mut v = MyVec(vec![0, 1, 2]);
    update_with!(&mut v, Annoying(vec![67]));
    println!("{v:?}");

Unfortunately, hardening macros is not my forte. Perhaps start a new topic if you want to pursue that route.


  1. You could have methods that support more than one in a single call. But without far-distant language features, you can't support an arbitrary number, because you need a distinct brand per instance. (The far-distance language feature is "generic variadics".) ↩︎

  2. say in fn(v: &MyVec<..>) ↩︎

1 Like