Golang's Defer in Rust

I am trying to recreate the defer keyword from Go in Rust. I did find this Github post that uses a macro. However, it doesn't work for calls to defer that have a smaller scope than the entire function. For example this example from go blog:

for i := 0; i < 4; i++ {
        defer fmt.Print(i)
}

This should output: 3 2 1 0

for x in 0..4 {
        defer!(println!("{}", x));
}

The macro outputs: 0 1 2 3 because x goes out of scope before the next loop of for. So is there a way to get the correct function no matter where the scope of the defer is called? I've tried this but I'm a rust noob so pardon my ignorance as it doesn't work.

You'll need to either make it a Box<dyn FnMut()>, or store a plain fn(), or make it generic. FnMut allows closures with varying amounts of capture data, so different closures implementing FnMut are different sizes. You can fix this by doing one of the following:

  • Allow different sized FnMuts: change dyn FnMut() to Box<dyn FnMut()>

    This works because the Box will make separate allocations for each one, and then also store a vtable to differentiate between them

  • Don't allow captured variables at all: change dyn FnMut() to fn()

    fn() with a lowercase f is the type for a literal function pointer. This won't allow closures which capture variables, but that comes with the advantage of being able to store everything as a simple pointer to the function

  • Restrict to a specific type of closure: make Defers generic over T: FnMut(). This would fix the problem of different closures having different size by simply saying "each Defers can only store a specific closure". Since you're adding to it in a loop (and thus only add from one closure definition), this should work for this case.

    Since this is a bit more involved, I'll include a code snippet - it'd be roughly:

    struct Defers<T: FnMut()> {
        stack: Vec<T>,
    }
    
    impl<T: FnMut()> Defers<T> {
        fn Defer(&mut self, F: T) {
            self.stack.Push(F);
        }
    
        fn New() -> Defers {
            Defers { stack: Vec::new() }
        }
    }
    
    impl<T: FnMut()> Drop for Defers<T> { ...
    

On top of that, you have one other problem: when calling d.Defers, you're just putting the function call inside. This will call println!() immediately, and then pass the result into Defers - I don't think this is what you want. If I'm understanding correctly, this will do what you intend:

        d.Defer(|| println!("{}", x));
3 Likes

Thank you for your help. Here's what I came up with if any future person is trying to do this. Also, if there are some improvements possible let me know!

struct Defers {
    stack: Vec<Box<dyn Fn()>>,
}

impl Defers {
    fn defer(&mut self, f: Box<dyn Fn()>) {
        self.stack.push(f);
    }

    fn new() -> Defers {
        Defers { stack: Vec::new() }
    }
}

impl Drop for Defers {
    fn drop(&mut self) {
        while let Some(f) = self.stack.pop() {
            f();
        }
    }
}

fn main() {
    let mut d = Defers::new();
    for x in 0..4 {
        d.defer(Box::new(move || println!("{}", x)));
    }
}
4 Likes

The main improvement is that you can loosen the Fn bound into a FnOnce one :slightly_smiling_face:

3 Likes

For what it is worth, there is a popular crate which provides a version of this:
https://docs.rs/scopeguard/1.1.0/scopeguard/#defer
I don't know about it's ordering. (Edit: Nevermind. I see this was mentioned in the original SO post.) But it does have the advantage that it handles panics.