Search code examples
rustiterator

Increase the size of iterator on a per element basis / Expand and Map iterator


Im trying to find or implement a method that works on any iterator, which allows me to iterate over all elements and for every element create multiple new elements that then become part of the old iterator. This effectively increases the iterator in size.

First of all, there is already a way to do this, just that is not feasible for the program I'm writing, since this is quite long and tedious and I would to do this very often.

//This is not the actual usecase, just example code
fn main() {
   let numbers: Vec<i32> = vec![10, 20, 30];
   
   //this bit here
   let result_nums: Vec<i32> = Vec::new();
   for num in number.iter() {
      result_nums.push(num);
      result_nums.push(num + 1);
      result_nums.push(num + 2);
   }
}

The main operation in this piece of code I need, the code in the for loop, which 'expands' the numbers array into a new array result_nums.

The way I want to write this code is the following:

fn main() {
   let numbers: Vec<i32> = vec![10, 20, 30];

   //this bit
   let result: Vec<i32> = numbers.iter().expand(|x, container| -> {
       container.push(x);
       container.push(x + 1);
       container.push(x + 2);
   }).collect()

   //the iterator will now contain the elements: [10, 11, 12, 20, 21, 22, 30, 31, 32]
}

This is the most ergonomic way I could think of. Keep in mind I'm going to be chaining the expand() operations and I won't just collect after every single one, this is just for the sake of the example. It would really save time and would be way less verbose, since I won't have to keep making temporary arrays with the whole type-signature, aswell as write the entire for loop.

If anyone knows of a crate that can do this or a better, even more ergonomic way of writing these expand operation, then please let me know!


Solution

  • @user4815162342 Suggested the use of the standard-library method flat_map, which makes this exact implementation possible!

    Implementation:

    fn main() {
        let numbers: Vec<i32> = vec![10, 20, 30];
        let result: Vec<i32> = numbers
            .iter()
            .flat_map(|&x| [x, x + 1, x + 2])
            .collect();
        assert_eq!(result, [10, 11, 12, 20, 21, 22, 30, 31, 32])
    }
    

    Playground