Higher-Order Functions in Rust

Rust is an imperative language but it provides many tools in the standard library which adhere to a more functional style, like the Iterator trait and its methods like map, for_each, and filter. This is a quick run-down of how to define your own higher-order functions in Rust which can both take closures as parameters and return closures in such a way that you can use the two together.

edit @theodesp rightly pointed out that these examples do not follow a purely functional style - I'm mutating the input in place. I wrote it this way because the pure solutions are well covered in the standard library and this use case is less clear, but I've updated the Rust Playground with some hand-rolled pure implementations as well for comparison.

To demonstrate these we're going to work with a 2D grid of numbers:

type Grid = Vec<Vec<i32>>;

fn prepare_grid(rows: i32, columns: i32) -> Grid {
    let mut ret = Vec::new();
    for _ in 0..rows {
        let mut row = Vec::new();
        for _ in 0..columns {
            row.push(0i32)
        }
        ret.push(row);
    }
    ret
}

This function will initialize the grid to all 0s:

fn main() {
    let mut grid = prepare_grid(5, 5);
    println!("{:?}", grid);
}
// [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]

For whatever nefarious (?) purposes we might have in mind for our number grid it may be useful to act on a row at a time. We can accept a closure as a parameter using a generic data type implementing the Fn trait:

fn map_rows<F>(grid: &mut Grid, func: F)
where
    F: Fn(&mut Vec<i32>)
{
    for row in grid {
        func(row)
    }
}

Now we can, for example, increment the first element of each row:

map_rows(&mut grid, |row: &mut Vec<i32>| row[0] += 1);

println!("{:?}", grid);
// [[1, 0, 0, 0, 0], [1, 0, 0, 0, 0], [1, 0, 0, 0, 0], [1, 0, 0, 0, 0], [1, 0, 0, 0, 0]]

Returning a function is a tad trickier because of how Rust manages lifetimes. What say you wanted to decide how much to increment that first value by. You need to return a Fn(&mut Vec<i32>). Thing is, Rust needs to calculate the lifetime of that return function at compile time. We have to explicitly tell the compiler that this function will live as long as the input parameter lives by using a reference and assigning it the lifetime 'a:

fn make_incrementer<'a>(amount:&'a i32) -> Box<Fn(&mut Vec<i32>) + 'a> {
    Box::new(move |row: &mut Vec<i32>| row[0] += amount)
}

We're using the Box pointer type so that make_incrementer's return type's size is known at compile time, and using a move closure to ensure a new stack frame is allocated for the closure and we copy amount into it - thus escaping make_incrementer's stack frame.

Using it with map_rows requires some changes:

fn map_rows2<F>(grid: &mut Grid, func: Box<F>)
where
    F: for<'a> Fn(&'a mut Vec<i32>) + ?Sized
{
    for row in grid {
        (*func)(row)
    }
}

We now have an explicit lifetime to deal with but generally it would need to be at least as long as the whole function. In this case the compiler will complain because it only comes into scope when the closure is called inside the function, way down in our for loop - clearly shorter than the whole function. The for<...> syntax is a feature called a Higher-Ranked Trait Bound, which tells the compiler to explicitly calculate the minimum lifetime to invoke our closure instead of defining it for the whole function with map_rows2<'a, F>(...), circumventing this problem.

We also need to pass the whole Box in as our argument because all local variables must be Sized (though apparently unsized locals are in the works). You don't need to take ownership - &Box<F> is fine. This, though, will cause Rust to freak out because now we don't have a Sized value for F like before. We needed to create one in order to get make_incrementer to compile but we've gone and undone it all by unwrapping it. Luckily there's an escape hatch - by adding ?Sized we can relax that requirement. The only other change is getting at the actual closure with (*func) inside our for loop.

Now we can go wild:

map_rows2(&mut grid, make_incrementer(&2));

println!("{:?}", grid);
// [[3, 0, 0, 0, 0], [3, 0, 0, 0, 0], [3, 0, 0, 0, 0], [3, 0, 0, 0, 0], [3, 0, 0, 0, 0]]
// we already added 1!

Note: This example would work fine without the reference/lifetimes because our amount has type i32, which implements Copy - the move closure will just copy it for you. This means you can omit them in the map_rows2 too. The reference is only strictly necessary if your parameter to your factory function is not Copy - say, a String.

All in all higher-order functions are definitely more cumbersome to use in Rust than in a garbage-collected functional language, but absolutely not impossible. You just need to massage the borrow checker a little and make sure you know what exactly you're asking for!

Here's the full sample on the Rust Playground to tinker with.

Comments

[Theofanis Despoudis]

I had the impression that map does not modify the original array but returns a new array with all the elements mapped to the function's domain. Also if you modify something in functional programming then it's unsafe and not really functional.

Can you show an example of this? (not using mut for example)?

[Ben Lovy]

I've implemented pure versions of map_rows, map_rows2, and make_incrementer - here's an updated Playground link.

The signatures are almost the same - I've just replaced &mut Grid with &[Vec<i32>]. The only difference is the body of the functions themselves but you're correct - this set adheres more strictly to a functional paradigm whereas the originals do not.

[Alexander Payne]

Rust isn’t functional; it’s just dressed like it. The ownership rules allow the compiler to prove that mut self -> Self functions and &mut self -> () methods are wholly equivalent, and therefore eliminate the need to construct a new object and destroy the old.

[Ben Lovy]

Definitely, good catch - map was probably a bad choice of name. This example is not a pure functional style, but higher-order functions do not need to be pure. I thought that example was more interesting - if you do need a pure map I generally recommend using the standard library. I'd argue the above functions are still "functional programming", just not "pure functional programming". I don't know that I'd call it unsafe - fearless mutation like this is more idiomatic Rust.

I'll write up an example later this morning once I've had some coffee!

[Jay]

I implemented a simplified version, of your pure functions. Used the standard map to map over rows, and some other goodies.

I suggest to use vec![init_value;size] instead of creating and pushing vec(s).

I used pass by value in Boxed case as passing by value for primitives is trivial, they implement Copy trait.

Rust Playground

[Ben Lovy]

Way better! I was specifically trying to avoid map but you're definitely right about using vec! and passing by value for this function.

Using the trait object in the parameter list is better for the local closure version too.

[jeikabu]

Don't think I've come across for<> before. Interesting.
Haven't needed to dig very deep on some aspects of Rust, and this certainly touches on a few of them.

[Ben Lovy]

I hadn't either. Useful, but not commonly necessary. I couldn't find a reference to it in the current edition of the book, here is the link to the relevant section of the first edition.