Building an iterator recursively results in "recursive opaque type" error

user_10048799 0
Viewed 129

I have a struct which owns a Vec and has an optional reference to a parent struct of its own type. I want a iterator that iterates over the items and then if there is a parent recursively iterate.

I currently have the following, but it allocates on the heap and uses dynamic dispatch:

struct Struct1 {
    items: Vec<(String, i32)>, 
    parent: Option<Rc<Struct1>>
}

impl Struct1 {
    pub fn keys<'a>(&'a self) -> Box<dyn Iterator<Item = &'a str> + 'a> {
        let iter = self.items.iter().map(|(key, _)| key.as_str());
        if let Some(child) = &self.parent {
            Box::new(iter.chain(child.keys()))
        } else {
            Box::new(iter)
        }
    }
}

Is there a alternative way to write the keys() implementation which doesn't allocate on the heap or use dynamic dispatch?

I am thinking of the following:

impl Struct1 {
    fn keys<'a>(&'a self) -> impl Iterator<Item = &'a str> {
        let iter = self.0.iter().map(|(key, _)| key.as_str());
        if let Some(child) = &self.1 {
            Either2::A(iter.chain(child.keys()))
        } else {
            Either2::B(iter)
        }
    }
}

pub enum Either2<A, B> {
    A(A),
    B(B),
}

impl<T, A: Iterator<Item = T>, B: Iterator<Item = T>> Iterator for Either2<A, B> {
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        match self {
            Self::A(a) => a.next(),
            Self::B(b) => b.next(),
        }
    }
}

Unfortunately this does not compile due to there being a recursive opaque type here impl Iterator<Item = &'a str>. And I can see why as the Either2::A is needed for the opaque type but also uses the opaque type.

Is there a way I could restructure the impl ... type reference to be recursive aware (possibly by using a type statement)? Or is there another way to rewrite the whole function to get it to work that doesn't involve intermediately collecting (and allocating) into a vec?

rust
1 Answers

These recursive structures also have a small performance problem: they build the entire parent-iterator chain before the first item is examined, even if the iterator will never reach the parent items. Instead, we can follow the parent chain and, separately, map each &Struct1 we iterate over to an iterator of its own keys. This way, the recursive iterator structure has been “flattened” into an “iterative” one that is efficient and accepted by the compiler:

use std::rc::Rc;

struct Struct1 {
    items: Vec<(String, i32)>,
    parent: Option<Rc<Struct1>>,
}

impl Struct1 {
    pub fn keys(&self) -> impl Iterator<Item = &str> + '_ {
        // Construct an iterator that follows parent references.
        let mut current: Option<&Struct1> = Some(self);
        let parents_iter = std::iter::from_fn(move || {
            let result = current;
            if let Some(Struct1 { parent, .. }) = current {
                // .as_deref() turns &Option<Rc<Struct1>> into Option<&Struct1>
                current = parent.as_deref();
            }
            result
        });

        parents_iter.flat_map(|p| p.only_self_keys())
    }

    fn only_self_keys(&self) -> impl Iterator<Item = &str> {
        self.items.iter().map(|(key, _)| key.as_str())
    }
}

fn main() {
    let input = Struct1 {
        items: vec![(String::from("a"), 1)],
        parent: Some(Rc::new(Struct1 {
            items: vec![(String::from("b"), 2), (String::from("c"), 3)],
            parent: None,
        })),
    };

    dbg!(input.keys().collect::<Vec<_>>());
}

One could also write a separate parent iterator struct that implements Iterator explicitly rather than using std::iter::from_fn.

user_996920
Related