rustfunctional-programmingfunction-composition

How to compose functions that modify a struct in Rust


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?


Solution

  • Can this be made to work?

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

    You can make it work if either:

    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))
    }