How to compose functions that modify a struct in Rust

Question

Take a super-simple struct Foo:

#[derive(Debug)]
struct Foo {
    a: i32
}

and a compose macro I got here:

macro_rules! compose {
    ( $last:expr ) => { $last };
    ( $head:expr, $($tail:expr), +) => {
        compose_two($head, compose!($($tail),+))
    };
}

fn compose_two<A, B, C, G, F>(f: F, g: G) -> impl Fn(A) -> C
where
    F: Fn(A) -> B,
    G: Fn(B) -> C,
{
    move |x| g(f(x))
}

I can define a simple function that takes a mutable reference and modifies the struct and returns the reference it was handed:

fn foo(x: &mut Foo) -> &mut Foo {
    x.a = x.a * 2;
    x
}

and it works as expected:

fn main() {
    let mut x = Foo { a: 3 };
    let y = foo(&mut x);
    println!("{:?}", y.a); // prints 6
    y.a = 7;
    println!("{:?}", x); // prints Foo { a: 7 }
}

The problem comes when I try to define a second simple function and compose the two:

fn bar(x: &mut Foo) -> &mut Foo {
    x.a = x.a + 1;
    x
}

fn main() {
    let baz = compose!(foo, bar);
    let mut x = Foo { a: 3 };
    let y = baz(&mut x);
    println!("{:?}", y.a);
}

I get an error that the mutable borrow of x in main let y = baz(&mut x); doesn't live long enough. I don't think I understand that compose macro well enough to understand what's gong wrong.

Also when I print the struct bound to x in the first version it works because it's after the last use of the mutable borrow y so I can immutably borrow x to print it. But in the second version if I try to print x at the end it says it's still borrowed mutably. Something in the compose macro seems to be "holding on" to that mutable borrow of x?

How do I make this work? Can this be made to work?

Playground

Edit based on comments:

It seems that while the closure in compose_two doesn't actually hold on to the mutable reference to the struct the return type doesn't specify that it doesn't (closures close over captured variables right?), and so the compiler is forced to assume that it might. How do I convince the compiler that I'm not holding that reference?

Answer 1

Can this be made to work?

No. But depending on your use-case, maybe yes.

You can make it work if either:

• You can constrain the function types to be Copy (if you are using function items (fn), they are always Copy, but for closures this may be a problem if they capture non-Copy types).
• You can use nightly.
• You can change the user of compose!() (main()).
• You can limit compose!() to references (mutable references, to be precise, but you can make a version for shared references too. Of course, if you want to make separate versions for references and owned types, this is fine).

There are three factors here, that joined together they convince the compiler x can be used after its lifetime. If we break one of them, it will work. Two of them are actually false, but the compiler doesn't know that (or doesn't want to rely on that). The factors are:

1. The compiler believes that the returned closure can capture its parameter. This is false, as I will explain in a minute, but the compiler does not know that.
2. The compiler believes the closure has a drop implementation, and can use x (captured in step 1) inside this drop. In fact, the compiler knows it doesn't, but because we used impl Trait, it is forced to treat it as if it implemented drop, so it will not be a breaking change to add one.
3. x is dropped before baz. This is true (variables are dropped in reversed order to their declaration), and combined with the two previous beliefs of the compiler it means that when baz will (potentially) use its captured x in its drop, it will be after x's lifetimes.

Let's start with the last claim. It is the easiest to break, because you only need to swap the order of x and baz:

fn main() {
    let mut x = 3;
    let baz = compose_two(foo, bar);
    let y = baz(&mut x);
    println!("{:?}", y);
}

But it is not always possible to change main(), or it may not be possible to declare x before baz.

So let's return to the second claim. The compiler believes the closure has a Drop impl because it is in impl Trait. What if it would not be?

Unfortunately, this requires nightly, because writing closures manually requires the features fn_traits and unboxed_closures. But it is definitely possible (and a nice side benefit is that the function can be conditionally FnOnce/FnMut/Fn depending on what its input functions are):

#![feature(fn_traits, unboxed_closures)]

struct ComposeTwo<G, F>(F, G);

impl<A, B, C, G, F> std::ops::FnOnce<(A,)> for ComposeTwo<G, F>
where
    F: FnOnce(A) -> B,
    G: FnOnce(B) -> C,
{
    type Output = C;

    extern "rust-call" fn call_once(self, (x,): (A,)) -> Self::Output {
        (self.1)((self.0)(x))
    }
}

impl<A, B, C, G, F> std::ops::FnMut<(A,)> for ComposeTwo<G, F>
where
    F: FnMut(A) -> B,
    G: FnMut(B) -> C,
{
    extern "rust-call" fn call_mut(&mut self, (x,): (A,)) -> Self::Output {
        (self.1)((self.0)(x))
    }
}

impl<A, B, C, G, F> std::ops::Fn<(A,)> for ComposeTwo<G, F>
where
    F: Fn(A) -> B,
    G: Fn(B) -> C,
{
    extern "rust-call" fn call(&self, (x,): (A,)) -> Self::Output {
        (self.1)((self.0)(x))
    }
}

fn compose_two<G, F>(f: F, g: G) -> ComposeTwo<G, F> {
    ComposeTwo(f, g)
}

Another way to break this assumption is by making the returned closure Copy. Copy type can never implement Drop, and the compiler knows that, and assumes they don't. Unfortunately, because the closure captures f and g, they need to be Copy too:

fn compose_two<A, B, C, G, F>(f: F, g: G) -> impl Fn(A) -> C + Copy
where
    F: Fn(A) -> B + Copy,
    G: Fn(B) -> C + Copy,
{
    move |x| g(f(x))
}

The last way is the most complicated to explain. First, I need to explain why the compiler thinks the closure can capture x, while in fact it cannot.

Let's first think why the closure cannot do that: what lifetime will it put in place of the '? below?

struct Closure {
    f: some_function_type,
    g: some_function_type,
    captured_x: Option<&'? mut Foo>,
}

When baz was defined (where we must decide what lifetime we'll use), we still don't know what will be passed to the closure, and so we don't know what lifetime we should use!

This knowledge, which is essentially "the closure can be called with any lifetime", is passed through Higher-Ranked Trait Bounds (HRTB) in Rust, spelled for<'lifetime>. So, A in compose_two() should've been HRTB.

But here lies the problem: generic parameters cannot be HRTB. They must be instantiated with a concrete lifetime. So, the compiler chooses some lifetime 'x for baz, and this lifetime must be bigger than baz itself - otherwise it would contain a dangling lifetime - and therefore it theoretically could have a member with that lifetime, and so the compiler believes baz can store the reference to x, while in reality it cannot.

If only we could make it HRTB...

We can! If we does not make it completely generic, and instead specify it as a reference:

fn compose_two<A, B, C, G, F>(f: F, g: G) -> impl for<'a> Fn(&'a mut A) -> &'a mut C
where
    F: for<'a> Fn(&'a mut A) -> &'a mut B,
    G: for<'a> Fn(&'a mut B) -> &'a mut C,
{
    move |x| g(f(x))
}

Or, using elided form, since HRTB is the default for Fn trait bounds:

fn compose_two<A, B, C, G, F>(f: F, g: G) -> impl Fn(&mut A) -> &mut C
where
    F: Fn(&mut A) -> &mut B,
    G: Fn(&mut B) -> &mut C,
{
    move |x| g(f(x))
}

It unfortunately also requires B: 'static, because the compiler cannot conclude B will live long enough (another limitation of the language), but then it works!

fn compose_two<A, B: 'static, C, G, F>(f: F, g: G) -> impl Fn(&mut A) -> &mut C
where
    F: Fn(&mut A) -> &mut B,
    G: Fn(&mut B) -> &mut C,
{
    move |x| g(f(x))
}

Answer 2

First, it helps to inline the macro a couple of times using the rust-analyzer quick fix (keybind ctrl+. in vscode):

fn main() {
    let baz = compose_two(foo, bar);
    let mut x = Foo { a: 3 };
    let y = baz(&mut x);
    println!("{:?}", y.a);
}

In this case it just calls compose_two with the given functions. The full compiler diagnostic for the error is:

error[E0597]: `x` does not live long enough
  --> src/main.rs:34:17
   |
34 |     let y = baz(&mut x);
   |                 ^^^^^^ borrowed value does not live long enough
35 |     println!("{:?}", y.a);
36 | }
   | -
   | |
   | `x` dropped here while still borrowed
   | borrow might be used here, when `baz` is dropped and runs the destructor for type `impl Fn(&mut Foo) -> &mut Foo`
   |
   = note: values in a scope are dropped in the opposite order they are defined

I'm not entirely sure what's going on here but I think the problem is that the compiler can't infer that the mutable reference isn't used after the call to baz for some reason. It works if you inline the call to compose_two:

    let baz = move |x| bar(foo(x));

It works if you define baz as a separate function:

fn baz(x: &mut Foo) -> &mut Foo{
    compose_two(foo, bar)(x)
}

fn main() {
    let bazz = baz;
    let mut x = Foo { a: 3 };
    let y = bazz(&mut x);
    println!("{:?}", y.a);
}

But I don't know why it can't infer that from the impl Fn(&mut Foo) -> &mut Foo type of baz (the variable).

Edit: some discussion in the question's comments has shed some light here. As Chayim Friedman described, the compiler can't assume much from the impl Fn(&mut Foo) -> &mut Foo type so it assumes it could be holding the reference (but it isn't). Which is why it works if you inline the closure or extract baz to it's own function, because then the compiler has more information.

In this case, it can be fixed by declaring x first so that baz (which the compiler thinks is holding a mutable reference to x) is dropped before x:

fn main() {
    let mut x = Foo { a: 3 };
    let baz = compose_two(foo,bar);
    let y = baz(&mut x);
    println!("{:?}", y.a);
}