SmoothTraderKen
SmoothTraderKen

Reputation: 652

Function composition chain in Rust

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

Upvotes: 1

Views: 1009

Answers (1)

cdhowie
cdhowie

Reputation: 169318

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

Upvotes: 6

Related Questions