How do I flatten a recursive structure using recursive iterators?

Question

I'm trying to flatten a recursive structure but I'm having trouble with recursive iterators.

Here's what the struct looks like:

#[derive(Debug, Clone)]
pub struct C {
    name: String,
    vb: Option<Vec<B>>,
}

#[derive(Debug, Clone)]
pub struct B {
    c: Option<C>,
}

#[derive(Debug, Clone)]
pub struct A {
    vb: Option<Vec<B>>,
    flat_c: Option<Vec<C>>,
}

My plan is to traverse the vb vector and flatten it into flat_c. I want it to look like this, or at least, be a Vec<String>:

Some([
    C {
        name: "foo",
        vb: None,
    },
    C {
        name: "bar",
        vb: None,
    },
    C {
        name: "fizz",
        vb: None,
    },
    C {
        name: "buzz",
        vb: None,
    },
])

Here what I managed to do, somewhat flattening the struct, but only for the last element, as the recursion is not implemented.

impl A {
    fn flat_c(self) -> Self {
        let fc: Vec<C> = self
            .vb
            .clone()
            .unwrap()
            .iter()
            .flat_map(|x| x.c.as_ref().unwrap().vb.as_ref().unwrap().iter())
            .cloned()
            .map(|x| x.c.unwrap())
            .collect();

        Self {
            flat_c: Some(fc),
            ..self
        }
    }
}

fn main() {
    let a = A {
        vb: Some(vec![
            B {
                c: Some(C {
                    name: "foo".to_string(),
                    vb: Some(vec![B {
                        c: Some(C {
                            name: "bar".to_string(),
                            vb: None,
                        }),
                    }]),
                }),
            },
            B {
                c: Some(C {
                    name: "fiz".to_string(),
                    vb: Some(vec![B {
                        c: Some(C {
                            name: "buzz".to_string(),
                            vb: None,
                        }),
                    }]),
                }),
            },
        ]),
        flat_c: None,
    };

    let a = a.flat_c();
    println!("a: {:#?}", a);
}

playground

The output for flat_c:

Some([
    C {
        name: "bar",
        vb: None,
    },
    C {
        name: "buzz",
        vb: None,
    },
])

I haven't dived into the Iterator trait implementation that might be required for this problem.

How would I tackle this problem? Maybe using a fold? Perhaps a recursive approach is not even needed? I'm at loss.

Answer 1

It's a good idea to be familiar with common data structures. You have a tree, and there are several ways to traverse a tree. You haven't precisely specified which method to use, so I chose one arbitrarily that's easy to implement.

The key here is to implement an iterator that keeps track of some state: all of the nodes yet to be visited. On each call to Iterator::next, we take the next value, save aside any new nodes to visit, and return the value.

Once you have the iterator, you can collect it into a Vec.

use std::collections::VecDeque;

impl IntoIterator for A {
    type IntoIter = IntoIter;
    type Item = String;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter {
            remaining: self.vb.into_iter().flatten().collect(),
        }
    }
}

struct IntoIter {
    remaining: VecDeque<B>,
}

impl Iterator for IntoIter {
    type Item = String;

    fn next(&mut self) -> Option<Self::Item> {
        self.remaining.pop_front().and_then(|b| {
            b.c.map(|C { name, vb }| {
                self.remaining.extend(vb.into_iter().flatten());

                name
            })
        })
    }
}

fn to_strings(a: A) -> Vec<String> {
    a.into_iter().collect()
}

#[derive(Debug, Clone)]
struct A {
    vb: Option<Vec<B>>,
}

#[derive(Debug, Clone)]
struct B {
    c: Option<C>,
}

#[derive(Debug, Clone)]
struct C {
    name: String,
    vb: Option<Vec<B>>,
}

fn main() {
    let example: A = A {
        vb: Some(vec![
            B {
                c: Some(C {
                    name: "Hello ".to_string(),
                    vb: None,
                }),
            },
            B {
                c: Some(C {
                    name: "World!".to_string(),
                    vb: None,
                }),
            },
        ]),
    };
    println!("The example struct: {:?}", example);
    //clone a copy for a second example, because to_strings() takes ownership of the example A struct
    let receipt: A = example.clone();
    println!("Iterated: {:?}", to_strings(example));
    // another example of using to_strings()
    println!(
        "As a string: {:?}",
        to_strings(receipt).into_iter().collect::<String>()
    );
}

From here, it should be straight-forward to create an iterator of B if that's what you need. Having all of the None values seemed silly, so I left them out and directly returned Strings.

I also made this a by-value iterator. You could follow the same pattern to create an iterator that returned references to the B / String and only clone them as needed.

See also:

• How to implement Iterator and IntoIterator for a simple struct?
• Implement IntoIterator for binary tree
• Cannot obtain a mutable reference when iterating a recursive structure: cannot borrow as mutable more than once at a time
• Recursive inorder traversal of a binary search tree

Answer 2

There is my solution:

impl C {
    fn flat(&self) -> Vec<C> {
        let mut result = Vec::new();
        result.push(C {
            name: self.name.clone(),
            vb: None,
        });
        if self.vb.is_some() {
            result.extend(
                (self.vb.as_ref().unwrap().iter())
                    .flat_map(|b| b.c.as_ref().map(|c| c.flat()).unwrap_or(Vec::new())),
            );
        }
        return result;
    }
}

impl A {
    fn flat_c(self) -> Self {
        let fc = (self.vb.as_ref().unwrap().iter())
            .flat_map(|b| b.c.as_ref().unwrap().flat())
            .collect();

        Self {
            flat_c: Some(fc),
            ..self
        }
    }
}

It adds flat function for C because the C is the source of the recursion and only this struct may properly handle it.

Because of those Options it looks scary and there is hard to deal with cryptic error messages. This solution supposes that all b.cs of initial a is not a None. Otherwise, it will panic. My advice is to avoid using Option<Vec> and use just empty vector instead of None.

https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=09ea11342cdd733b03172c0fc13c85fd