Function composition chain in Rust

Question

I want to implement function composition in an object method chain in Rust.

Before that, I confirmed that implementation of an object method chain for "pipe" of function application is possible in Rust as follows:

https://docs.rs/apply/0.2.2/apply/trait.Apply.html

trait Pipe: Sized {
    fn pipe<F, B>(self, f: F) -> B
    where
        F: FnOnce(Self) -> B,
    {
        f(self)
    }
}

impl<T> Pipe for T {}
 
fn main() {
    let string = 1.pipe(|x| x * 2).pipe(|x| x + 1).pipe(|x| x.to_string());

    println!("{}", string);
}

https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=d6210a499e2522ff04e0cbbeef8cedc5

Now, the code of function composition is How to compose functions in Rust?

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

fn main() {
    let multiply_and_add = compose(|x| x * 2, |x| x + 2);
    let divide_and_subtract = compose(|x| x / 2, |x| x - 2);

    let finally = compose(multiply_and_add, divide_and_subtract);
    println!("Result is {}", finally(10));
}

So the goal would be

let composed_f = (|x| x * 2).compose(|x| x + 1).compose(|x| x.to_string());

and what is the proper code to implement function composition in an object method chain in Rust like "pipe"?

EDIT

for function composition macro quoted from https://stackoverflow.com/a/45792463/11316608

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

fn main() {
    let add = |x| x + 2;
    let multiply = |x| x * 2;
    let divide = |x| x / 2;
    let intermediate = compose!(add, multiply, divide);

    let subtract = |x| x - 2;
    let finally = compose!(intermediate, subtract);

    println!("Result is {}", finally(10));
}

Side note: you should take FnMut when possible. There's no reason to restrict callers to FnOnce. — cdhowie
– cdhowie, Commented Jun 13, 2022 at 4:45
@cdhowie Thanks for your side-note. It looks true and I only takes the first method from the example trait: docs.rs/apply/0.2.2/src/apply/lib.rs.html#8-23 — SmoothTraderKen
– SmoothTraderKen, Commented Jun 13, 2022 at 4:47

cdhowie · Accepted Answer · 2022-06-13 06:51:07Z

6

Implementing this directly on closures is tricky, but we can make it vastly simpler with a newtype that represents a composeable function. This adds a step to convert into a composeable function and then out again at the end, but it winds up being pretty reasonable (especially considering its expressive power), it requires no boxing or other clever tricks, and the implementation is very easy to understand.

Because it winds up with unboxed, nested closure invocations, it is also very likely that the compiler will be able to optimize away all of the composition calls and inline the closures into each other, likely even into the final call site.

use std::marker::PhantomData;

struct ComposeableFn<F, A, B>(F, PhantomData<*mut A>, PhantomData<*mut B>);

impl<F, A, B> ComposeableFn<F, A, B>
    where F: FnMut(A) -> B
{
    pub fn new(f: F) -> Self {
        Self(f, Default::default(), Default::default())
    }
    
    pub fn compose<C>(self, mut next: impl FnMut(B) -> C)
        -> ComposeableFn<impl FnMut(A) -> C, A, C>
    {
        let mut prior = self.0;
        
        ComposeableFn::new(move |v| next(prior(v)))
    }
    
    pub fn into_inner(self) -> F {
        self.0
    }
}

Used like so:

fn main() {
    let mut composed_f = ComposeableFn::new(|x: i32| x * 2)
        .compose(|x| x + 1)
        .compose(|x| x.to_string())
        .into_inner();
        
    println!("{:?}", composed_f(21)); // "43"
}

(Playground)

The PhantomData members are required to allow the implementation to name A and B, otherwise they would be considered unused. I chose *mut _ as this makes the type invariant in A and B.

You can make the expression of a composition a bit nicer with a macro:

macro_rules! compose {
    ({ $fh:expr } $( { $ft:expr } )*) => {
        ComposeableFn::new($fh)
        $( .compose($ft) )*
        .into_inner()
    }
}

fn main() {
    let mut composed_f = compose!(
        {|x: i32| x * 2}
        {|x| x + 1}
        {|x| x.to_string()}
    );
        
    println!("{:?}", composed_f(21));
}

edited Jun 13, 2022 at 6:51

answered Jun 13, 2022 at 5:01

cdhowie

172k25 gold badges303 silver badges324 bronze badges

Sign up to request clarification or add additional context in comments.

14 Comments

SmoothTraderKen Over a year ago

This is like one of "wow" because your implementation is well-considered for simplicity (that I also consider it most) among other possibilities of implementations that I need to learn all of them that I don't know anyway. Thank you! I really appreciate it.

cdhowie Over a year ago

@SmoothTraderKen I'm glad you find it useful! Note that I didn't extensively test with closures that borrow -- you may need to tweak compose to introduce a named lifetime parameter to handle cases where closures borrow from their environment. I didn't run into any such issues, but I'm not confident enough to say they won't come up.

SmoothTraderKen Over a year ago

Thanks again for your side-note, @cdhowie. I have an impression that the example-code of trait Pipe of function application may be refined by using your code somehow, but I don't know so far. If there's something let me mention to you here. I appreciate it again.

cdhowie Over a year ago

@SmoothTraderKen The roadblock there is that you'd need your trait's compose function to return impl FnMut(A) -> C but trait functions can't (currently) have impl return types, and associated types won't work either since closure types can't be named. RFC 2071 might help here but it's not stable yet. You can work around this by boxing the closure, but then you're adding indirection and heap allocation per composed function, and severely limiting the optimizer's ability to inline the composed functions.

cdhowie Over a year ago

All of that to say, the newtype approach works today and should have pretty darn close to zero runtime impact compared to composing the functions by hand, doesn't depend on unstable features, and the two bookends of create-newtype and unwrap-newtype are, in my opinion, both trivial and easily understood. And, on second thought, you can probably hide that all in a macro anyway.

|

Collectives™ on Stack Overflow

Function composition chain in Rust

1 Answer 1

14 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

14 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related