CODEWITHSUNDEEP
rusttypesiterator
user1002430
user1002430

Reputation:

Rust function to combine items of a slice of iterators using a function

I'd like to create a function which takes a slice of iterators, and combines their elements with a function. The first result would be a function of the first elements of each iterator, the second would be a function of the second elements, etc ...

Visually, vertically three iterators, to the right is the output iterator:

it1   it2   it3         output it
a1    b1    c1    ->    f(a1, b1, c1)
a2    b2    c2    ->    f(a2, b2, c2)
a3    b3    c3    ->    f(a3, b3, c3)

The following does not compile but shows the idea:

  pub fn combine<I, A, F, B>(iters: &mut [I], combine_fn: F) -> impl Iterator<Item = B>
  where
      I: Iterator<Item = A>,
      F: Fn(impl Iterator<Item = Option<A>>) -> B,
  {
      std::iter::from_fn(|| {
          let row = iters.iter_mut().map(|it| it.next());
          Some(combine_fn(row))
      })
  }

What am I missing in order to make this function work?

Upvotes: 0

Views: 168

Answers (1)

harmic
harmic

Reputation: 30597

Here is a version that works:

pub fn combine<'a, I, A, F, B>(iters: &'a mut [I], combine_fn: F) -> impl Iterator<Item = B> + 'a
where
     I: Iterator<Item = A>,
     F: Fn(&mut dyn Iterator<Item = A>) -> B + 'a,
{
     std::iter::from_fn(move || {
         let mut row = iters.iter_mut().map(|it| it.next()).flatten();
         Some(combine_fn(&mut row))
     })
}

Playground

Key things that needed to be fixed:

  • The closure passed into from_fn had to be annotated with move so that the argument combine_fn gets moved into the closure rather than referenced from it.

  • You can't use impl Trait in a Fn trait in argument position. An alternative is to use dyn Trait instead (although that does imply dynamic dispatch).

  • The row iterator passed into the combine_fn has to be flattened, otherwise the values returned would be Option<Option<A>> instead of Option<A>.

  • Lifetime annotations are needed to ensure that the output iterator has a lifetime no longer than the lifetime of both the iters array and any references captured by combine_fn.

Upvotes: 1

Related Questions