Home > database >  How do I return an iterator of arbitrary type from a trait function?
How do I return an iterator of arbitrary type from a trait function?

Time:12-25

I have the following code where I would like to implement a trait that returns an iterator:

trait MyTrait<V> {
    type Iter: Iterator<Item = V>;
    fn produce_iter(&self) -> Self::Iter;
}

struct MyStruct<V> {
    values: Vec<V>,
}

impl<V> MyTrait<V> for MyStruct<V>
where
    V: Clone,
{
    type Iter = std::vec::IntoIter<V>;

    fn produce_iter(&self) -> Self::Iter {
        self.values.clone().into_iter()
    }
}

fn main() {
    let s = MyStruct {
        values: vec![1, 2, 3],
    };

    s.produce_iter().for_each(|v| println!("{:?}", v));
}

However if I want my implementation of fn produce_iter(&self) to be more complicated such as

fn produce_iter(&mut self) -> Self::Iter {
    self.values
        .clone()
        .into_iter()
        .flat_map(|v| vec![v, v].into_iter())
        .into_iter()
}

instead of

fn produce_iter(&mut self) -> Self::Iter {
    self.values.clone().into_iter()
}

I run into error[E0308]: mismatched types issues trying to figure out what type Iter = std::vec::IntoIter<V>; should be instead.

Is there a way of 'generifying' that error away so fn produce_iter(&self) can produce arbitrarily complex iterators without having to worry about what I set for type Iter? Maybe I should be looking at the problem from a different angle?

CodePudding user response:

There's no way to be generic over a return value, the return value is always either a concrete type or determined base on input type. The return type must be fully spelled out.

However, there is a way to return multiple types through trait object. Read more about trait object from The Rust Book.

Instead of returning a concrete type from produce_iter, return a Box<dyn Iterator<Item = V>> trait object.

trait MyTrait<V> {
    fn produce_iter(&self) -> Box<dyn Iterator<Item = V>>;
}

struct MyStruct<V> {
    values: Vec<V>,
}

impl<V> MyTrait<V> for MyStruct<V>
where
    V: Clone   'static,
{
    fn produce_iter(&self) -> Box<dyn Iterator<Item = V>> {
    Box::new(
        self.values
            .clone()
            .into_iter()
            .flat_map(|v| vec![v.clone(), v].into_iter())
        )
    }
}

fn main() {
    let s = MyStruct {
        values: vec![1, 2, 3],
    };

    s.produce_iter().for_each(|v| println!("{:?}", v));
}

But beware that trait objects are slower than concrete types, becuase trait objects use dynamic dispatch at runtime.

CodePudding user response:

On nightly you can use type_alias_impl_trait:

#![feature(type_alias_impl_trait)]

impl<V> MyTrait<V> for MyStruct<V>
where
    V: Clone   std::ops::Mul,
{
    type Iter = impl Iterator<Item = V>;

    fn produce_iter(&self) -> Self::Iter {
        self.values
            .clone()
            .into_iter()
            .flat_map(|v| vec![v, v].into_iter())
            .into_iter()
    }
}

On stable you must use Box<dyn Iterator>.

  • Related