Iterating over (x,y) positions of a grid using functional programming methods - functional-programming

Question

I want to iterate over positions (x, y) in a 2-dimensional grid and call a function for each position (providing x and y as parameters). I know how this can be done using for loops, but I'd like to write this using functional programming methods such that I can later take advantage of libraries like Rayon. I've managed to build a functional version, but it seems complicated and I wanted to ask whether there's a cleaner way to achieve this.
Here's a small example of computing the maximum sum of all 2x2 patches in a grid:
use std::cmp::max;
const WIDTH: usize = 4;
const HEIGHT: usize = 3;
type Grid = [[u32; WIDTH]; HEIGHT];
fn main() {
let grid: Grid = [
[1, 3, 5, 8],
[3, 9, 4, 2],
[3, 4, 5, 0],
];
let coords = (0..WIDTH - 1).flat_map(|x| (0..HEIGHT - 1).map(move |y| (x, y)));
let max_sum = coords.map(|(x, y)| sum_2x2(x, y, &grid)).max().unwrap();
println!("Max 2x2 patch: {}", max_sum);
}
fn sum_2x2(x: usize, y: usize, grid: &Grid) -> u32 {
[
grid[y][x],
grid[y][x + 1],
grid[y + 1][x],
grid[y + 1][x + 1],
]
.iter()
.sum()
}
The line let coords = (0..WIDTH - 1).flat_map(|x| (0..HEIGHT - 1).map(move |y| (x, y))); seems pretty complex for such a simple task. In Python, I'd do the following to get the positions (which I think is much cleaner):
>>> from itertools import product
>>> product(range(HEIGHT-1), range(WIDTH-1))
Is there a better way to write this or do I just need to get used to it?

Answer 1

You can possibly simplify the code with one of the loop comprehension crates like map_for or mdo.
For example with map_for:
let max_sum = map_for!(
move;
x <- 0..WIDTH-1;
y <- 0..HEIGHT-1;
=> sum_2x2(x, y, &grid)
).max().unwrap();
I haven't tried it with Rayon, but since map_for is simply syntax sugar for a sequence of calls to map and flat_map and works with any type that implements those functions, it should work with Rayon as well.
Full disclosure: I am the author of the map_for crate.

Answer 2

Another option is to use cartesian_product from the Itertools crate, which has many convenient methods involving iterators.
use itertools::Itertools;
(0..WIDTH-1).cartesian_product(0..HEIGHT-1).map(|(x, y)| sum_2x2(x, y, &grid)).max()

Answer 3

As #hellow pointed out in the comments, there's an iproduct macro in the itertools crate that does what I want. Using it, my code can be rewritten like this:
use itertools::iproduct;
...
let coords = iproduct!(0..WIDTH - 1, 0..HEIGHT - 1);
...
That's exactly what I was looking for. Thanks to everyone who commented / posted answers here.
PS: If you want to use an itertools iterator with rayon, you can do this by using the par_bridge method, e.g. iproduct!(0..WIDTH - 1, 0..HEIGHT - 1).par_iter();. Took me a while to figure that out since par_iter and into_par_iter don't work.